Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 28 additions & 1 deletion src/google/adk/flows/llm_flows/base_llm_flow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand All @@ -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]
Comment on lines +124 to +126
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This loop can be made more concise and Pythonic by using zip. It simplifies iterating over function_calls and prior_fc_ids together and automatically stops when the shorter list is exhausted, matching the current logic.

Suggested change
for idx, fc in enumerate(function_calls):
if idx < len(prior_fc_ids) and prior_fc_ids[idx]:
fc.id = prior_fc_ids[idx]
for fc, prior_id in zip(function_calls, prior_fc_ids):
if prior_id:
fc.id = prior_id


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
Expand Down
8 changes: 3 additions & 5 deletions src/google/adk/tools/transfer_to_agent_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
"""
Expand Down
75 changes: 75 additions & 0 deletions tests/unittests/flows/llm_flows/test_base_llm_flow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'