From de37c99021e58705d642c24d915555ba910b3906 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 15 Sep 2023 14:57:24 -0700 Subject: [PATCH] add test cases for checking role in streaming reqs --- litellm/tests/test_completion.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index b500d804277..7198e3607ff 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -106,6 +106,8 @@ def test_completion_claude_stream(): for chunk in response: print(chunk["choices"][0]["delta"]) # same as openai format print(chunk["choices"][0]["finish_reason"]) + print(chunk["choices"][0]["delta"]["content"]) + print(chunk["choices"][0]["delta"]["role"]) except Exception as e: pytest.fail(f"Error occurred: {e}") # test_completion_claude_stream() @@ -138,8 +140,11 @@ def test_completion_nlp_cloud_streaming(): for chunk in response: print(chunk["choices"][0]["delta"]["content"]) # same as openai format print(chunk["choices"][0]["finish_reason"]) + print(chunk["choices"][0]["delta"]["content"]) + print(chunk["choices"][0]["delta"]["role"]) except Exception as e: pytest.fail(f"Error occurred: {e}") +# test_completion_nlp_cloud_streaming() # test_completion_nlp_cloud_streaming() # def test_completion_hf_api(): @@ -206,6 +211,8 @@ def test_completion_cohere_stream(): for chunk in response: print(chunk["choices"][0]["delta"]) # same as openai format print(chunk["choices"][0]["finish_reason"]) + print(chunk["choices"][0]["delta"]["content"]) + print(chunk["choices"][0]["delta"]["role"]) except Exception as e: pytest.fail(f"Error occurred: {e}") # test_completion_cohere_stream() @@ -364,6 +371,8 @@ def test_completion_openai_with_stream(): for chunk in response: print(chunk) print(chunk["choices"][0]["finish_reason"]) + print(chunk["choices"][0]["delta"]["content"]) + print(chunk["choices"][0]["delta"]["role"]) except Exception as e: pytest.fail(f"Error occurred: {e}") # test_completion_openai_with_stream() @@ -396,6 +405,8 @@ def test_completion_openai_with_functions(): for chunk in response: print(chunk) print(chunk["choices"][0]["finish_reason"]) + print(chunk["choices"][0]["delta"]["content"]) + print(chunk["choices"][0]["delta"]["role"]) except Exception as e: pytest.fail(f"Error occurred: {e}") # test_completion_openai_with_functions() @@ -562,9 +573,10 @@ def test_completion_replicate_llama_stream(): try: response = completion(model=model_name, messages=messages, stream=True) # Add any assertions here to check the response - for result in response: - print(result) - # chunk_text = result['choices'][0]['delta']['content'] + for chunk in response: + print(chunk) + print(chunk["choices"][0]["delta"]["content"]) + print(chunk["choices"][0]["delta"]["role"]) except Exception as e: pytest.fail(f"Error occurred: {e}") # test_completion_replicate_llama_stream()