From 987a8fcf48b0a8de787b4405290fb8f4cd3c108b Mon Sep 17 00:00:00 2001 From: mgeorgaklis Date: Wed, 29 Jul 2026 03:31:54 +0000 Subject: [PATCH 1/2] fix(gemini): do not send duplicate thoughtSignature copies to Gemini Gemini returns each thoughtSignature on exactly one part. LiteLLM stores a function-call signature both message-level (thought_signatures) and on the tool call itself, then re-attached it to BOTH the text part and the function-call part when serializing history. gemini-3 and newer models bill every replayed copy as the previous turn's full reasoning token count, so long agentic sessions doubled their context growth and hit the 1,048,576-token limit Only attach a message-level signature to the text part when the same signature is not already carried by a tool-call part: - compare signature values instead of boolean presence so a distinct text-part signature is never dropped - ignore the gemini-3 dummy-signature fallback during detection so replaying gemini-2.5 history to a newer model keeps the real text signature - count signatures carried by server-side tool invocations so they are not re-attached to the text part gemini-2.5 responses (signature on the text part, function call unsigned) are unaffected: the text signature is preserved as before --- .../llms/vertex_ai/gemini/transformation.py | 61 +++- .../test_vertex_ai_gemini_transformation.py | 316 ++++++++++++++++++ 2 files changed, 375 insertions(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 0db1118a7b4..f465a54b265 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -20,6 +20,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( _get_image_mime_type_from_url, ) from litellm.litellm_core_utils.prompt_templates.factory import ( + _get_thought_signature_from_tool, convert_generic_image_chunk_to_openai_image_obj, convert_to_anthropic_image_obj, convert_to_gemini_tool_call_invoke, @@ -635,6 +636,52 @@ def check_if_part_exists_in_parts(parts: List[PartType], part: PartType, exclude return False +def _collect_tool_call_thought_signatures( + assistant_msg: ChatCompletionAssistantMessage, +) -> frozenset[str]: + """Thought signatures already carried by this message's tool-call parts. + + Gemini returns each thoughtSignature on exactly one part. When the signed + part is a function call, the signature is replayed on that tool-call part + by convert_to_gemini_tool_call_invoke, so attaching the same signature to + the text part as well would send two copies and double-bill the previous + turn's reasoning tokens on gemini-3 and newer models. + + Detection deliberately calls _get_thought_signature_from_tool without the + model argument: with a gemini-3 model that helper synthesizes a dummy + signature for unsigned tool calls, which must not suppress a real + text-part signature (e.g. replaying gemini-2.5 history to a newer model). + """ + signatures: tuple[str, ...] = () + + tool_calls = assistant_msg.get("tool_calls") + if isinstance(tool_calls, list): + for tool in tool_calls: + if isinstance(tool, dict): + signature = _get_thought_signature_from_tool(tool) + if signature: + signatures += (signature,) + + function_call = assistant_msg.get("function_call") + if isinstance(function_call, dict): + signature = _get_thought_signature_from_tool({"function": function_call}) + if signature: + signatures += (signature,) + + provider_specific_fields = assistant_msg.get("provider_specific_fields") + if isinstance(provider_specific_fields, dict): + invocations = provider_specific_fields.get("server_side_tool_invocations") + if isinstance(invocations, list): + for invocation in invocations: + if isinstance(invocation, dict): + for key in ("thought_signature", "response_thought_signature"): + invocation_signature = invocation.get(key) + if isinstance(invocation_signature, str) and invocation_signature: + signatures += (invocation_signature,) + + return frozenset(signatures) + + def _gemini_convert_messages_with_history( messages: List[AllMessageValues], model: Optional[str] = None, @@ -864,8 +911,18 @@ def _gemini_convert_messages_with_history( 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: + # A signature that is already carried by one of this message's + # tool-call parts must not be attached to the text part too: + # Gemini bills every replayed copy as the previous turn's full + # reasoning token count on gemini-3 and newer models + tool_call_signatures = _collect_tool_call_thought_signatures(assistant_msg) + + if ( + thought_signatures + and isinstance(thought_signatures, list) + and len(thought_signatures) > 0 + and thought_signatures[0] not in tool_call_signatures + ): # Use the first signature for the text part (Gemini expects one signature per part) assistant_content.append( PartType( diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index d99c190c6e5..8ee8186f6bb 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -2097,3 +2097,319 @@ def test_multi_turn_function_calling_roles(): assert ( content["role"] == "user" ), f"Content block {i} with function_response has role='{content['role']}', expected 'user'" + + +def test_gemini_thought_signature_preservation_real_response(): + """Test that thought signatures are preserved on the text part if originally there, without dropping or duplicating (real response case).""" + 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, + ) + + real_candidate = { + "content": { + "parts": [ + { + "text": "I will explain and then list files.", + "thoughtSignature": "mock_signature_from_text_part", + }, + { + "functionCall": { + "name": "list_files", + "args": {}, + } + }, + ] + } + } + + parts = real_candidate["content"]["parts"] + + content, reasoning_content = ( + VertexGeminiConfig().get_assistant_content_message(parts=parts) + ) + thought_signatures = ( + VertexGeminiConfig()._extract_thought_signatures_from_parts( + parts=parts + ) + ) + functions, tools, _ = VertexGeminiConfig._transform_parts( + parts=parts, + cumulative_tool_call_idx=0, + is_function_call=False, + ) + + msg: dict = {"role": "assistant"} + if content is not None: + msg["content"] = content + if tools: + msg["tool_calls"] = tools + if functions is not None: + msg["function_call"] = functions + if thought_signatures is not None: + msg["provider_specific_fields"] = { + "thought_signatures": thought_signatures + } + + converted_real = _gemini_convert_messages_with_history( + messages=[msg], + model="gemini-2.5-pro", + ) + + assert len(converted_real) == 1 + assert "parts" in converted_real[0] + parts_out = converted_real[0]["parts"] + assert len(parts_out) == 2 + assert "text" in parts_out[0] + assert ( + parts_out[0]["thoughtSignature"] == "mock_signature_from_text_part" + ) + assert "function_call" in parts_out[1] + assert "thoughtSignature" not in parts_out[1] + + +def test_gemini_thought_signature_deduplication_assumed_response(): + """Test that thought signatures are deduplicated and not attached to the text part if already present in the tool call (assumed response case).""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + pr_assumed_msg = { + "role": "assistant", + "content": "I will list the directory.", + "provider_specific_fields": { + "thought_signatures": ["mock_signature_63k"] + }, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "list_files", "arguments": "{}"}, + "provider_specific_fields": { + "thought_signature": "mock_signature_63k" + }, + } + ], + } + + converted_pr = _gemini_convert_messages_with_history( + messages=[pr_assumed_msg], + model="gemini-2.5-pro", + ) + + assert len(converted_pr) == 1 + assert "parts" in converted_pr[0] + parts_out = converted_pr[0]["parts"] + assert len(parts_out) == 2 + assert "text" in parts_out[0] + assert "thoughtSignature" not in parts_out[0] + assert "function_call" in parts_out[1] + assert parts_out[1]["thoughtSignature"] == "mock_signature_63k" + + +def test_gemini_thought_signature_pure_text(): + """Test that thought signatures are preserved on the text part for responses with no tool calls.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "Hello, I am a model.", + "provider_specific_fields": { + "thought_signatures": ["pure_text_signature"] + }, + } + + converted = _gemini_convert_messages_with_history( + messages=[msg], + model="gemini-2.5-pro", + ) + + assert len(converted) == 1 + assert "parts" in converted[0] + parts_out = converted[0]["parts"] + assert len(parts_out) == 1 + assert "text" in parts_out[0] + assert parts_out[0]["thoughtSignature"] == "pure_text_signature" + + +def test_gemini_thought_signature_pure_tool_call(): + """Test that thought signatures are preserved on the tool call for responses with no intermediate text.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": None, + "provider_specific_fields": { + "thought_signatures": ["pure_tool_signature"] + }, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "list_files", "arguments": "{}"}, + "provider_specific_fields": { + "thought_signature": "pure_tool_signature" + }, + } + ], + } + + converted = _gemini_convert_messages_with_history( + messages=[msg], + model="gemini-2.5-pro", + ) + + assert len(converted) == 1 + assert "parts" in converted[0] + parts_out = converted[0]["parts"] + assert len(parts_out) == 1 + assert "function_call" in parts_out[0] + assert parts_out[0]["thoughtSignature"] == "pure_tool_signature" + + +def test_gemini_distinct_text_and_tool_signatures_are_both_preserved(): + """A text-part signature that differs from the tool-call signature must stay on the text part.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "Some analysis.", + "provider_specific_fields": { + "thought_signatures": ["text_signature", "tool_signature"] + }, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "list_files", "arguments": "{}"}, + "provider_specific_fields": {"thought_signature": "tool_signature"}, + } + ], + } + + parts = _gemini_convert_messages_with_history( + messages=[msg], model="gemini-2.5-pro" + )[0]["parts"] + + assert parts[0]["text"] == "Some analysis." + assert parts[0]["thoughtSignature"] == "text_signature" + assert "function_call" in parts[1] + assert parts[1]["thoughtSignature"] == "tool_signature" + + +def test_gemini_25_text_signature_survives_replay_to_gemini_3(): + """gemini-2.5 history (signed text, unsigned tool call) replayed to gemini-3 keeps the real + text signature; the dummy signature synthesized for the unsigned tool call must not suppress it.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + _get_dummy_thought_signature, + ) + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "I will list the directory.", + "provider_specific_fields": {"thought_signatures": ["real_25_signature"]}, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "list_files", "arguments": "{}"}, + } + ], + } + + parts = _gemini_convert_messages_with_history(messages=[msg], model="gemini-3-pro")[ + 0 + ]["parts"] + + assert parts[0]["text"] == "I will list the directory." + assert parts[0]["thoughtSignature"] == "real_25_signature" + assert "function_call" in parts[1] + assert parts[1]["thoughtSignature"] == _get_dummy_thought_signature() + + +def test_gemini_function_call_signature_round_trip_no_duplicate(): + """End to end: a gemini-3-style response (unsigned text + signed functionCall) parsed and + re-serialized sends the signature exactly once, on the function-call part.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + response_parts = [ + {"text": "I will calculate the result for you."}, + { + "functionCall": {"name": "add_numbers", "args": {"a": 17, "b": 25}}, + "thoughtSignature": "signature_from_function_call", + }, + ] + + config = VertexGeminiConfig() + content, _ = config.get_assistant_content_message(parts=response_parts) + thought_signatures = config._extract_thought_signatures_from_parts( + parts=response_parts + ) + _, tools, _ = VertexGeminiConfig._transform_parts( + parts=response_parts, cumulative_tool_call_idx=0, is_function_call=False + ) + + msg = { + "role": "assistant", + "content": content, + "tool_calls": tools, + "provider_specific_fields": {"thought_signatures": thought_signatures}, + } + + parts = _gemini_convert_messages_with_history(messages=[msg], model="gemini-3-pro")[ + 0 + ]["parts"] + + signatures = [p["thoughtSignature"] for p in parts if "thoughtSignature" in p] + assert signatures == ["signature_from_function_call"] + assert "thoughtSignature" not in parts[0] + assert "function_call" in parts[1] + + +def test_gemini_server_side_tool_signature_not_duplicated_on_text(): + """A signature already re-injected on a server-side toolCall part is not attached to the text part again.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "The weather in Buenos Aires is sunny.", + "provider_specific_fields": { + "thought_signatures": ["server_side_signature"], + "server_side_tool_invocations": [ + { + "tool_type": "GOOGLE_SEARCH_WEB", + "id": "abc123", + "args": {"queries": ["weather Buenos Aires"]}, + "response": {"weather": "Sunny"}, + "thought_signature": "server_side_signature", + } + ], + }, + } + + parts = _gemini_convert_messages_with_history( + messages=[msg], model="gemini-2.5-pro" + )[0]["parts"] + + text_part = next(p for p in parts if "text" in p) + assert "thoughtSignature" not in text_part + tool_call_part = next(p for p in parts if "toolCall" in p) + assert tool_call_part["thoughtSignature"] == "server_side_signature" From 67969303fceb053184039c6792131a77f09df4f5 Mon Sep 17 00:00:00 2001 From: mgeorgaklis Date: Fri, 31 Jul 2026 14:48:40 +0000 Subject: [PATCH 2/2] refactor(gemini): simplify thought signature collection --- .../llms/vertex_ai/gemini/transformation.py | 32 +++++++++++-------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index f465a54b265..1f568496041 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -657,10 +657,11 @@ def _collect_tool_call_thought_signatures( tool_calls = assistant_msg.get("tool_calls") if isinstance(tool_calls, list): for tool in tool_calls: - if isinstance(tool, dict): - signature = _get_thought_signature_from_tool(tool) - if signature: - signatures += (signature,) + if not isinstance(tool, dict): + continue + signature = _get_thought_signature_from_tool(tool) + if signature: + signatures += (signature,) function_call = assistant_msg.get("function_call") if isinstance(function_call, dict): @@ -669,15 +670,20 @@ def _collect_tool_call_thought_signatures( signatures += (signature,) provider_specific_fields = assistant_msg.get("provider_specific_fields") - if isinstance(provider_specific_fields, dict): - invocations = provider_specific_fields.get("server_side_tool_invocations") - if isinstance(invocations, list): - for invocation in invocations: - if isinstance(invocation, dict): - for key in ("thought_signature", "response_thought_signature"): - invocation_signature = invocation.get(key) - if isinstance(invocation_signature, str) and invocation_signature: - signatures += (invocation_signature,) + if not isinstance(provider_specific_fields, dict): + return frozenset(signatures) + + invocations = provider_specific_fields.get("server_side_tool_invocations") + if not isinstance(invocations, list): + return frozenset(signatures) + + for invocation in invocations: + if not isinstance(invocation, dict): + continue + for key in ("thought_signature", "response_thought_signature"): + invocation_signature = invocation.get(key) + if isinstance(invocation_signature, str) and invocation_signature: + signatures += (invocation_signature,) return frozenset(signatures)