Skip to content

Commit 5e1d3eb

Browse files
committed
fix(models): route interleaved interactions function-call deltas by step index
The Interactions streaming converter attached arguments_delta and step.stop events to state.parts[-1] instead of the function-call step identified by the event's index. When multiple function-call steps are interleaved, deltas for an earlier step were appended to the most recently started call, concatenating their raw JSON and producing invalid arguments (or swapped/merged names) in the final response. Track function-call parts by step index in _StreamState and route arguments_delta and step.stop to the matching step, retaining the existing last-part fallback for events without an index. Fixes #6832
1 parent 4599a52 commit 5e1d3eb

2 files changed

Lines changed: 99 additions & 11 deletions

File tree

src/google/adk/models/interactions_utils.py

Lines changed: 40 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -772,6 +772,12 @@ class _StreamState:
772772
"""
773773

774774
parts: list[types.Part] = dataclasses.field(default_factory=list)
775+
# Maps a function-call step's ``index`` to the part started at that step, so
776+
# interleaved calls route their argument deltas and stops to the matching
777+
# step instead of always landing on the most recently started call.
778+
fc_parts_by_index: dict[int, types.Part] = dataclasses.field(
779+
default_factory=dict
780+
)
775781
web_search_queries: list[str] = dataclasses.field(default_factory=list)
776782
grounding_chunks: list[types.GroundingChunk] = dataclasses.field(
777783
default_factory=list
@@ -836,23 +842,43 @@ def _handle_media(
836842
return _partial_part_response(part, interaction_id)
837843

838844

845+
def _resolve_streaming_function_call_part(
846+
index: int | None, state: _StreamState
847+
) -> types.Part | None:
848+
"""Resolve the function-call part a streaming event applies to.
849+
850+
Streaming events carry the ``index`` of the step they belong to. When that
851+
index maps to a known function-call part we use it directly, so argument
852+
deltas and step stops for interleaved calls route to the correct step instead
853+
of always landing on the most recently started call. Events without an index
854+
(or from builds that don't track one) fall back to the last started function
855+
call to preserve the previous behavior.
856+
"""
857+
if index is not None and index in state.fc_parts_by_index:
858+
return state.fc_parts_by_index[index]
859+
if state.parts and state.parts[-1].function_call:
860+
return state.parts[-1]
861+
return None
862+
863+
839864
def _handle_arguments_delta(
840-
delta: StepDeltaData, state: _StreamState, interaction_id: str | None
865+
delta: StepDeltaData,
866+
state: _StreamState,
867+
interaction_id: str | None,
868+
index: int | None = None,
841869
) -> LlmResponse | None:
842-
if not state.parts:
843-
return None
844-
last_part = state.parts[-1]
845-
if not last_part.function_call:
870+
target_part = _resolve_streaming_function_call_part(index, state)
871+
if target_part is None or not target_part.function_call:
846872
return None
847873
delta_args = delta.arguments
848-
if delta_args is None or last_part.function_call.partial_args is None:
874+
if delta_args is None or target_part.function_call.partial_args is None:
849875
return None
850-
last_part.function_call.partial_args.append(
876+
target_part.function_call.partial_args.append(
851877
types.PartialArg(string_value=delta_args)
852878
)
853879
chunk_part = types.Part(
854880
function_call=types.FunctionCall(
855-
name=last_part.function_call.name,
881+
name=target_part.function_call.name,
856882
partial_args=[types.PartialArg(string_value=delta_args)],
857883
)
858884
)
@@ -1066,6 +1092,8 @@ def convert_interaction_event_to_llm_response(
10661092
)
10671093
part = types.Part(function_call=fc)
10681094
state.parts.append(part)
1095+
if event.index is not None:
1096+
state.fc_parts_by_index[event.index] = part
10691097

10701098
return LlmResponse(
10711099
content=types.Content(role='model', parts=[part]),
@@ -1087,7 +1115,7 @@ def convert_interaction_event_to_llm_response(
10871115
elif delta_type in ('image', 'audio', 'video', 'document'):
10881116
return _handle_media(delta, state, interaction_id)
10891117
elif delta_type == 'arguments_delta':
1090-
return _handle_arguments_delta(delta, state, interaction_id)
1118+
return _handle_arguments_delta(delta, state, interaction_id, event.index)
10911119
elif delta_type == 'code_execution_call':
10921120
return _handle_code_execution_call(delta, state, interaction_id)
10931121
elif delta_type == 'code_execution_result':
@@ -1104,8 +1132,9 @@ def convert_interaction_event_to_llm_response(
11041132
return _handle_unknown_delta(delta, state, interaction_id)
11051133

11061134
elif isinstance(event, StepStop):
1107-
if state.parts and state.parts[-1].function_call:
1108-
fc = state.parts[-1].function_call
1135+
target_part = _resolve_streaming_function_call_part(event.index, state)
1136+
if target_part is not None and target_part.function_call:
1137+
fc = target_part.function_call
11091138
if fc.partial_args is not None:
11101139
arg_str = ''.join(pa.string_value or '' for pa in fc.partial_args)
11111140

tests/unittests/models/test_interactions_utils.py

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2074,6 +2074,65 @@ def test_function_call_streaming_json_parse_error(self, caplog):
20742074
# The logging check can remain to ensure the raw exception is still logged.
20752075
assert 'Failed to parse function call args' in caplog.text
20762076

2077+
def test_interleaved_function_call_streaming_routes_by_index(self):
2078+
"""Interleaved function-call steps route deltas/stops by their index.
2079+
2080+
Two calls start at indexes 0 and 1, then arguments for both arrive before
2081+
either stops. Without index-based routing, both deltas would be appended to
2082+
the most recently started call (index 1), so the first call ends up with no
2083+
arguments while the second receives two concatenated JSON objects.
2084+
"""
2085+
state = interactions_utils._StreamState()
2086+
2087+
# Start two function calls at different indexes.
2088+
for idx, (call_id, name) in enumerate(
2089+
[('call_0', 'get_weather'), ('call_1', 'get_time')]
2090+
):
2091+
interactions_utils.convert_interaction_event_to_llm_response(
2092+
StepStart(
2093+
event_type='step.start',
2094+
index=idx,
2095+
step=FunctionCallStep(
2096+
type='function_call', id=call_id, name=name, arguments={}
2097+
),
2098+
),
2099+
state,
2100+
interaction_id='int_multi',
2101+
)
2102+
2103+
# Interleave argument deltas: index 0 first, then index 1.
2104+
interactions_utils.convert_interaction_event_to_llm_response(
2105+
StepDelta(
2106+
event_type='step.delta',
2107+
index=0,
2108+
delta={'type': 'arguments_delta', 'arguments': '{"city": "Paris"}'},
2109+
),
2110+
state,
2111+
interaction_id='int_multi',
2112+
)
2113+
interactions_utils.convert_interaction_event_to_llm_response(
2114+
StepDelta(
2115+
event_type='step.delta',
2116+
index=1,
2117+
delta={'type': 'arguments_delta', 'arguments': '{"zone": "UTC"}'},
2118+
),
2119+
state,
2120+
interaction_id='int_multi',
2121+
)
2122+
2123+
# Stop both steps.
2124+
for idx in (0, 1):
2125+
interactions_utils.convert_interaction_event_to_llm_response(
2126+
StepStop(event_type='step.stop', index=idx),
2127+
state,
2128+
interaction_id='int_multi',
2129+
)
2130+
2131+
assert state.parts[0].function_call.name == 'get_weather'
2132+
assert state.parts[0].function_call.args == {'city': 'Paris'}
2133+
assert state.parts[1].function_call.name == 'get_time'
2134+
assert state.parts[1].function_call.args == {'zone': 'UTC'}
2135+
20772136

20782137
@pytest.mark.parametrize(
20792138
('streamed_events_factory', 'expected_ids'),

0 commit comments

Comments
 (0)