diff --git a/src/google/adk/flows/llm_flows/contents.py b/src/google/adk/flows/llm_flows/contents.py index e9233f2b..d3e9e481 100644 --- a/src/google/adk/flows/llm_flows/contents.py +++ b/src/google/adk/flows/llm_flows/contents.py @@ -14,6 +14,7 @@ from __future__ import annotations +from bisect import bisect_left import copy import logging from typing import AsyncGenerator @@ -147,19 +148,42 @@ def _rearrange_events_for_async_function_responses_in_history( events: list[Event], ) -> list[Event]: """Rearrange the async function_response events in the history.""" - function_call_id_to_response_events_index: dict[str | None, int] = {} + # A model may hand out the same function call id more than once in a session, + # so an id on its own does not identify a single call. Each response is + # attributed to the newest call that precedes it and carries the same id, and + # a call then takes the last response attributed to it. Taking the last one + # keeps the closing update of a long-running tool, which reports progress + # several times under one id, while attributing first stops a reused id from + # handing a call the response that belongs to a different call. + call_event_indices_by_id: dict[str | None, list[int]] = {} for i, event in enumerate(events): - function_responses = event.get_function_responses() - if function_responses: - for function_response in function_responses: - function_call_id = function_response.id - function_call_id_to_response_events_index[function_call_id] = i + if event.get_function_responses(): + continue + for function_call in event.get_function_calls(): + call_event_indices_by_id.setdefault(function_call.id, []).append(i) - if not function_call_id_to_response_events_index: + response_event_index_by_call: dict[tuple[str | None, int], int] = {} + history_has_function_responses = False + for i, event in enumerate(events): + for function_response in event.get_function_responses(): + history_has_function_responses = True + call_event_indices = call_event_indices_by_id.get(function_response.id) + if not call_event_indices: + continue + # Indices are collected in ascending order, so the call that owns this + # response is the one just before it. A response preceding every call + # that carries its id keeps the first, as it did before ids could repeat. + preceding_calls = bisect_left(call_event_indices, i) + owning_call_event_index = call_event_indices[max(preceding_calls - 1, 0)] + response_event_index_by_call[ + (function_response.id, owning_call_event_index) + ] = i + + if not history_has_function_responses: return events result_events: list[Event] = [] - for event in events: + for i, event in enumerate(events): if event.get_function_responses(): # function_response should be handled together with function_call below. continue @@ -167,11 +191,11 @@ def _rearrange_events_for_async_function_responses_in_history( function_response_events_indices = set() for function_call in event.get_function_calls(): - function_call_id = function_call.id - if function_call_id in function_call_id_to_response_events_index: - function_response_events_indices.add( - function_call_id_to_response_events_index[function_call_id] - ) + response_event_index = response_event_index_by_call.get( + (function_call.id, i) + ) + if response_event_index is not None: + function_response_events_indices.add(response_event_index) result_events.append(event) if not function_response_events_indices: continue diff --git a/tests/unittests/flows/llm_flows/test_contents.py b/tests/unittests/flows/llm_flows/test_contents.py index 445cc162..dcef2a98 100644 --- a/tests/unittests/flows/llm_flows/test_contents.py +++ b/tests/unittests/flows/llm_flows/test_contents.py @@ -1964,6 +1964,99 @@ def test_rearrange_async_function_responses_early_returns_when_no_responses(): assert result is events +def _function_call_event(call_id: str, name: str) -> Event: + return Event( + invocation_id="inv1", + author="test_agent", + content=types.Content( + role="model", + parts=[ + types.Part( + function_call=types.FunctionCall( + id=call_id, name=name, args={} + ) + ) + ], + ), + ) + + +def _function_response_event(call_id: str, name: str, result: str) -> Event: + return Event( + invocation_id="inv1", + author="user", + content=types.Content( + role="user", + parts=[ + types.Part( + function_response=types.FunctionResponse( + id=call_id, name=name, response={"result": result} + ) + ) + ], + ), + ) + + +def test_rearrange_async_function_responses_reused_id_across_tools(): + """A reused call id must not pair a call with a different tool's response.""" + events = [ + _function_call_event("call_807", "site_posture"), + _function_response_event("call_807", "site_posture", "site"), + _function_call_event("call_807", "fleet_summary"), + _function_response_event("call_807", "fleet_summary", "fleet"), + ] + + result = contents._rearrange_events_for_async_function_responses_in_history( # pylint: disable=protected-access + events + ) + + assert len(result) == 4 + assert result[0].get_function_calls()[0].name == "site_posture" + assert result[1].get_function_responses()[0].name == "site_posture" + assert result[2].get_function_calls()[0].name == "fleet_summary" + assert result[3].get_function_responses()[0].name == "fleet_summary" + + +def test_rearrange_async_function_responses_reused_id_same_tool(): + """A call id reused by one tool must keep each call's own response.""" + events = [ + _function_call_event("call_42", "lookup"), + _function_response_event("call_42", "lookup", "first"), + _function_call_event("call_42", "lookup"), + _function_response_event("call_42", "lookup", "second"), + ] + + result = contents._rearrange_events_for_async_function_responses_in_history( # pylint: disable=protected-access + events + ) + + assert len(result) == 4 + assert result[1].get_function_responses()[0].response == {"result": "first"} + assert result[3].get_function_responses()[0].response == {"result": "second"} + + +def test_rearrange_async_function_responses_reused_id_keeps_last_update(): + """A call reporting progress twice keeps its last update, not a later call's.""" + events = [ + _function_call_event("call_7", "watch"), + _function_response_event("call_7", "watch", "progress"), + _function_response_event("call_7", "watch", "done"), + _function_call_event("call_7", "watch"), + _function_response_event("call_7", "watch", "second_call"), + ] + + result = contents._rearrange_events_for_async_function_responses_in_history( # pylint: disable=protected-access + events + ) + + assert len(result) == 4 + assert result[1].get_function_responses()[0].response == {"result": "done"} + assert result[3].get_function_responses()[0].response == { + "result": "second_call" + } + + def _long_running_call_event() -> Event: return Event( invocation_id="inv2",