From 504cd98fd168e4e0f87b7dd467941b249ad74b97 Mon Sep 17 00:00:00 2001
From: Bahtya <bahtayr@gmail.com>
Date: Thu, 9 Apr 2026 18:31:09 +0800
Subject: [PATCH] fix(gemini): clarify thought-signature design; add
 sibling-thought-part test
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Addresses greptile P1 + P2 review comments:

P1 — Clarify that _thought_signatures collects ONLY from thought parts
(thought==True) and that a functionCall part's own thoughtSignature is
consumed separately and never bleeds into sibling calls. Added a more
detailed comment in _transform_parts explaining the design.

P2 — Added test_sibling_thought_part_propagates_to_all_parallel_calls:
covers the regression scenario where Gemini returns a separate thought
part before multiple parallel functionCall parts, and both function calls
must receive the thought signature so multi-turn context is preserved.

Local test run: pytest tests/test_litellm/llms/vertex_ai/gemini/
test_thought_signature_in_tool_call_id.py → 12 passed.

Bahtya
---
 .../vertex_and_google_ai_studio_gemini.py     | 12 +++--
 .../test_thought_signature_in_tool_call_id.py | 52 +++++++++++++++++++
 2 files changed, 60 insertions(+), 4 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 b49cb58decc..40da10d6c01 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
@@ -1585,10 +1585,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
     ]:
         function: Optional[ChatCompletionToolCallFunctionChunk] = None
         _tools: List[ChatCompletionToolCallChunk] = []
-        # Collect thoughtSignature values only from thought parts (thought: true).
-        # Gemini may place thoughtSignature on a separate thought part rather than
-        # the functionCall part. We exclude functionCall parts to avoid bleeding
-        # signatures across parallel function calls.
+        # Collect thoughtSignature values ONLY from thought parts (thought: true).
+        # Gemini may return: [{thought: true, thoughtSignature: "..."}, {functionCall: A}, ...]
+        # In this layout the signature belongs to ALL function calls in the turn (it
+        # encodes the model's reasoning context for multi-turn preservation), so we
+        # propagate it to every parallel call via the fallback below.
+        # We intentionally exclude functionCall parts here — a signature on a
+        # functionCall part is already consumed via `part.get("thoughtSignature")`
+        # in the loop and must NOT bleed into sibling calls.
         _thought_signatures = [
             p.get("thoughtSignature")
             for p in parts
diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py b/tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py
index 208cba519f3..78d4ba91cab 100644
--- a/tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py
+++ b/tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py
@@ -285,3 +285,55 @@ def test_parallel_tool_calls_with_signatures(enable_preview_features):
     assert THOUGHT_SIGNATURE_SEPARATOR not in tools[1]["id"]
     sig2 = _get_thought_signature_from_tool({"id": tools[1]["id"], "type": "function"})
     assert sig2 is None
+
+
+def test_sibling_thought_part_propagates_to_all_parallel_calls() -> None:
+    """
+    Regression test for the scenario this PR addresses.
+
+    Gemini may return a structure like:
+        [{thought: true, thoughtSignature: "..."}, {functionCall: A}, {functionCall: B}]
+
+    The thoughtSignature is on a *sibling* thought part, not on the functionCall parts.
+    It represents the model's reasoning for the entire turn and should be propagated to
+    ALL parallel function calls so that multi-turn context is preserved on the next turn.
+
+    This also verifies that a functionCall part's OWN thoughtSignature (a separate scenario
+    where Gemini puts the signature directly on the functionCall) does NOT bleed into
+    sibling calls — _thought_signatures only collects from parts where thought==True,
+    so functionCall-level signatures stay scoped to their own call.
+    """
+    thought_signature = "ErcN4Ad4qImqBhBo8LkVzGuqK7Y="
+
+    # Gemini response: one thought part with signature, two parallel function calls
+    parts_sibling_thought = [
+        HttpxPartType(
+            thought=True,
+            thoughtSignature=thought_signature,
+        ),
+        HttpxPartType(
+            functionCall={"name": "get_temperature", "args": {"location": "Paris"}},
+        ),
+        HttpxPartType(
+            functionCall={"name": "get_temperature", "args": {"location": "London"}},
+        ),
+    ]
+
+    function, tools, _ = VertexGeminiConfig._transform_parts(
+        parts=parts_sibling_thought,
+        cumulative_tool_call_idx=0,
+        is_function_call=False,
+    )
+
+    assert tools is not None
+    assert len(tools) == 2
+
+    # Both function calls must carry the thought signature from the sibling thought part.
+    for i, tool in enumerate(tools):
+        sig = tool.get("provider_specific_fields", {}).get("thought_signature")
+        assert sig == thought_signature, (
+            f"tool[{i}] expected thought_signature={thought_signature!r}, got {sig!r}"
+        )
+        assert THOUGHT_SIGNATURE_SEPARATOR in tool["id"], (
+            f"tool[{i}] thought signature should be embedded in ID"
+        )