diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index 5bf71169d9..3459d750d3 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -3328,6 +3328,15 @@ def _reset_stream_buffers() -> None: part_grounding = _extract_grounding_metadata(part) if part_grounding: grounding_metadata = part_grounding + # One delta can carry several complete tool calls, and + # _model_response_to_chunk yields a FunctionChunk per call, each paired + # with the same choice-level finish_reason. Record what this part asks + # for and finalize once, after every chunk in it has been accumulated: + # finalizing inside the loop rebuilds the response per call and keeps + # only the last one. + finalize_tool_calls = False + finalize_text = False + finalize_finish_reason: str | None = None for chunk, finish_reason in _model_response_to_chunk(part): if finish_reason: last_finish_reason = finish_reason @@ -3401,13 +3410,8 @@ def _reset_stream_buffers() -> None: or finish_reason == "length" or (finish_reason == "stop" and chunk is None) ): - aggregated_llm_response_with_tool_call = ( - _finalize_tool_call_response( - model_version=part.model, - finish_reason=finish_reason, - ) - ) - _reset_stream_buffers() + finalize_tool_calls = True + finalize_finish_reason = finish_reason elif (text_parts or reasoning_parts) and ( finish_reason == "length" or ( @@ -3416,11 +3420,24 @@ def _reset_stream_buffers() -> None: and not function_calls ) ): - aggregated_llm_response = _finalize_text_response( - model_version=part.model, - finish_reason=finish_reason, - ) - _reset_stream_buffers() + finalize_text = True + finalize_finish_reason = finish_reason + + # Tool calls take precedence for a part that asks for both: + # _finalize_tool_call_response carries the buffered text and reasoning + # into the tool-call response, so nothing accumulated here is dropped. + if finalize_tool_calls: + aggregated_llm_response_with_tool_call = _finalize_tool_call_response( + model_version=part.model, + finish_reason=finalize_finish_reason, + ) + _reset_stream_buffers() + elif finalize_text: + aggregated_llm_response = _finalize_text_response( + model_version=part.model, + finish_reason=finalize_finish_reason, + ) + _reset_stream_buffers() # The in-loop finalizers only fire on the reasons known to end a stream, # so any other terminal reason ("content_filter" above all) reaches the diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index 2efa54d566..bb3371566b 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -1401,6 +1401,47 @@ def mock_response(): ] +def _same_chunk_multiple_function_calls_stream( + finish_reason, content=None, reasoning_content=None +): + """A single delta carrying two complete tool calls and a finish reason. + + Providers and litellm transformations that emit a whole assistant message as + one stream chunk produce this shape, as do custom ``LiteLlm.llm_client`` + implementations. ``_model_response_to_chunk`` yields one ``FunctionChunk`` per + tool call, each paired with the same choice-level finish reason. + """ + delta = Delta( + role="assistant", + content=content, + tool_calls=[ + ChatCompletionDeltaToolCall( + type="function", + id="call_1", + function=Function( + name="function_1", arguments='{"arg": "value1"}' + ), + index=0, + ), + ChatCompletionDeltaToolCall( + type="function", + id="call_2", + function=Function( + name="function_2", arguments='{"arg": "value2"}' + ), + index=1, + ), + ], + ) + if reasoning_content: + delta.reasoning_content = reasoning_content + return [ + ModelResponseStream( + choices=[StreamingChoices(finish_reason=finish_reason, delta=delta)] + ) + ] + + @pytest.fixture def mock_acompletion(mock_response): return AsyncMock(return_value=mock_response) @@ -5155,6 +5196,151 @@ async def test_generate_content_async_non_compliant_multiple_function_calls( assert final_response.content.parts[1].function_call.args == {"arg": "value2"} +@pytest.mark.asyncio +@pytest.mark.parametrize("finish_reason", ["tool_calls", "length"]) +async def test_generate_content_async_same_chunk_multiple_function_calls( + mock_completion, lite_llm_instance, finish_reason +): + """Every tool call in one finish-bearing delta reaches the final response. + + The finalization is decided once per streamed part. Deciding it per chunk + rebuilt the aggregated response for each tool call and kept only the last. + """ + mock_completion.return_value = _same_chunk_multiple_function_calls_stream( + finish_reason + ) + + llm_request = LlmRequest( + contents=[ + types.Content( + role="user", + parts=[types.Part.from_text(text="Test same-chunk calls")], + ) + ], + config=types.GenerateContentConfig( + tools=[ + types.Tool( + function_declarations=[ + types.FunctionDeclaration( + name="function_1", + description="First test function", + parameters=types.Schema( + type=types.Type.OBJECT, + properties={ + "arg": types.Schema(type=types.Type.STRING), + }, + ), + ), + types.FunctionDeclaration( + name="function_2", + description="Second test function", + parameters=types.Schema( + type=types.Type.OBJECT, + properties={ + "arg": types.Schema(type=types.Type.STRING), + }, + ), + ), + ] + ) + ], + ), + ) + + responses = [ + response + async for response in lite_llm_instance.generate_content_async( + llm_request, stream=True + ) + if not response.partial + ] + + function_calls = [ + part.function_call + for response in responses + if response.content + for part in response.content.parts + if part.function_call + ] + assert [call.name for call in function_calls] == [ + "function_1", + "function_2", + ] + assert [call.id for call in function_calls] == ["call_1", "call_2"] + assert [call.args for call in function_calls] == [ + {"arg": "value1"}, + {"arg": "value2"}, + ] + + +@pytest.mark.asyncio +async def test_generate_content_async_same_chunk_calls_keep_text_and_reasoning( + mock_completion, lite_llm_instance +): + """Text and reasoning sharing the delta survive alongside every tool call.""" + mock_completion.return_value = _same_chunk_multiple_function_calls_stream( + "tool_calls", content="Calling both.", reasoning_content="Thinking." + ) + + llm_request = LlmRequest( + contents=[ + types.Content( + role="user", + parts=[types.Part.from_text(text="Test same-chunk calls")], + ) + ], + config=types.GenerateContentConfig( + tools=[ + types.Tool( + function_declarations=[ + types.FunctionDeclaration( + name="function_1", + description="First test function", + parameters=types.Schema( + type=types.Type.OBJECT, + properties={ + "arg": types.Schema(type=types.Type.STRING), + }, + ), + ), + types.FunctionDeclaration( + name="function_2", + description="Second test function", + parameters=types.Schema( + type=types.Type.OBJECT, + properties={ + "arg": types.Schema(type=types.Type.STRING), + }, + ), + ), + ] + ) + ], + ), + ) + + responses = [ + response + async for response in lite_llm_instance.generate_content_async( + llm_request, stream=True + ) + if not response.partial + ] + + parts = [ + part + for response in responses + if response.content + for part in response.content.parts + ] + assert [part.function_call.name for part in parts if part.function_call] == [ + "function_1", + "function_2", + ] + assert any(part.text == "Calling both." for part in parts) + assert any(part.thought and part.text == "Thinking." for part in parts) + + @pytest.mark.asyncio async def test_generate_content_async_stream_with_empty_chunk( mock_completion, lite_llm_instance