diff --git a/src/google/adk/models/interactions_utils.py b/src/google/adk/models/interactions_utils.py index 43687f09ce..25f3390d52 100644 --- a/src/google/adk/models/interactions_utils.py +++ b/src/google/adk/models/interactions_utils.py @@ -772,6 +772,12 @@ class _StreamState: """ parts: list[types.Part] = dataclasses.field(default_factory=list) + # Maps a function-call step's ``index`` to the part started at that step, so + # interleaved calls route their argument deltas and stops to the matching + # step instead of always landing on the most recently started call. + fc_parts_by_index: dict[int, types.Part] = dataclasses.field( + default_factory=dict + ) web_search_queries: list[str] = dataclasses.field(default_factory=list) grounding_chunks: list[types.GroundingChunk] = dataclasses.field( default_factory=list @@ -836,23 +842,43 @@ def _handle_media( return _partial_part_response(part, interaction_id) +def _resolve_streaming_function_call_part( + index: int | None, state: _StreamState +) -> types.Part | None: + """Resolve the function-call part a streaming event applies to. + + Streaming events carry the ``index`` of the step they belong to. When that + index maps to a known function-call part we use it directly, so argument + deltas and step stops for interleaved calls route to the correct step instead + of always landing on the most recently started call. Events without an index + (or from builds that don't track one) fall back to the last started function + call to preserve the previous behavior. + """ + if index is not None and index in state.fc_parts_by_index: + return state.fc_parts_by_index[index] + if state.parts and state.parts[-1].function_call: + return state.parts[-1] + return None + + def _handle_arguments_delta( - delta: StepDeltaData, state: _StreamState, interaction_id: str | None + delta: StepDeltaData, + state: _StreamState, + interaction_id: str | None, + index: int | None = None, ) -> LlmResponse | None: - if not state.parts: - return None - last_part = state.parts[-1] - if not last_part.function_call: + target_part = _resolve_streaming_function_call_part(index, state) + if target_part is None or not target_part.function_call: return None delta_args = delta.arguments - if delta_args is None or last_part.function_call.partial_args is None: + if delta_args is None or target_part.function_call.partial_args is None: return None - last_part.function_call.partial_args.append( + target_part.function_call.partial_args.append( types.PartialArg(string_value=delta_args) ) chunk_part = types.Part( function_call=types.FunctionCall( - name=last_part.function_call.name, + name=target_part.function_call.name, partial_args=[types.PartialArg(string_value=delta_args)], ) ) @@ -1066,6 +1092,8 @@ def convert_interaction_event_to_llm_response( ) part = types.Part(function_call=fc) state.parts.append(part) + if event.index is not None: + state.fc_parts_by_index[event.index] = part return LlmResponse( content=types.Content(role='model', parts=[part]), @@ -1087,7 +1115,7 @@ def convert_interaction_event_to_llm_response( elif delta_type in ('image', 'audio', 'video', 'document'): return _handle_media(delta, state, interaction_id) elif delta_type == 'arguments_delta': - return _handle_arguments_delta(delta, state, interaction_id) + return _handle_arguments_delta(delta, state, interaction_id, event.index) elif delta_type == 'code_execution_call': return _handle_code_execution_call(delta, state, interaction_id) elif delta_type == 'code_execution_result': @@ -1104,8 +1132,9 @@ def convert_interaction_event_to_llm_response( return _handle_unknown_delta(delta, state, interaction_id) elif isinstance(event, StepStop): - if state.parts and state.parts[-1].function_call: - fc = state.parts[-1].function_call + target_part = _resolve_streaming_function_call_part(event.index, state) + if target_part is not None and target_part.function_call: + fc = target_part.function_call if fc.partial_args is not None: arg_str = ''.join(pa.string_value or '' for pa in fc.partial_args) diff --git a/tests/unittests/models/test_interactions_utils.py b/tests/unittests/models/test_interactions_utils.py index 109c3be544..04ef666833 100644 --- a/tests/unittests/models/test_interactions_utils.py +++ b/tests/unittests/models/test_interactions_utils.py @@ -2074,6 +2074,65 @@ def test_function_call_streaming_json_parse_error(self, caplog): # The logging check can remain to ensure the raw exception is still logged. assert 'Failed to parse function call args' in caplog.text + def test_interleaved_function_call_streaming_routes_by_index(self): + """Interleaved function-call steps route deltas/stops by their index. + + Two calls start at indexes 0 and 1, then arguments for both arrive before + either stops. Without index-based routing, both deltas would be appended to + the most recently started call (index 1), so the first call ends up with no + arguments while the second receives two concatenated JSON objects. + """ + state = interactions_utils._StreamState() + + # Start two function calls at different indexes. + for idx, (call_id, name) in enumerate( + [('call_0', 'get_weather'), ('call_1', 'get_time')] + ): + interactions_utils.convert_interaction_event_to_llm_response( + StepStart( + event_type='step.start', + index=idx, + step=FunctionCallStep( + type='function_call', id=call_id, name=name, arguments={} + ), + ), + state, + interaction_id='int_multi', + ) + + # Interleave argument deltas: index 0 first, then index 1. + interactions_utils.convert_interaction_event_to_llm_response( + StepDelta( + event_type='step.delta', + index=0, + delta={'type': 'arguments_delta', 'arguments': '{"city": "Paris"}'}, + ), + state, + interaction_id='int_multi', + ) + interactions_utils.convert_interaction_event_to_llm_response( + StepDelta( + event_type='step.delta', + index=1, + delta={'type': 'arguments_delta', 'arguments': '{"zone": "UTC"}'}, + ), + state, + interaction_id='int_multi', + ) + + # Stop both steps. + for idx in (0, 1): + interactions_utils.convert_interaction_event_to_llm_response( + StepStop(event_type='step.stop', index=idx), + state, + interaction_id='int_multi', + ) + + assert state.parts[0].function_call.name == 'get_weather' + assert state.parts[0].function_call.args == {'city': 'Paris'} + assert state.parts[1].function_call.name == 'get_time' + assert state.parts[1].function_call.args == {'zone': 'UTC'} + @pytest.mark.parametrize( ('streamed_events_factory', 'expected_ids'),