phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 19:24:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

test(test_custom_logger.py): fix async streaming test

Browse source
This commit is contained in:
Krrish Dholakia 2024-03-15 18:18:54 -07:00
parent 3072137739
commit 4912ad5cf4
1 changed files with 11 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

12
litellm/tests/test_custom_logger.py
View file

@ -97,27 +97,23 @@ class TmpFunction:
) )
def test_async_chat_openai_stream(): @pytest.mark.asyncio
async def test_async_chat_openai_stream():
try: try:
tmp_function = TmpFunction() tmp_function = TmpFunction()
litellm.set_verbose = True litellm.set_verbose = True
litellm.success_callback = [tmp_function.async_test_logging_fn] litellm.success_callback = [tmp_function.async_test_logging_fn]
complete_streaming_response = "" complete_streaming_response = ""
async def call_gpt():
nonlocal complete_streaming_response
response = await litellm.acompletion( response = await litellm.acompletion(
model="gpt-3.5-turbo", model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}], messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}],
stream=True, stream=True,
) )
async for chunk in response: async for chunk in response:
complete_streaming_response += ( complete_streaming_response += chunk["choices"][0]["delta"]["content"] or ""
chunk["choices"][0]["delta"]["content"] or ""
)
print(complete_streaming_response) print(complete_streaming_response)
asyncio.run(call_gpt())
complete_streaming_response = complete_streaming_response.strip("'") complete_streaming_response = complete_streaming_response.strip("'")
response1 = tmp_function.complete_streaming_response_in_callback["choices"][0][ response1 = tmp_function.complete_streaming_response_in_callback["choices"][0][
"message" "message"
@ -130,7 +126,7 @@ def test_async_chat_openai_stream():
assert tmp_function.async_success == True assert tmp_function.async_success == True
except Exception as e: except Exception as e:
print(e) print(e)
pytest.fail(f"An error occurred - {str(e)}") pytest.fail(f"An error occurred - {str(e)}\n\n{traceback.format_exc()}")
# test_async_chat_openai_stream() # test_async_chat_openai_stream()