From dd96f91254e89d7880f98aacb1ccbbbbb3cdcf42 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 2 Jan 2026 12:31:31 +0530 Subject: [PATCH] Add thought signature for non tool call requests --- .../llms/vertex_ai/gemini/transformation.py | 13 ++- .../vertex_and_google_ai_studio_gemini.py | 34 ++++++ .../test_gemini_reasoning_content.py | 106 ++++++++++++++++++ 3 files changed, 152 insertions(+), 1 deletion(-) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index baa825bfcc..2bbdfa17cd 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -383,7 +383,18 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 and isinstance(_message_content, str) ): assistant_text = _message_content - assistant_content.append(PartType(text=assistant_text)) # type: ignore + # Check if message has thought_signatures in provider_specific_fields + provider_specific_fields = assistant_msg.get("provider_specific_fields") + thought_signatures = None + if provider_specific_fields and isinstance(provider_specific_fields, dict): + thought_signatures = provider_specific_fields.get("thought_signatures") + + # If we have thought signatures, add them to the part + if thought_signatures and isinstance(thought_signatures, list) and len(thought_signatures) > 0: + # Use the first signature for the text part (Gemini expects one signature per part) + assistant_content.append(PartType(text=assistant_text, thoughtSignature=thought_signatures[0])) # type: ignore + else: + assistant_content.append(PartType(text=assistant_text)) # type: ignore ## HANDLE ASSISTANT FUNCTION CALL if ( 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 a5cc3dca8c..35fa09e34f 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 @@ -1210,6 +1210,25 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): thinking_blocks.append(block) return thinking_blocks + def _extract_thought_signatures_from_parts( + self, parts: List[HttpxPartType] + ) -> Optional[List[str]]: + """Extract thoughtSignature values from parts. + + Per Google's docs, thoughtSignature is returned for multi-turn context preservation + and can appear on parts even without thought: true (e.g., regular text responses, + function calls). This method extracts all thoughtSignature values from parts. + + Returns: + List of thoughtSignature strings if any are found, None otherwise + """ + signatures: List[str] = [] + for part in parts: + signature = part.get("thoughtSignature") + if signature is not None: + signatures.append(signature) + return signatures if signatures else None + def _extract_image_response_from_parts( self, parts: List[HttpxPartType] ) -> Optional[List[ImageURLListItem]]: @@ -1620,6 +1639,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): from litellm.types.utils import Delta, StreamingChoices annotations = chat_completion_message.get("annotations") # type: ignore + provider_specific_fields = chat_completion_message.get("provider_specific_fields") # type: ignore # create a streaming choice object choice = StreamingChoices( finish_reason=VertexGeminiConfig._check_finish_reason( @@ -1633,6 +1653,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): images=image_response, function_call=functions, annotations=annotations, # type: ignore + provider_specific_fields=provider_specific_fields, ), logprobs=chat_completion_logprobs, enhancements=None, @@ -1811,6 +1832,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) ) + # Extract thoughtSignatures from parts (can exist without thought: true) + thought_signatures = ( + VertexGeminiConfig()._extract_thought_signatures_from_parts( + parts=candidate["content"]["parts"] + ) + ) + if audio_response is not None: cast(Dict[str, Any], chat_completion_message)[ "audio" @@ -1876,6 +1904,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): reasoning_content = "\n".join(reasoning_content_parts) chat_completion_message["reasoning_content"] = reasoning_content + # Store thoughtSignatures in provider_specific_fields + if thought_signatures is not None: + if "provider_specific_fields" not in chat_completion_message: + chat_completion_message["provider_specific_fields"] = {} + chat_completion_message["provider_specific_fields"]["thought_signatures"] = thought_signatures # type: ignore + if isinstance(model_response, ModelResponseStream): choice = VertexGeminiConfig._create_streaming_choice( chat_completion_message=chat_completion_message, diff --git a/tests/local_testing/test_gemini_reasoning_content.py b/tests/local_testing/test_gemini_reasoning_content.py index 758616b4c9..f1f9c2ab51 100644 --- a/tests/local_testing/test_gemini_reasoning_content.py +++ b/tests/local_testing/test_gemini_reasoning_content.py @@ -1,4 +1,5 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig +from litellm.llms.vertex_ai.gemini.transformation import _gemini_convert_messages_with_history def test_thought_true_creates_thinking_block(): @@ -36,3 +37,108 @@ def test_thought_signature_without_thought_does_not_create_block(): config = VertexGeminiConfig() thinking_blocks = config._extract_thinking_blocks_from_parts(parts) assert thinking_blocks == [] + + +def test_extract_thought_signatures_from_regular_parts(): + """ + Test that thoughtSignatures are extracted from regular text parts (without thought=True). + This is the key feature for Gemini 3 multi-turn context preservation. + """ + parts = [{"text": "I am Gemini", "thoughtSignature": "sig-regular-123"}] + config = VertexGeminiConfig() + + # Should NOT create thinking block + thinking_blocks = config._extract_thinking_blocks_from_parts(parts) + assert thinking_blocks == [] + + # Should extract thought signature + signatures = config._extract_thought_signatures_from_parts(parts) + assert signatures is not None + assert len(signatures) == 1 + assert signatures[0] == "sig-regular-123" + + +def test_extract_multiple_thought_signatures(): + """ + Test extraction of multiple thoughtSignatures from different parts. + """ + parts = [ + {"text": "Part 1", "thoughtSignature": "sig-1"}, + {"text": "Part 2", "thoughtSignature": "sig-2"}, + {"text": "Part 3"} # No signature + ] + config = VertexGeminiConfig() + signatures = config._extract_thought_signatures_from_parts(parts) + + assert signatures is not None + assert len(signatures) == 2 + assert signatures[0] == "sig-1" + assert signatures[1] == "sig-2" + + +def test_round_trip_thought_signature_in_conversation(): + """ + Test that thoughtSignatures are properly round-tripped through conversation history. + This ensures multi-turn context preservation works correctly. + """ + messages = [ + {"role": "user", "content": "Hello"}, + { + "role": "assistant", + "content": "Hi there", + "provider_specific_fields": { + "thought_signatures": ["sig-round-trip-abc"] + } + }, + {"role": "user", "content": "How are you?"} + ] + + gemini_contents = _gemini_convert_messages_with_history(messages) + + # Find the assistant (model) message + model_message = None + for content in gemini_contents: + if content.get("role") == "model": + model_message = content + break + + assert model_message is not None + assert len(model_message["parts"]) >= 1 + + # Check that the text part has the thoughtSignature + text_part = model_message["parts"][0] + assert text_part["text"] == "Hi there" + assert "thoughtSignature" in text_part + assert text_part["thoughtSignature"] == "sig-round-trip-abc" + + +def test_round_trip_without_thought_signature_still_works(): + """ + Test that messages without thoughtSignatures continue to work normally. + This ensures backward compatibility. + """ + messages = [ + {"role": "user", "content": "Hello"}, + { + "role": "assistant", + "content": "Hi there" + }, + {"role": "user", "content": "How are you?"} + ] + + gemini_contents = _gemini_convert_messages_with_history(messages) + + # Find the assistant (model) message + model_message = None + for content in gemini_contents: + if content.get("role") == "model": + model_message = content + break + + assert model_message is not None + assert len(model_message["parts"]) >= 1 + + # Check that the text part works without thoughtSignature + text_part = model_message["parts"][0] + assert text_part["text"] == "Hi there" + assert "thoughtSignature" not in text_part