From b63e23d259431b5bfb71962d51d72a157142a38f Mon Sep 17 00:00:00 2001 From: Vetle Leinonen-Roeim Date: Thu, 8 Oct 2026 09:32:40 +0200 Subject: [PATCH] fix(litellm): keep a streamed thought signature sent without the call's name A tool-call delta carrying only a signature was dropped as empty, so a provider that sends the signature apart from the name and arguments would still lose it. Keep such deltas, attach the signature to the call at the delta's own index, or hold it for the call that starts there next, and never open a call for it, which would reach the model without a name. In live runs with gemini-3-flash-preview, neither Vertex AI's OpenAI-compatible endpoint nor LiteLLM's vertex_ai provider sent this shape; both put the signature on the delta that names the call. Follow-up to #7441, for the case raised on #7438. --- src/google/adk/models/lite_llm.py | 39 +++++++++--- tests/unittests/models/test_litellm.py | 83 ++++++++++++++++++++++++++ 2 files changed, 115 insertions(+), 7 deletions(-) diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index 5c79acd3fa..0a0f1d11cb 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -2845,9 +2845,14 @@ def _has_meaningful_signal(message: Message | Delta | None) -> bool: func_args = function_obj.get("arguments") func_index = tool_call.get("index", idx) tool_call_id = tool_call.get("id") + thought_signature = _extract_thought_signature_from_tool_call( + tool_call + ) - # Ignore empty chunks that don't carry any information. - if not func_name and not func_args: + # Ignore empty chunks that don't carry any information. A chunk + # with only a signature still counts: the signature belongs to the + # call another chunk names. + if not func_name and not func_args and not thought_signature: continue yield FunctionChunk( @@ -2855,9 +2860,7 @@ def _has_meaningful_signal(message: Message | Delta | None) -> bool: name=func_name, args=func_args, index=func_index, - thought_signature=_extract_thought_signature_from_tool_call( - tool_call - ), + thought_signature=thought_signature, ), finish_reason if finish_reason and not (message_content or tool_calls or reasoning_parts): @@ -3941,6 +3944,8 @@ async def generate_content_async( function_calls: dict[int, dict[str, Any]] = ( {} ) # index -> {name, args_parts, id, thought_signature} + # Signatures that arrived before any chunk naming their call. + pending_signatures: dict[int, bytes] = {} tool_call_trackers: Dict[int, _BraceDepthTracker] = {} completion_args["stream"] = True completion_args["stream_options"] = {"include_usage": True} @@ -4078,6 +4083,7 @@ def _reset_stream_buffers() -> None: text_parts.clear() reasoning_parts = [] function_calls.clear() + pending_signatures.clear() tool_call_trackers.clear() # The reason belongs to the segment just finalized; carrying it into # the next one would stamp the wrong reason on the next response. @@ -4104,14 +4110,33 @@ def _reset_stream_buffers() -> None: for chunk, finish_reason in _model_response_to_chunk(part): if finish_reason: last_finish_reason = finish_reason - if isinstance(chunk, FunctionChunk): + if ( + isinstance(chunk, FunctionChunk) + and not chunk.name + and not chunk.args + and chunk.thought_signature + ): + # Only a signature. Opening a call for it would send the model a + # call with no name, so attach it to the call at its own index + # (the fallback index may already point past that call), or hold + # it for the call that starts there next. + signed_call = function_calls.get( + chunk.index if chunk.index is not None else fallback_index + ) + if signed_call is not None: + signed_call["thought_signature"] = chunk.thought_signature + else: + pending_signatures[chunk.index or fallback_index] = ( + chunk.thought_signature + ) + elif isinstance(chunk, FunctionChunk): index = chunk.index or fallback_index if index not in function_calls: function_calls[index] = { "name": "", "args_parts": [], "id": None, - "thought_signature": None, + "thought_signature": pending_signatures.pop(index, None), } if chunk.name: diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index e6616afdbe..24d6a241eb 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -4345,6 +4345,89 @@ async def test_streaming_parallel_tool_calls_keep_signature_per_call( assert parts[1].thought_signature is None +def test_model_response_to_chunk_keeps_a_signature_only_delta(): + """A delta with only a signature is kept so it can reach its call.""" + chunks = list( + _model_response_to_chunk( + _streamed_tool_call_chunk(index=0, signature=b"late_sig") + ) + ) + + function_chunk = chunks[0][0] + assert isinstance(function_chunk, FunctionChunk) + assert not function_chunk.name and not function_chunk.args + assert function_chunk.thought_signature == b"late_sig" + + +@pytest.mark.asyncio +async def test_streaming_signature_after_its_call_signs_that_call( + mock_completion, lite_llm_instance +): + """A signature-only delta that trails a later call still signs its own.""" + mock_completion.return_value = iter([ + _streamed_tool_call_chunk( + index=0, + call_id="call_1", + name="get_weather", + arguments='{"city": "Oslo"}', + ), + _streamed_tool_call_chunk( + index=1, + call_id="call_2", + name="get_weather", + arguments='{"city": "Bergen"}', + ), + _streamed_tool_call_chunk(index=0, signature=b"late_sig"), + _streamed_finish_chunk(), + ]) + + responses = [ + response + async for response in lite_llm_instance.generate_content_async( + LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True + ) + ] + + parts = responses[-1].content.parts + assert [p.function_call.id for p in parts] == ["call_1", "call_2"] + assert [p.thought_signature for p in parts] == [b"late_sig", None] + + +@pytest.mark.asyncio +async def test_streaming_signature_before_its_call_signs_that_call( + mock_completion, lite_llm_instance +): + """A signature-only delta ahead of its call opens no call of its own.""" + mock_completion.return_value = iter([ + _streamed_tool_call_chunk(index=0, signature=b"early_sig"), + _streamed_tool_call_chunk( + index=0, + call_id="call_1", + name="get_weather", + arguments='{"city": "Oslo"}', + ), + _streamed_finish_chunk(), + ]) + + responses = [ + response + async for response in lite_llm_instance.generate_content_async( + LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True + ) + ] + + streamed_calls = [ + part.function_call + for response in responses + for part in (response.content.parts if response.content else []) + if part.function_call + ] + assert all(call.name == "get_weather" for call in streamed_calls) + parts = responses[-1].content.parts + assert [p.function_call.id for p in parts] == ["call_1"] + assert parts[0].thought_signature == b"early_sig" + + def test_message_to_generate_content_response_no_thought_signature(): """Parts without thought_signature have thought_signature=None.""" message = ChatCompletionAssistantMessage(