From a866bdb01ed1ed1b9cab4080c15d31fd5de15c2e Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 14 Nov 2024 10:53:37 -0800 Subject: [PATCH] fix ModelResponseIterator --- litellm/llms/anthropic/chat/handler.py | 4 +++- tests/llm_translation/base_llm_unit_tests.py | 3 +++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index d56834c6e..6976ea0e5 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -644,7 +644,9 @@ class AnthropicChatCompletion(BaseLLM): class ModelResponseIterator: - def __init__(self, streaming_response, sync_stream: bool, json_mode: bool): + def __init__( + self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False + ): self.streaming_response = streaming_response self.response_iterator = self.streaming_response self.content_blocks: List[ContentBlockDelta] = [] diff --git a/tests/llm_translation/base_llm_unit_tests.py b/tests/llm_translation/base_llm_unit_tests.py index acb764ba1..bcdbe9e95 100644 --- a/tests/llm_translation/base_llm_unit_tests.py +++ b/tests/llm_translation/base_llm_unit_tests.py @@ -45,6 +45,9 @@ class BaseLLMChatTest(ABC): ) assert response is not None + # for OpenAI the content contains the JSON schema, so we need to assert that the content is not None + assert response.choices[0].message.content is not None + def test_message_with_name(self): base_completion_call_args = self.get_base_completion_call_args() messages = [