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
34 changes: 30 additions & 4 deletions src/google/adk/flows/llm_flows/contents.py
Original file line number Diff line number Diff line change
Expand Up @@ -148,7 +148,7 @@ async def run_async(
def _rearrange_events_for_async_function_responses_in_history(
events: list[Event],
) -> list[Event]:
"""Rearrange the async function_response events in the history."""
"""Rearrange async function responses and their model replies in history."""
# 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
Expand All @@ -164,6 +164,7 @@ def _rearrange_events_for_async_function_responses_in_history(
call_event_indices_by_id.setdefault(function_call.id, []).append(i)

response_event_index_by_call: dict[tuple[str | None, int], int] = {}
response_event_indices_by_call: dict[tuple[str | None, int], list[int]] = {}
history_has_function_responses = False
for i, event in enumerate(events):
for function_response in event.get_function_responses():
Expand All @@ -176,15 +177,40 @@ def _rearrange_events_for_async_function_responses_in_history(
# 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
call_key = (function_response.id, owning_call_event_index)
response_event_index_by_call[call_key] = i
response_event_indices_by_call.setdefault(call_key, []).append(i)

if not history_has_function_responses:
return events

# A text-only model reply immediately after an old tool update was generated
# from the result that is about to be discarded. Remove that reply as well,
# while stopping at any user input, function call, or other event boundary.
latest_response_event_indices = set(response_event_index_by_call.values())
superseded_response_event_indices = {
response_event_index
for response_event_indices in response_event_indices_by_call.values()
for response_event_index in response_event_indices
if response_event_index not in latest_response_event_indices
}
stale_model_event_indices: set[int] = set()
for response_event_index in superseded_response_event_indices:
for i in range(response_event_index + 1, len(events)):
event = events[i]
if (
not event.content
or event.content.role != 'model'
or event.get_function_calls()
or not any(part.text for part in event.content.parts or [])
):
break
stale_model_event_indices.add(i)

result_events: list[Event] = []
for i, event in enumerate(events):
if i in stale_model_event_indices:
continue
if event.get_function_responses():
# function_response should be handled together with function_call below.
continue
Expand Down
33 changes: 33 additions & 0 deletions tests/unittests/flows/llm_flows/test_contents.py
Original file line number Diff line number Diff line change
Expand Up @@ -1971,6 +1971,39 @@ def _function_response_event(call_id: str, name: str, result: str) -> Event:
)


def _model_text_event(text: str) -> Event:
return Event(
invocation_id="inv1",
author="test_agent",
content=types.ModelContent(text),
)


def test_rearrange_async_function_responses_drops_stale_model_reply():
"""A model reply to a superseded tool update must not remain in history."""
events = [
_function_call_event("call_1", "watch"),
_function_response_event("call_1", "watch", "progress"),
_model_text_event("Still working."),
_function_response_event("call_1", "watch", "done"),
_model_text_event("Finished."),
Event(
invocation_id="inv2",
author="user",
content=types.UserContent("What happened?"),
),
]

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[2].content.parts[0].text == "Finished."
assert result[3].content.parts[0].text == "What happened?"


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 = [
Expand Down