Add thought signature support to v1/messages api (#16812)

* Add thought signature support to v1/messages api

* update the thinking level handling logic

* update the thinking level handling logic

* Add streaming support

* fix intalling litellm error

Author: Sameerlite

litellm/litellm_core_utils/prompt_templates/factory.py

@@ -1162,16 +1162,58 @@ def _gemini_tool_call_invoke_helper(
     return function_call
 
 
-def _get_thought_signature_from_tool(tool: dict) -> Optional[str]:
-    """Extract thought signature from tool call's provider_specific_fields"""
+def _get_thought_signature_from_tool(tool: dict, model: Optional[str] = None) -> Optional[str]:
+    """Extract thought signature from tool call's provider_specific_fields.
+    
+    Checks both tool.provider_specific_fields and tool.function.provider_specific_fields.
+    If no signature is found and model is gemini-3, returns a dummy signature.
+    """
+    # First check tool's provider_specific_fields
     provider_fields = tool.get("provider_specific_fields") or {}
     if isinstance(provider_fields, dict):
-        return provider_fields.get("thought_signature")
+        signature = provider_fields.get("thought_signature")
+        if signature:
+            return signature
+    
+    # Then check function's provider_specific_fields
+    function = tool.get("function")
+    if function:
+        if isinstance(function, dict):
+            func_provider_fields = function.get("provider_specific_fields") or {}
+            if isinstance(func_provider_fields, dict):
+                signature = func_provider_fields.get("thought_signature")
+                if signature:
+                    return signature
+        elif hasattr(function, "provider_specific_fields") and function.provider_specific_fields:
+            if isinstance(function.provider_specific_fields, dict):
+                signature = function.provider_specific_fields.get("thought_signature")
+                if signature:
+                    return signature
+    
+    # If no signature found and model is gemini-3, return dummy signature
+    from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
+    if model and VertexGeminiConfig._is_gemini_3_or_newer(model):
+        return _get_dummy_thought_signature()
+    
     return None
 
 
+def _get_dummy_thought_signature() -> str:
+    """Generate a dummy thought signature for models that require it.
+    
+    This is used when transferring conversation history from older models
+    (like gemini-2.5-flash) to gemini-3, which requires thought_signature
+    for strict validation.
+    """
+    # Return a base64-encoded dummy signature string
+    # Below dummy signature is recommended by google - https://ai.google.dev/gemini-api/docs/thought-signatures#faqs
+    dummy_data = b"skip_thought_signature_validator"
+    return base64.b64encode(dummy_data).decode("utf-8")
+
+
 def convert_to_gemini_tool_call_invoke(
     message: ChatCompletionAssistantMessage,
+    model: Optional[str] = None,
 ) -> List[VertexPartType]:
     """
     OpenAI tool invokes:
@@ -1229,7 +1271,7 @@ def convert_to_gemini_tool_call_invoke(
                         part_dict: VertexPartType = {
                             "function_call": gemini_function_call
                         }
-                        thought_signature = _get_thought_signature_from_tool(dict(tool))
+                        thought_signature = _get_thought_signature_from_tool(dict(tool), model=model)
                         if thought_signature:
                             part_dict["thoughtSignature"] = thought_signature
 
@@ -1250,11 +1292,18 @@ def convert_to_gemini_tool_call_invoke(
                 }
 
                 # Extract thought signature from function_call's provider_specific_fields
+                thought_signature = None
                 provider_fields = function_call.get("provider_specific_fields") if isinstance(function_call, dict) else {}
                 if isinstance(provider_fields, dict):
                     thought_signature = provider_fields.get("thought_signature")
-                    if thought_signature:
-                        part_dict_function["thoughtSignature"] = thought_signature
+                
+                # If no signature found and model is gemini-3, use dummy signature
+                from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
+                if not thought_signature and model and VertexGeminiConfig._is_gemini_3_or_newer(model):
+                    thought_signature = _get_dummy_thought_signature()
+                
+                if thought_signature:
+                    part_dict_function["thoughtSignature"] = thought_signature
 
                 _parts_list.append(part_dict_function)
             else:  # don't silently drop params. Make it clear to user what's happening.

litellm/litellm_core_utils/streaming_chunk_builder_utils.py

@@ -137,6 +137,7 @@ def get_combined_tool_content(
                             "name": None,
                             "type": None,
                             "arguments": [],
+                            "provider_specific_fields": None,
                         }
 
                     if hasattr(tool_call, "id") and tool_call.id:
@@ -156,22 +157,48 @@ def get_combined_tool_content(
                             tool_call_map[index]["arguments"].append(
                                 tool_call.function.arguments
                             )
+                    
+                    # Preserve provider_specific_fields from streaming chunks
+                    provider_fields = None
+                    if hasattr(tool_call, "provider_specific_fields") and tool_call.provider_specific_fields:
+                        provider_fields = tool_call.provider_specific_fields
+                    elif hasattr(tool_call, "function") and hasattr(tool_call.function, "provider_specific_fields") and tool_call.function.provider_specific_fields:
+                        provider_fields = tool_call.function.provider_specific_fields
+                    
+                    if provider_fields:
+                        # Merge provider_specific_fields if multiple chunks have them
+                        if tool_call_map[index]["provider_specific_fields"] is None:
+                            tool_call_map[index]["provider_specific_fields"] = {}
+                        if isinstance(provider_fields, dict):
+                            tool_call_map[index]["provider_specific_fields"].update(
+                                provider_fields
+                            )
 
         # Convert the map to a list of tool calls
         for index in sorted(tool_call_map.keys()):
             tool_call_data = tool_call_map[index]
             if tool_call_data["id"] and tool_call_data["name"]:
                 combined_arguments = "".join(tool_call_data["arguments"]) or "{}"
-                tool_calls_list.append(
-                    ChatCompletionMessageToolCall(
-                        id=tool_call_data["id"],
-                        function=Function(
-                            arguments=combined_arguments,
-                            name=tool_call_data["name"],
-                        ),
-                        type=tool_call_data["type"] or "function",
-                    )
+                
+                # Build function - provider_specific_fields should be on tool_call level, not function level
+                function = Function(
+                    arguments=combined_arguments,
+                    name=tool_call_data["name"],
                 )
+                
+                # Prepare params for ChatCompletionMessageToolCall
+                tool_call_params = {
+                    "id": tool_call_data["id"],
+                    "function": function,
+                    "type": tool_call_data["type"] or "function",
+                }
+                
+                # Add provider_specific_fields if present (for thought signatures in Gemini 3)
+                if tool_call_data.get("provider_specific_fields"):
+                    tool_call_params["provider_specific_fields"] = tool_call_data["provider_specific_fields"]
+                
+                tool_call = ChatCompletionMessageToolCall(**tool_call_params)
+                tool_calls_list.append(tool_call)
 
         return tool_calls_list
 

litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py

@@ -3,6 +3,7 @@
     TYPE_CHECKING,
     Any,
     AsyncIterator,
+    Dict,
     List,
     Literal,
     Optional,
@@ -129,6 +130,39 @@ def __init__(self):
 
     ### FOR [BETA] `/v1/messages` endpoint support
 
+    def _extract_signature_from_tool_call(
+        self, tool_call: Any
+    ) -> Optional[str]:
+        """
+        Extract signature from a tool call's provider_specific_fields.
+        Only checks provider_specific_fields, not thinking blocks.
+        """
+        signature = None
+        
+        if hasattr(tool_call, "provider_specific_fields") and tool_call.provider_specific_fields:
+            if "thought_signature" in tool_call.provider_specific_fields:
+                signature = tool_call.provider_specific_fields["thought_signature"]
+        elif (
+            hasattr(tool_call.function, "provider_specific_fields")
+            and tool_call.function.provider_specific_fields
+        ):
+            if "thought_signature" in tool_call.function.provider_specific_fields:
+                signature = tool_call.function.provider_specific_fields["thought_signature"]
+        
+        return signature
+
+    def _extract_signature_from_tool_use_content(
+        self, content: Dict[str, Any]
+    ) -> Optional[str]:
+        """
+        Extract signature from a tool_use content block's provider_specific_fields.
+        """
+        provider_specific_fields = content.get("provider_specific_fields", {})
+        if provider_specific_fields:
+            return provider_specific_fields.get("signature")
+        return None
+
+
     def translatable_anthropic_params(self) -> List:
         """
         Which anthropic params, we need to translate to the openai format.
@@ -263,10 +297,18 @@ def translate_anthropic_messages_to_openai(  # noqa: PLR0915
                                 else:
                                     assistant_message_str += content.get("text", "")
                             elif content.get("type") == "tool_use":
-                                function_chunk = ChatCompletionToolCallFunctionChunk(
-                                    name=content.get("name", ""),
-                                    arguments=json.dumps(content.get("input", {})),
-                                )
+                                function_chunk: ChatCompletionToolCallFunctionChunk = {
+                                    "name": content.get("name", ""),
+                                    "arguments": json.dumps(content.get("input", {})),
+                                }
+                                signature = self._extract_signature_from_tool_use_content(content)
+                                
+                                if signature:
+                                    provider_specific_fields: Dict[str, Any] = (
+                                        function_chunk.get("provider_specific_fields") or {}
+                                    )
+                                    provider_specific_fields["thought_signature"] = signature
+                                    function_chunk["provider_specific_fields"] = provider_specific_fields
 
                                 tool_calls.append(
                                     ChatCompletionAssistantToolCall(
@@ -512,18 +554,27 @@ def _translate_openai_content_to_anthropic(self, choices: List[Choices]) -> List
                 and len(choice.message.tool_calls) > 0
             ):
                 for tool_call in choice.message.tool_calls:
-                    new_content.append(
-                        AnthropicResponseContentBlockToolUse(
-                            type="tool_use",
-                            id=tool_call.id,
-                            name=tool_call.function.name or "",
-                            input=(
-                                json.loads(tool_call.function.arguments)
-                                if tool_call.function.arguments
-                                else {}
-                            ),
-                        )
+                    # Extract signature from provider_specific_fields only
+                    signature = self._extract_signature_from_tool_call(tool_call)
+                    
+                    provider_specific_fields = {}
+                    if signature:
+                        provider_specific_fields["signature"] = signature
+                    
+                    tool_use_block = AnthropicResponseContentBlockToolUse(
+                        type="tool_use",
+                        id=tool_call.id,
+                        name=tool_call.function.name or "",
+                        input=(
+                            json.loads(tool_call.function.arguments)
+                            if tool_call.function.arguments
+                            else {}
+                        ),
                     )
+                    # Add provider_specific_fields if signature is present
+                    if provider_specific_fields:
+                        tool_use_block.provider_specific_fields = provider_specific_fields
+                    new_content.append(tool_use_block)
             # Handle text content
             elif choice.message.content is not None:
                 new_content.append(

litellm/llms/gemini/chat/transformation.py

@@ -140,4 +140,4 @@ def _transform_messages(
                             except Exception:
                                 # If conversion fails, leave as is and let the API handle it
                                 pass
-        return _gemini_convert_messages_with_history(messages=messages)
+        return _gemini_convert_messages_with_history(messages=messages, model=model)

litellm/llms/vertex_ai/context_caching/transformation.py

@@ -173,7 +173,7 @@ def transform_openai_messages_to_gemini_context_caching(
         supports_system_message=supports_system_message, messages=messages
     )
 
-    transformed_messages = _gemini_convert_messages_with_history(messages=new_messages)
+    transformed_messages = _gemini_convert_messages_with_history(messages=new_messages, model=model)
 
     model_name = "models/{}".format(model)
 

litellm/llms/vertex_ai/gemini/transformation.py

@@ -195,6 +195,7 @@ def check_if_part_exists_in_parts(
 
 def _gemini_convert_messages_with_history(  # noqa: PLR0915
     messages: List[AllMessageValues],
+    model: Optional[str] = None,
 ) -> List[ContentType]:
     """
     Converts given messages from OpenAI format to Gemini format
@@ -379,7 +380,7 @@ def _gemini_convert_messages_with_history(  # noqa: PLR0915
                     or assistant_msg.get("function_call") is not None
                 ):  # support assistant tool invoke conversion
                     gemini_tool_call_parts = convert_to_gemini_tool_call_invoke(
-                        assistant_msg
+                        assistant_msg, model=model
                     )
                     ## check if gemini_tool_call already exists in assistant_content
                     for gemini_tool_call_part in gemini_tool_call_parts:

litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py

@@ -898,11 +898,11 @@ def map_openai_params(  # noqa: PLR0915
         if VertexGeminiConfig._is_gemini_3_or_newer(model):
             if "temperature" not in optional_params:
                 optional_params["temperature"] = 1.0
+            thinking_config = optional_params.get("thinkingConfig", {})
             if (
-                "thinkingConfig" not in optional_params
-                or "thinkingLevel" not in optional_params.get("thinkingConfig", {})
+                "thinkingLevel" not in thinking_config
+                and "thinkingBudget" not in thinking_config
             ):
-                thinking_config = optional_params.get("thinkingConfig", {})
                 thinking_config["thinkingLevel"] = "low"
                 optional_params["thinkingConfig"] = thinking_config
 
@@ -1892,7 +1892,7 @@ def _transform_google_generate_content_to_openai_model_response(
     def _transform_messages(
         self, messages: List[AllMessageValues], model: Optional[str] = None
     ) -> List[ContentType]:
-        return _gemini_convert_messages_with_history(messages=messages)
+        return _gemini_convert_messages_with_history(messages=messages, model=model)
 
     def get_error_class(
         self, error_message: str, status_code: int, headers: Union[Dict, httpx.Headers]