phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 11:14:04 +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

26
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(): response = await litellm.acompletion(
nonlocal complete_streaming_response model="gpt-3.5-turbo",
response = await litellm.acompletion( messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}],
model="gpt-3.5-turbo", stream=True,
messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}], )
stream=True, async for chunk in response:
) complete_streaming_response += chunk["choices"][0]["delta"]["content"] or ""
async for chunk in response: print(complete_streaming_response)
complete_streaming_response += (
chunk["choices"][0]["delta"]["content"] or ""
)
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()