diff --git a/src/eva/assistant/agentic/system.py b/src/eva/assistant/agentic/system.py index 0ad2b652..9c6fdc15 100644 --- a/src/eva/assistant/agentic/system.py +++ b/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) diff --git a/src/eva/assistant/pipeline/alm_base.py b/src/eva/assistant/pipeline/alm_base.py index ddfcbe8f..be5d3126 100644 --- a/src/eva/assistant/pipeline/alm_base.py +++ b/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 diff --git a/tests/unit/assistant/test_alm_base.py b/tests/unit/assistant/test_alm_base.py new file mode 100644 index 00000000..af711fe4 --- /dev/null +++ b/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"