From 58568d51e8404e5d09a1a460d33e9df557a8fe88 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 15 Sep 2023 18:18:55 -0700 Subject: [PATCH] fix streaming stuff --- litellm/tests/test_completion.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 9feeea1b23..d9ac1fcdf9 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -214,7 +214,7 @@ def test_completion_cohere_stream(): pass except Exception as e: pytest.fail(f"Error occurred: {e}") -test_completion_cohere_stream() +# test_completion_cohere_stream() def test_completion_openai(): @@ -369,10 +369,10 @@ def test_completion_openai_with_stream(): print(response) for chunk in response: print(chunk) + if chunk["choices"][0]["finish_reason"] == "stop": + break print(chunk["choices"][0]["finish_reason"]) print(chunk["choices"][0]["delta"]["content"]) - except KeyError: - pass except Exception as e: pytest.fail(f"Error occurred: {e}") # test_completion_openai_with_stream() @@ -404,8 +404,11 @@ def test_completion_openai_with_functions(): print(response) for chunk in response: print(chunk) + if chunk["choices"][0]["finish_reason"] == "stop": + break print(chunk["choices"][0]["finish_reason"]) print(chunk["choices"][0]["delta"]["content"]) + except Exception as e: pytest.fail(f"Error occurred: {e}") # test_completion_openai_with_functions()