From 16efb8db67a019f164c42a91d085967f5419ee3c Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 31 May 2025 13:29:51 -0700 Subject: [PATCH] Revert "Make gemini stream thinking as reasoning_content (#11290)" This reverts commit e0daa3da6857a9abc2cd95741ca356dbfae6ac9c. --- .../vertex_and_google_ai_studio_gemini.py | 76 ++++++++----------- ...test_vertex_and_google_ai_studio_gemini.py | 48 ++---------- 2 files changed, 36 insertions(+), 88 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index a706dff4c5..1ab851b9b4 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -43,6 +43,7 @@ from litellm.types.llms.openai import ( ChatCompletionToolCallChunk, ChatCompletionToolCallFunctionChunk, ChatCompletionToolParamFunctionChunk, + ChatCompletionUsageBlock, OpenAIChatCompletionFinishReason, ) from litellm.types.llms.vertex_ai import ( @@ -63,10 +64,8 @@ from litellm.types.utils import ( ChatCompletionTokenLogprob, ChoiceLogprobs, CompletionTokensDetailsWrapper, - Delta, - ModelResponseStream, + GenericStreamingChunk, PromptTokensDetailsWrapper, - StreamingChoices, TopLogprob, Usage, ) @@ -1651,15 +1650,14 @@ class ModelResponseIterator: self.accumulated_json = "" self.sent_first_chunk = False - def chunk_parser(self, chunk: dict) -> ModelResponseStream: + def chunk_parser(self, chunk: dict) -> GenericStreamingChunk: try: processed_chunk = GenerateContentResponseBody(**chunk) # type: ignore text = "" - reasoning_content = None tool_use: Optional[ChatCompletionToolCallChunk] = None finish_reason = "" - usage: Optional[Usage] = None + usage: Optional[ChatCompletionUsageBlock] = None _candidates: Optional[List[Candidates]] = processed_chunk.get("candidates") gemini_chunk: Optional[Candidates] = None if _candidates and len(_candidates) > 0: @@ -1671,11 +1669,7 @@ class ModelResponseIterator: and "parts" in gemini_chunk["content"] ): if "text" in gemini_chunk["content"]["parts"][0]: - if gemini_chunk["content"]["parts"][0].get("thought"): - reasoning_content = gemini_chunk["content"]["parts"][0]["text"] - else: - text = gemini_chunk["content"]["parts"][0]["text"] - + text = gemini_chunk["content"]["parts"][0]["text"] elif "functionCall" in gemini_chunk["content"]["parts"][0]: function_call = ChatCompletionToolCallFunctionChunk( name=gemini_chunk["content"]["parts"][0]["functionCall"][ @@ -1701,7 +1695,7 @@ class ModelResponseIterator: ## GEMINI SETS FINISHREASON ON EVERY CHUNK! if "usageMetadata" in processed_chunk: - usage = Usage( + usage = ChatCompletionUsageBlock( prompt_tokens=processed_chunk["usageMetadata"].get( "promptTokenCount", 0 ), @@ -1711,26 +1705,20 @@ class ModelResponseIterator: total_tokens=processed_chunk["usageMetadata"].get( "totalTokenCount", 0 ), - completion_tokens_details=CompletionTokensDetailsWrapper( - reasoning_tokens=processed_chunk["usageMetadata"].get( + completion_tokens_details={ + "reasoning_tokens": processed_chunk["usageMetadata"].get( "thoughtsTokenCount", 0 ) - ), + }, ) - returned_chunk = ModelResponseStream( - choices=[ - StreamingChoices( - index=0, - delta=Delta( - content=text, - tool_calls=[tool_use] if tool_use is not None else None, - reasoning_content=reasoning_content, - ), - finish_reason=finish_reason, - ) - ], + returned_chunk = GenericStreamingChunk( + text=text, + tool_use=tool_use, + is_finished=False, + finish_reason=finish_reason, usage=usage, + index=0, ) return returned_chunk except json.JSONDecodeError: @@ -1741,7 +1729,7 @@ class ModelResponseIterator: self.response_iterator = self.streaming_response return self - def handle_valid_json_chunk(self, chunk: str) -> ModelResponseStream: + def handle_valid_json_chunk(self, chunk: str) -> GenericStreamingChunk: chunk = chunk.strip() try: json_chunk = json.loads(chunk) @@ -1759,7 +1747,7 @@ class ModelResponseIterator: return self.chunk_parser(chunk=json_chunk) - def handle_accumulated_json_chunk(self, chunk: str) -> ModelResponseStream: + def handle_accumulated_json_chunk(self, chunk: str) -> GenericStreamingChunk: chunk = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "" message = chunk.replace("\n\n", "") @@ -1773,18 +1761,16 @@ class ModelResponseIterator: return self.chunk_parser(chunk=_data) except json.JSONDecodeError: # If it's not valid JSON yet, continue to the next event - return ModelResponseStream( - choices=[ - StreamingChoices( - index=0, - delta=Delta(content=""), - finish_reason="", - ) - ], + return GenericStreamingChunk( + text="", + is_finished=False, + finish_reason="", usage=None, + index=0, + tool_use=None, ) - def _common_chunk_parsing_logic(self, chunk: str) -> ModelResponseStream: + def _common_chunk_parsing_logic(self, chunk: str) -> GenericStreamingChunk: try: chunk = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "" if len(chunk) > 0: @@ -1798,15 +1784,13 @@ class ModelResponseIterator: elif self.chunk_type == "accumulated_json": return self.handle_accumulated_json_chunk(chunk=chunk) - return ModelResponseStream( - choices=[ - StreamingChoices( - index=0, - delta=Delta(content=""), - finish_reason="", - ) - ], + return GenericStreamingChunk( + text="", + is_finished=False, + finish_reason="", usage=None, + index=0, + tool_use=None, ) except Exception: raise diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 6b573d84d4..18ffd7ca60 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -334,52 +334,15 @@ def test_streaming_chunk_includes_reasoning_tokens(): } iterator = ModelResponseIterator(streaming_response=[], sync_stream=True) streaming_chunk = iterator.chunk_parser(chunk) - assert streaming_chunk.usage is not None - assert streaming_chunk.usage.prompt_tokens == 5 - assert streaming_chunk.usage.completion_tokens == 7 - assert streaming_chunk.usage.total_tokens == 12 + assert streaming_chunk["usage"] is not None + assert streaming_chunk["usage"]["prompt_tokens"] == 5 + assert streaming_chunk["usage"]["completion_tokens"] == 7 + assert streaming_chunk["usage"]["total_tokens"] == 12 assert ( - streaming_chunk.usage.completion_tokens_details.reasoning_tokens == 3 + streaming_chunk["usage"]["completion_tokens_details"]["reasoning_tokens"] == 3 ) -def test_streaming_chunk_includes_reasoning_content(): - """ - Ensure that when Gemini returns a chunk with `thought=True`, the parser maps it to `reasoning_content`. - """ - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - ModelResponseIterator, - ) - - # Simulate a streaming chunk from Gemini which contains reasoning (thought) content - chunk = { - "candidates": [ - { - "content": { - "parts": [ - { - "text": "I'm thinking through the problem...", - "thought": True, - } - ] - } - } - ], - "usageMetadata": {}, - } - - iterator = ModelResponseIterator(streaming_response=[], sync_stream=True) - streaming_chunk = iterator.chunk_parser(chunk) - - # The text content should be empty and reasoning_content should be populated - assert streaming_chunk.choices[0].delta.content == "" - assert ( - streaming_chunk.choices[0].delta.reasoning_content - == "I'm thinking through the problem..." - ) - - - def test_check_finish_reason(): config = VertexGeminiConfig() finish_reason_mappings = config.get_finish_reason_mapping() @@ -483,3 +446,4 @@ def test_vertex_ai_map_tool_with_anyof(): ] == { "anyOf": [{"type": "string", "nullable": True, "title": "Base Branch"}] }, f"Expected only anyOf field and its contents to be kept, but got {tools[0]['function_declarations'][0]['parameters']['properties']['base_branch']}" +