diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 51cad297a19..b6a3303c330 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -32,6 +32,7 @@ def test_completion_custom_provider_model_name(): ) # Add any assertions here to check the response print(response) + print(response['choices'][0]['finish_reason']) except Exception as e: pytest.fail(f"Error occurred: {e}") @@ -107,8 +108,10 @@ def test_completion_claude_stream(): # Add any assertions here to check the response for chunk in response: print(chunk["choices"][0]["delta"]) # same as openai format + print(chunk["choices"][0]["finish_reason"]) except Exception as e: pytest.fail(f"Error occurred: {e}") +# test_completion_claude_stream() @@ -173,8 +176,10 @@ def test_completion_cohere_stream(): # Add any assertions here to check the response for chunk in response: print(chunk["choices"][0]["delta"]) # same as openai format + print(chunk["choices"][0]["finish_reason"]) except Exception as e: pytest.fail(f"Error occurred: {e}") +# test_completion_cohere_stream() def test_completion_openai(): @@ -293,8 +298,12 @@ def test_completion_openai_with_stream(): ) # Add any assertions here to check the response print(response) + for chunk in response: + print(chunk) + print(chunk["choices"][0]["finish_reason"]) except Exception as e: pytest.fail(f"Error occurred: {e}") +# test_completion_openai_with_stream() def test_completion_openai_with_functions(): @@ -317,12 +326,16 @@ def test_completion_openai_with_functions(): ] try: response = completion( - model="gpt-3.5-turbo", messages=messages, functions=function1 + model="gpt-3.5-turbo", messages=messages, functions=function1, stream=True ) # Add any assertions here to check the response print(response) + for chunk in response: + print(chunk) + print(chunk["choices"][0]["finish_reason"]) except Exception as e: pytest.fail(f"Error occurred: {e}") +# test_completion_openai_with_functions() def test_completion_azure(): diff --git a/litellm/utils.py b/litellm/utils.py index 507e130f554..34c618b178a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2198,12 +2198,22 @@ class CustomStreamWrapper: completion_obj["content"] = self.handle_openai_text_completion_chunk(chunk) else: # openai chat/azure models chunk = next(self.completion_stream) - completion_obj["content"] = self.handle_openai_chat_completion_chunk(chunk) + return chunk # open ai returns finish_reason, we should just return the openai chunk + + #completion_obj["content"] = self.handle_openai_chat_completion_chunk(chunk) # LOGGING threading.Thread(target=self.logging_obj.success_handler, args=(completion_obj,)).start() # return this for all models - return {"choices": [{"delta": completion_obj}]} + return { + "choices": + [ + { + "delta": completion_obj, + "finish_reason": "stop" + }, + ] + } except Exception as e: raise StopIteration diff --git a/pyproject.toml b/pyproject.toml index 8f8b312b2c5..0907ff1acb7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "0.1.580" +version = "0.1.582" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT License"