diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index a854c0104b..f9f7b24d68 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -2395,7 +2395,20 @@ def _message_to_generate_content_response( raise ValueError( "Function-call part factory returned no function call" ) - function_call.id = tool_call.id + # The signature is carried on the part, so strip it back out of the id. + # Leaving it there appends a few hundred characters of base64 to every + # consumer of `function_call_id`, and `_content_to_message_param` + # re-attaches it to the outgoing tool call from `thought_signature` + # anyway. Only an id whose suffix decodes as a signature is LiteLLM's + # own; every other id is the provider's and stays opaque. + call_id, separator, suffix = tool_call.id.partition( + _THOUGHT_SIGNATURE_SEPARATOR + ) + function_call.id = ( + call_id + if separator and _decode_thought_signature(suffix) + else tool_call.id + ) if thought_signature: part.thought_signature = thought_signature parts.append(part) diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index 68b0e74c17..3c011988f2 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -2977,6 +2977,146 @@ def test_message_to_generate_content_response_preserves_thought_signature(): assert fc_part.thought_signature == b"round_trip_sig" +def test_message_to_generate_content_response_strips_signature_from_id(): + """An id carrying an embedded signature is split before it reaches the part.""" + sig_b64 = base64.b64encode(b"embedded_sig").decode("utf-8") + message = ChatCompletionAssistantMessage( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + type="function", + id=f"call_ts_2{_THOUGHT_SIGNATURE_SEPARATOR}{sig_b64}", + function=Function( + name="load_skill", + arguments='{"skill": "my_skill"}', + ), + ) + ], + ) + + response = _message_to_generate_content_response(message) + fc_part = response.content.parts[0] + assert fc_part.function_call.id == "call_ts_2" + assert fc_part.thought_signature == b"embedded_sig" + + +def test_message_to_generate_content_response_keeps_a_non_signature_id_whole(): + """A provider id is opaque, so a separator whose suffix is not a signature stays.""" + message = ChatCompletionAssistantMessage( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + type="function", + id=f"job{_THOUGHT_SIGNATURE_SEPARATOR}not-base64", + function=Function(name="load_skill", arguments="{}"), + ) + ], + ) + + response = _message_to_generate_content_response(message) + fc_part = response.content.parts[0] + assert ( + fc_part.function_call.id == f"job{_THOUGHT_SIGNATURE_SEPARATOR}not-base64" + ) + assert fc_part.thought_signature is None + + +def test_message_to_generate_content_response_keeps_the_id_when_the_signature_came_from_extra_content(): + """The id is only LiteLLM's to split when its own suffix decodes as the signature.""" + sig_b64 = base64.b64encode(b"channel_sig").decode("utf-8") + message = ChatCompletionAssistantMessage( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + type="function", + id=f"job{_THOUGHT_SIGNATURE_SEPARATOR}not-base64", + function=Function(name="load_skill", arguments="{}"), + extra_content={"google": {"thought_signature": sig_b64}}, + ) + ], + ) + + response = _message_to_generate_content_response(message) + fc_part = response.content.parts[0] + assert ( + fc_part.function_call.id == f"job{_THOUGHT_SIGNATURE_SEPARATOR}not-base64" + ) + assert fc_part.thought_signature == b"channel_sig" + + +@pytest.mark.asyncio +async def test_embedded_signature_round_trips_from_a_clean_id(): + """The repro end to end: the id is cleaned inbound and the signature still goes back.""" + sig_b64 = base64.b64encode(b"embedded_sig").decode("utf-8") + message = ChatCompletionAssistantMessage( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + type="function", + id=f"call_abc{_THOUGHT_SIGNATURE_SEPARATOR}{sig_b64}", + function=Function(name="load_skill", arguments='{"skill": "s"}'), + ) + ], + ) + + response = _message_to_generate_content_response(message) + outbound = await _content_to_message_param(response.content) + + tool_call = outbound["tool_calls"][0] + assert tool_call["id"] == "call_abc" + assert tool_call["provider_specific_fields"]["thought_signature"] == sig_b64 + assert tool_call["extra_content"]["google"]["thought_signature"] == sig_b64 + + +@pytest.mark.asyncio +async def test_generate_content_async_cleans_an_embedded_signature_id(): + """The reported path end to end: a Gemini tool call arrives with a clean id.""" + sig_b64 = base64.b64encode(b"embedded_sig").decode("utf-8") + model_response = ModelResponse( + model="test_model", + choices=[ + Choices( + message=ChatCompletionAssistantMessage( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + type="function", + id=f"call_abc{_THOUGHT_SIGNATURE_SEPARATOR}{sig_b64}", + function=Function( + name="test_function", + arguments='{"test_arg": "test_value"}', + ), + ) + ], + ) + ) + ], + ) + llm = LiteLlm( + model="test_model", + llm_client=MockLLMClient( + AsyncMock(return_value=model_response), + Mock(return_value=model_response), + ), + ) + + responses = [ + response + async for response in llm.generate_content_async( + LLM_REQUEST_WITH_FUNCTION_DECLARATION + ) + ] + + part = responses[0].content.parts[0] + assert part.function_call.id == "call_abc" + assert part.thought_signature == b"embedded_sig" + + def test_message_to_generate_content_response_no_thought_signature(): """Parts without thought_signature have thought_signature=None.""" message = ChatCompletionAssistantMessage(