Preserve Gemini thought signatures through streamed tool calls

src/eva/assistant/agentic/system.py

@@ -284,13 +284,19 @@ async def _run_tool_loop(
                     # Apply tool name cleaning for Harmony token leak bug
                     tc_dict["function"]["name"] = _clean_tool_name(tc_dict["function"]["name"])
 
-                    # Log if provider_specific_fields are present (e.g., Gemini thought signatures)
-                    if "provider_specific_fields" in tc_dict:
-                        fields = tc_dict["provider_specific_fields"]
-                        if "thought_signature" in fields:
-                            logger.info(
-                                "🔮 Gemini thought signature present in tool call (will be preserved for next turn)"
-                            )
+                    # Log if a Gemini thought signature is present, so it's visible that it will be
+                    # preserved for the next turn. LiteLLM-routed calls (CASCADE) surface it under
+                    # provider_specific_fields; ALMGeminiClient's raw OpenAI-compat calls to Gemini
+                    # surface it under extra_content (see alm_base._merge_streamed_tool_call_extras).
+                    provider_fields = tc_dict.get("provider_specific_fields") or {}
+                    extra_content = tc_dict.get("extra_content") or {}
+                    if provider_fields.get("thought_signature") or (
+                        isinstance(extra_content.get("google"), dict)
+                        and extra_content["google"].get("thought_signature")
+                    ):
+                        logger.info(
+                            "🔮 Gemini thought signature present in tool call (will be preserved for next turn)"
+                        )
 
                     tool_calls_dicts.append(tc_dict)
 

src/eva/assistant/pipeline/alm_base.py

@@ -102,6 +102,72 @@ def resample_pcm16(pcm_data: bytes, from_rate: int, to_rate: int) -> bytes:
     return struct.pack(f"<{len(out_samples)}h", *out_samples)
 
 
+# Keys litellm.stream_chunk_builder itself understands when merging a streamed
+# tool_call delta (see get_combined_tool_content in
+# litellm/litellm_core_utils/streaming_chunk_builder_utils.py). Anything else a
+# provider puts on the delta gets silently dropped during reassembly. This bites
+# Gemini's raw OpenAI-compatible endpoint (used directly by ALMGeminiClient,
+# bypassing litellm's own Gemini transport): it carries the thought signature
+# required for multi-turn function calling under an `extra_content` key that
+# litellm's merge doesn't recognize. CASCADE's LiteLLMClient doesn't hit this
+# because it talks to Gemini through litellm's native transport, which stamps
+# signatures onto `provider_specific_fields` itself before chunks reach here.
+_KNOWN_STREAMED_TOOL_CALL_KEYS = {"id", "type", "function", "index", "provider_specific_fields"}
+
+
+def _merge_streamed_tool_call_extras(dict_chunks: list[dict]) -> list[Any] | None:
+    """Rebuild tool_calls from raw stream chunks, preserving extra keys.
+
+    E.g. Gemini's `extra_content`, that litellm.stream_chunk_builder's merge drops.
+    Returns None if no such extra keys are present anywhere in the chunks (the
+    common case for other providers), so the caller can leave litellm's own
+    reconstruction untouched.
+    """
+    tool_call_map: dict[int, dict[str, Any]] = {}
+    for chunk in dict_chunks:
+        for choice in chunk.get("choices") or []:
+            delta = choice.get("delta") or {}
+            for tc in delta.get("tool_calls") or []:
+                if not isinstance(tc, dict):
+                    continue
+                index = tc.get("index", 0)
+                entry = tool_call_map.setdefault(
+                    index, {"id": None, "type": None, "name": None, "arguments": [], "extra": {}}
+                )
+                if tc.get("id"):
+                    entry["id"] = tc["id"]
+                if tc.get("type"):
+                    entry["type"] = tc["type"]
+                function = tc.get("function") or {}
+                if function.get("name"):
+                    entry["name"] = function["name"]
+                if function.get("arguments"):
+                    entry["arguments"].append(function["arguments"])
+                for key, value in tc.items():
+                    if key not in _KNOWN_STREAMED_TOOL_CALL_KEYS and value is not None:
+                        entry["extra"][key] = value
+
+    if not any(entry["extra"] for entry in tool_call_map.values()):
+        return None
+
+    from litellm.types.utils import ChatCompletionMessageToolCall, Function
+
+    tool_calls = []
+    for index in sorted(tool_call_map):
+        entry = tool_call_map[index]
+        if not (entry["id"] and entry["name"]):
+            continue
+        tc_obj = ChatCompletionMessageToolCall(
+            id=entry["id"],
+            type=entry["type"] or "function",
+            function=Function(name=entry["name"], arguments="".join(entry["arguments"]) or "{}"),
+        )
+        for key, value in entry["extra"].items():
+            setattr(tc_obj, key, value)
+        tool_calls.append(tc_obj)
+    return tool_calls or None
+
+
 def _assemble_stream_chunks(chunks: list, messages: list[dict[str, Any]]) -> tuple[Any, Any, str]:
     """Reconstruct the final message from raw OpenAI-SDK stream chunks.
 
@@ -117,6 +183,13 @@ def _assemble_stream_chunks(chunks: list, messages: list[dict[str, Any]]) -> tup
     dict_chunks = [c.model_dump() if hasattr(c, "model_dump") else c for c in chunks]
     full = litellm.stream_chunk_builder(dict_chunks, messages=messages)
     message = full.choices[0].message
+    # See _merge_streamed_tool_call_extras: litellm's own merge drops provider-
+    # specific extra keys (e.g. Gemini's extra_content.thought_signature).
+    # Rebuild tool_calls ourselves when present so the field survives into the
+    # message handed back to AgenticSystem._run_tool_loop.
+    merged_tool_calls = _merge_streamed_tool_call_extras(dict_chunks)
+    if merged_tool_calls is not None:
+        message.tool_calls = merged_tool_calls
     usage = getattr(full, "usage", None)
     finish_reason = getattr(full.choices[0], "finish_reason", None) or "unknown"
     return message, usage, finish_reason

tests/unit/assistant/test_alm_base.py

@@ -0,0 +1,172 @@
+"""Tests for stream-chunk reassembly in alm_base.py.
+
+Gemini's raw OpenAI-compatible endpoint (hit directly by ALMGeminiClient, bypassing
+litellm's own Gemini transport) carries the thought signature required for multi-turn
+function calling under an `extra_content` key on each tool_call delta. litellm's
+stream_chunk_builder only preserves id/type/function/provider_specific_fields when
+merging streamed tool_call deltas, so that key is silently dropped unless we re-merge it
+ourselves -- see _merge_streamed_tool_call_extras.
+"""
+
+from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
+
+from eva.assistant.pipeline.alm_base import _assemble_stream_chunks, _merge_streamed_tool_call_extras
+
+
+def _make_chunk(delta: dict, finish_reason: str | None = None, usage: dict | None = None) -> ChatCompletionChunk:
+    payload = {
+        "id": "chatcmpl-1",
+        "object": "chat.completion.chunk",
+        "created": 1234567890,
+        "model": "gemini-3.6-flash",
+        "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
+    }
+    if usage:
+        payload["usage"] = usage
+    return ChatCompletionChunk.model_validate(payload)
+
+
+def test_assemble_stream_chunks_preserves_gemini_thought_signature():
+    """extra_content.google.thought_signature must survive stream reassembly."""
+    chunks = [
+        _make_chunk(
+            {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_1",
+                        "type": "function",
+                        "function": {"name": "get_flight_status", "arguments": ""},
+                        "extra_content": {"google": {"thought_signature": "sig123"}},
+                    }
+                ],
+            }
+        ),
+        _make_chunk({"tool_calls": [{"index": 0, "function": {"arguments": '{"flight": "AA1"}'}}]}),
+        _make_chunk(
+            {},
+            finish_reason="tool_calls",
+            usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
+        ),
+    ]
+
+    message, _usage, finish_reason = _assemble_stream_chunks(chunks, messages=[{"role": "user", "content": "hi"}])
+
+    assert finish_reason == "tool_calls"
+    tc_dict = message.tool_calls[0].model_dump(exclude_none=True)
+    assert tc_dict["function"]["name"] == "get_flight_status"
+    assert tc_dict["function"]["arguments"] == '{"flight": "AA1"}'
+    assert tc_dict["extra_content"]["google"]["thought_signature"] == "sig123"
+
+
+def test_assemble_stream_chunks_unaffected_when_no_extra_fields():
+    """Providers with no unrecognized delta keys (e.g. vLLM) get litellm's own object untouched."""
+    chunks = [
+        _make_chunk(
+            {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_1",
+                        "type": "function",
+                        "function": {"name": "get_weather", "arguments": ""},
+                    }
+                ],
+            }
+        ),
+        _make_chunk({"tool_calls": [{"index": 0, "function": {"arguments": "{}"}}]}),
+        _make_chunk(
+            {},
+            finish_reason="tool_calls",
+            usage={"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8},
+        ),
+    ]
+
+    message, _usage, _finish_reason = _assemble_stream_chunks(chunks, messages=[{"role": "user", "content": "hi"}])
+
+    tc_dict = message.tool_calls[0].model_dump(exclude_none=True)
+    assert tc_dict == {
+        "id": "call_1",
+        "type": "function",
+        "function": {"name": "get_weather", "arguments": "{}"},
+    }
+
+
+def test_merge_streamed_tool_call_extras_returns_none_without_extra_keys():
+    dict_chunks = [
+        {
+            "choices": [
+                {
+                    "delta": {
+                        "tool_calls": [
+                            {
+                                "index": 0,
+                                "id": "call_1",
+                                "type": "function",
+                                "function": {"name": "f", "arguments": "{}"},
+                            }
+                        ]
+                    }
+                }
+            ]
+        }
+    ]
+    assert _merge_streamed_tool_call_extras(dict_chunks) is None
+
+
+def test_merge_streamed_tool_call_extras_handles_multiple_tool_calls():
+    """Two parallel tool calls, each streamed across multiple chunks, each with its own signature."""
+    dict_chunks = [
+        {
+            "choices": [
+                {
+                    "delta": {
+                        "tool_calls": [
+                            {
+                                "index": 0,
+                                "id": "call_1",
+                                "type": "function",
+                                "function": {"name": "get_flight_status", "arguments": ""},
+                                "extra_content": {"google": {"thought_signature": "sig_a"}},
+                            },
+                            {
+                                "index": 1,
+                                "id": "call_2",
+                                "type": "function",
+                                "function": {"name": "get_weather", "arguments": ""},
+                                "extra_content": {"google": {"thought_signature": "sig_b"}},
+                            },
+                        ]
+                    }
+                }
+            ]
+        },
+        {
+            "choices": [
+                {
+                    "delta": {
+                        "tool_calls": [
+                            {"index": 0, "function": {"arguments": '{"flight": "AA1"}'}},
+                            {"index": 1, "function": {"arguments": '{"city": "SF"}'}},
+                        ]
+                    }
+                }
+            ]
+        },
+    ]
+
+    tool_calls = _merge_streamed_tool_call_extras(dict_chunks)
+
+    assert tool_calls is not None
+    assert len(tool_calls) == 2
+    dumped = [tc.model_dump(exclude_none=True) for tc in tool_calls]
+    assert dumped[0]["id"] == "call_1"
+    assert dumped[0]["function"]["arguments"] == '{"flight": "AA1"}'
+    assert dumped[0]["extra_content"]["google"]["thought_signature"] == "sig_a"
+    assert dumped[1]["id"] == "call_2"
+    assert dumped[1]["function"]["arguments"] == '{"city": "SF"}'
+    assert dumped[1]["extra_content"]["google"]["thought_signature"] == "sig_b"