diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index 424bb580e1..d00980d084 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -90,11 +90,26 @@ def _finalize_model_response_event( Args: llm_request: The original LLM request. llm_response: The LLM response from the model. - model_response_event: The base event to populate. + model_response_event: The base event to populate. When SSE streaming is + active, this same object is passed for every partial *and* the final + response. After the first partial event is finalized, the assigned + ``adk-*`` function-call IDs are written back into + ``model_response_event.content`` so that subsequent calls (including the + final non-partial event) can reuse the same IDs instead of generating + new ones. Returns: The finalized Event with LLM response data merged in. """ + # Collect any function-call IDs that were already assigned during a previous + # partial-event finalization for this same logical LLM turn. We extract them + # *before* the merge so they are not lost when llm_response overwrites content. + prior_fc_ids: list[str | None] = [] + if model_response_event.content: + prior_fc_ids = [ + fc.id for fc in (model_response_event.get_function_calls() or []) + ] + finalized_event = Event.model_validate({ **model_response_event.model_dump(exclude_none=True), **llm_response.model_dump(exclude_none=True), @@ -103,7 +118,19 @@ def _finalize_model_response_event( if finalized_event.content: function_calls = finalized_event.get_function_calls() if function_calls: + # Restore previously-assigned IDs (by position) so that partial and + # final SSE events for the same function call share the same ID. + if prior_fc_ids: + for idx, fc in enumerate(function_calls): + if idx < len(prior_fc_ids) and prior_fc_ids[idx]: + fc.id = prior_fc_ids[idx] + functions.populate_client_function_call_id(finalized_event) + + # Persist the now-assigned IDs back into model_response_event so that + # the next call (e.g. the final non-partial event) can reuse them. + model_response_event.content = finalized_event.content + finalized_event.long_running_tool_ids = ( functions.get_long_running_function_calls( function_calls, llm_request.tools_dict diff --git a/src/google/adk/tools/transfer_to_agent_tool.py b/src/google/adk/tools/transfer_to_agent_tool.py index 4db933999a..9f51e7fe36 100644 --- a/src/google/adk/tools/transfer_to_agent_tool.py +++ b/src/google/adk/tools/transfer_to_agent_tool.py @@ -23,17 +23,15 @@ from .tool_context import ToolContext +# Note: For most use cases, you should use TransferToAgentTool instead of this +# function directly. TransferToAgentTool provides additional enum constraints +# that prevent LLMs from hallucinating invalid agent names. def transfer_to_agent(agent_name: str, tool_context: ToolContext) -> None: """Transfer the question to another agent. This tool hands off control to another agent when it's more suitable to answer the user's question according to the agent's description. - Note: - For most use cases, you should use TransferToAgentTool instead of this - function directly. TransferToAgentTool provides additional enum constraints - that prevent LLMs from hallucinating invalid agent names. - Args: agent_name: the agent name to transfer to. """ diff --git a/tests/unittests/flows/llm_flows/test_base_llm_flow.py b/tests/unittests/flows/llm_flows/test_base_llm_flow.py index 3dfadbcabf..c9a7b384a0 100644 --- a/tests/unittests/flows/llm_flows/test_base_llm_flow.py +++ b/tests/unittests/flows/llm_flows/test_base_llm_flow.py @@ -487,3 +487,78 @@ async def call(self, **kwargs): assert result1.grounding_metadata == {'foo': 'bar'} assert result2.grounding_metadata == {'foo': 'bar'} assert result3.grounding_metadata == {'foo': 'bar'} + + +# --------------------------------------------------------------------------- +# Tests for _finalize_model_response_event function-call ID consistency +# --------------------------------------------------------------------------- + +from google.adk.flows.llm_flows.base_llm_flow import _finalize_model_response_event + + +def _make_fc_response(fc_name: str, fc_id: str | None = None, partial: bool = False) -> LlmResponse: + """Helper: build an LlmResponse with a single function call.""" + return LlmResponse( + content=types.Content( + role='model', + parts=[ + types.Part( + function_call=types.FunctionCall( + name=fc_name, + args={'x': 1}, + id=fc_id, + ) + ) + ], + ), + partial=partial, + ) + + +def test_finalize_model_response_event_consistent_fc_id_across_partial_and_final(): + """Function call IDs must be identical in partial and final SSE events. + + Regression test for https://github.com/google/adk-python/issues/4609. + When SSE streaming is active, _finalize_model_response_event is called + once for the partial event and once for the final event, both sharing the + same model_response_event object. The assigned adk-* ID must be the same + in both calls. + """ + llm_request = LlmRequest() + llm_request.tools_dict = {} + base_event = Event( + invocation_id='inv1', + author='agent', + ) + + # First call: partial streaming event (function call has no ID from LLM) + partial_response = _make_fc_response('my_tool', fc_id=None, partial=True) + partial_finalized = _finalize_model_response_event( + llm_request, partial_response, base_event + ) + partial_fc_id = partial_finalized.get_function_calls()[0].id + assert partial_fc_id is not None + assert partial_fc_id.startswith('adk-') + + # Second call: final (non-partial) event for the same function call + final_response = _make_fc_response('my_tool', fc_id=None, partial=False) + final_finalized = _finalize_model_response_event( + llm_request, final_response, base_event + ) + final_fc_id = final_finalized.get_function_calls()[0].id + + assert final_fc_id == partial_fc_id, ( + f'Function call ID changed between partial ({partial_fc_id!r}) and ' + f'final ({final_fc_id!r}) SSE events — HITL workflows will break.' + ) + + +def test_finalize_model_response_event_preserves_llm_assigned_id(): + """If the LLM already assigned an ID, it must be preserved as-is.""" + llm_request = LlmRequest() + llm_request.tools_dict = {} + base_event = Event(invocation_id='inv1', author='agent') + + response = _make_fc_response('my_tool', fc_id='llm-assigned-id') + finalized = _finalize_model_response_event(llm_request, response, base_event) + assert finalized.get_function_calls()[0].id == 'llm-assigned-id'