From 3a0c1015ac8b1e553beb078563309127c8d6e291 Mon Sep 17 00:00:00 2001 From: Chaitanya Laxman Date: Mon, 7 Sep 2026 10:46:41 +0400 Subject: [PATCH 1/2] fix(models): stream LiteLlm function-call args instead of buffering LiteLlm appended FunctionChunk args until finish_reason. Text already streamed. Gemini already yields partial function-call events when PROGRESSIVE_SSE_STREAMING is on. Yield a partial LlmResponse after each FunctionChunk and keep the aggregated finish event. Fixes #5342 --- src/google/adk/models/lite_llm.py | 33 ++++++ tests/unittests/models/test_litellm.py | 152 +++++++++++++++++-------- 2 files changed, 139 insertions(+), 46 deletions(-) diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index e3560a795d0..b454ab06630 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -58,6 +58,8 @@ from typing_extensions import Required from . import _prompt_cache +from ..features import FeatureName +from ..features import is_feature_enabled from ..utils._google_client_headers import merge_tracking_headers from ..utils._schema_utils import lowercase_schema_types from ._capabilities import LlmCapabilities @@ -2571,6 +2573,32 @@ def _message_to_generate_content_response( ) +def _function_chunk_partial_response( + fc_state: dict[str, Any], + *, + model_version: str | None, +) -> LlmResponse: + """Builds a partial LlmResponse from the current FunctionChunk buffer.""" + accumulated_args = "".join(fc_state["args_parts"]) + function_call = types.FunctionCall( + id=fc_state["id"], + name=fc_state["name"] or None, + will_continue=True, + ) + if accumulated_args: + function_call.partial_args = [ + types.PartialArg(string_value=accumulated_args) + ] + return LlmResponse( + content=types.Content( + role="model", + parts=[types.Part(function_call=function_call)], + ), + partial=True, + model_version=model_version, + ) + + def _finish_reason_to_error_message( finish_reason: types.FinishReason, ) -> str: @@ -3390,6 +3418,11 @@ def _reset_stream_buffers() -> None: function_calls[index]["id"] = ( chunk.id or function_calls[index]["id"] or str(index) ) + if is_feature_enabled(FeatureName.PROGRESSIVE_SSE_STREAMING): + yield _function_chunk_partial_response( + function_calls[index], + model_version=part.model, + ) elif isinstance(chunk, TextChunk): if chunk.text: text_parts.append(chunk.text) diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index 2011908745e..55c46812f40 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -28,6 +28,8 @@ import warnings from google.adk.agents.context_cache_config import ContextCacheConfig +from google.adk.features._feature_registry import FeatureName +from google.adk.features._feature_registry import temporary_feature_override from google.adk.models.lite_llm import _aggregate_streaming_thought_parts from google.adk.models.lite_llm import _append_fallback_user_content_if_missing from google.adk.models.lite_llm import _BraceDepthTracker @@ -4708,7 +4710,7 @@ async def test_completion_additional_args(mock_completion, mock_client): LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True ) ] - assert len(responses) == 4 + assert len(responses) == 6 mock_completion.assert_called_once() _, kwargs = mock_completion.call_args @@ -4736,7 +4738,7 @@ async def test_completion_with_drop_params(mock_completion, mock_client): LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True ) ] - assert len(responses) == 4 + assert len(responses) == 6 mock_completion.assert_called_once() @@ -4806,7 +4808,7 @@ async def test_generate_content_async_stream_tool_call_includes_aggregated_text( LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True ) ] - assert len(responses) == 4 + assert len(responses) == 6 assert responses[0].content.role == "model" assert responses[0].content.parts[0].text == "zero, " assert responses[0].model_version == "test_model" @@ -4816,16 +4818,20 @@ async def test_generate_content_async_stream_tool_call_includes_aggregated_text( assert responses[2].content.role == "model" assert responses[2].content.parts[0].text == "two:" assert responses[2].model_version == "test_model" - assert responses[3].content.role == "model" - assert len(responses[3].content.parts) == 2 - assert responses[3].content.parts[0].text == "zero, one, two:" - assert responses[3].content.parts[1].function_call.name == "test_function" - assert responses[3].content.parts[-1].function_call.args == { + assert responses[3].partial is True + assert responses[3].get_function_calls() + assert responses[4].partial is True + assert responses[4].get_function_calls() + assert responses[5].content.role == "model" + assert len(responses[5].content.parts) == 2 + assert responses[5].content.parts[0].text == "zero, one, two:" + assert responses[5].content.parts[1].function_call.name == "test_function" + assert responses[5].content.parts[-1].function_call.args == { "test_arg": "test_value" } - assert responses[3].content.parts[-1].function_call.id == "test_tool_call_id" - assert responses[3].finish_reason == types.FinishReason.STOP - assert responses[3].model_version == "test_model" + assert responses[5].content.parts[-1].function_call.id == "test_tool_call_id" + assert responses[5].finish_reason == types.FinishReason.STOP + assert responses[5].model_version == "test_model" mock_completion.assert_called_once() _, kwargs = mock_completion.call_args @@ -5079,25 +5085,25 @@ async def test_generate_content_async_stream_with_reasoning_tokens( LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True ) ] - assert len(responses) == 4 + assert len(responses) == 6 assert responses[0].content.role == "model" assert responses[0].content.parts[0].text == "zero, " assert responses[1].content.role == "model" assert responses[1].content.parts[0].text == "one, " assert responses[2].content.role == "model" assert responses[2].content.parts[0].text == "two:" - assert responses[3].content.role == "model" - assert responses[3].content.parts[-1].function_call.name == "test_function" - assert responses[3].content.parts[-1].function_call.args == { + assert responses[-1].content.role == "model" + assert responses[-1].content.parts[-1].function_call.name == "test_function" + assert responses[-1].content.parts[-1].function_call.args == { "test_arg": "test_value" } - assert responses[3].content.parts[-1].function_call.id == "test_tool_call_id" - assert responses[3].finish_reason == types.FinishReason.STOP + assert responses[-1].content.parts[-1].function_call.id == "test_tool_call_id" + assert responses[-1].finish_reason == types.FinishReason.STOP - assert responses[3].usage_metadata.prompt_token_count == 10 - assert responses[3].usage_metadata.candidates_token_count == 5 - assert responses[3].usage_metadata.total_token_count == 15 - assert responses[3].usage_metadata.thoughts_token_count == 5 + assert responses[-1].usage_metadata.prompt_token_count == 10 + assert responses[-1].usage_metadata.candidates_token_count == 5 + assert responses[-1].usage_metadata.total_token_count == 15 + assert responses[-1].usage_metadata.thoughts_token_count == 5 mock_completion.assert_called_once() @@ -5151,12 +5157,12 @@ async def test_generate_content_async_stream_with_usage_metadata( LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True ) ] - assert len(responses) == 4 - assert responses[3].usage_metadata.prompt_token_count == 10 - assert responses[3].usage_metadata.candidates_token_count == 5 - assert responses[3].usage_metadata.total_token_count == 15 - assert responses[3].usage_metadata.cached_content_token_count == 8 - assert responses[3].usage_metadata.thoughts_token_count == 5 + assert len(responses) == 6 + assert responses[-1].usage_metadata.prompt_token_count == 10 + assert responses[-1].usage_metadata.candidates_token_count == 5 + assert responses[-1].usage_metadata.total_token_count == 15 + assert responses[-1].usage_metadata.cached_content_token_count == 8 + assert responses[-1].usage_metadata.thoughts_token_count == 5 @pytest.mark.asyncio @@ -5191,12 +5197,12 @@ async def test_generate_content_async_stream_with_bedrock_cache_tokens( LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True ) ] - assert len(responses) == 4 - assert responses[3].usage_metadata.prompt_token_count == 10 - assert responses[3].usage_metadata.candidates_token_count == 5 - assert responses[3].usage_metadata.total_token_count == 15 - assert responses[3].usage_metadata.cached_content_token_count == 8 - assert responses[3].usage_metadata.cache_creation_input_tokens == 4 + assert len(responses) == 6 + assert responses[-1].usage_metadata.prompt_token_count == 10 + assert responses[-1].usage_metadata.candidates_token_count == 5 + assert responses[-1].usage_metadata.total_token_count == 15 + assert responses[-1].usage_metadata.cached_content_token_count == 8 + assert responses[-1].usage_metadata.cache_creation_input_tokens == 4 @pytest.mark.asyncio @@ -5360,8 +5366,8 @@ async def test_generate_content_async_stream_with_empty_chunk( ) ] - assert len(responses) == 1 - final_response = responses[0] + assert len(responses) == 3 + final_response = responses[-1] assert final_response.content.role == "model" # Crucially, assert that only ONE tool call was generated, @@ -5414,8 +5420,8 @@ async def test_streaming_tool_call_truncated_by_max_tokens( ) ] - assert len(responses) == 1 - error_response = responses[0] + assert len(responses) == 2 + error_response = responses[-1] assert error_response.error_code == types.FinishReason.MAX_TOKENS assert error_response.finish_reason == types.FinishReason.MAX_TOKENS assert "truncated" in error_response.error_message @@ -5462,8 +5468,8 @@ async def test_streaming_tool_call_complete_with_length_finish_reason( ) ] - assert len(responses) == 1 - final_response = responses[0] + assert len(responses) == 2 + final_response = responses[-1] assert final_response.content.role == "model" assert len(final_response.content.parts) == 1 @@ -5515,8 +5521,8 @@ async def test_streaming_tool_call_malformed_arguments_returns_empty( ) ] - assert len(responses) == 1 - final_response = responses[0] + assert len(responses) == 2 + final_response = responses[-1] assert final_response.content.role == "model" function_call = final_response.content.parts[0].function_call assert function_call.name == "test_function" @@ -7353,8 +7359,8 @@ async def test_streaming_tool_call_args_assembled_from_many_fragments( ) ] - assert len(responses) == 1 - function_call = responses[0].content.parts[0].function_call + assert len(responses) == len(fragments) + 1 + function_call = responses[-1].content.parts[0].function_call assert function_call.name == "my_func" assert function_call.id == "call_xyz" assert function_call.args == json.loads(full_args) @@ -7439,14 +7445,65 @@ async def test_streaming_tool_call_brace_in_string_does_not_falsely_complete( ) ] - assert len(responses) == 1 - parts = responses[0].content.parts + assert len(responses) == len(function_chunks) + 1 + parts = responses[-1].content.parts assert len(parts) == 2 args_by_name = {p.function_call.name: p.function_call.args for p in parts} assert args_by_name["my_func"] == json.loads(full_args_a) assert args_by_name["other_func"] == json.loads(full_args_b) +@pytest.mark.asyncio +async def test_streaming_function_chunks_yield_partials_when_progressive_sse_on( + mock_completion, lite_llm_instance +): + fragments = ['{"city": "', "Paris", '"}'] + mock_completion.return_value = iter( + _stream_chunks_from_function_chunks(_function_chunks_for_args(fragments)) + ) + + responses = [ + r + async for r in lite_llm_instance.generate_content_async( + LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True + ) + ] + + partials = [r for r in responses if r.partial] + assert len(partials) == len(fragments) + accumulated = "" + for fragment, partial in zip(fragments, partials): + assert partial.get_function_calls() + function_call = partial.get_function_calls()[0] + assert function_call.will_continue is True + accumulated += fragment + assert function_call.partial_args[0].string_value == accumulated + assert responses[-1].partial is False + assert responses[-1].content.parts[0].function_call.args == {"city": "Paris"} + + +@pytest.mark.asyncio +async def test_streaming_function_chunks_stay_buffered_when_progressive_sse_off( + mock_completion, lite_llm_instance +): + fragments = ['{"city": "', "Paris", '"}'] + mock_completion.return_value = iter( + _stream_chunks_from_function_chunks(_function_chunks_for_args(fragments)) + ) + + with temporary_feature_override(FeatureName.PROGRESSIVE_SSE_STREAMING, False): + responses = [ + r + async for r in lite_llm_instance.generate_content_async( + LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True + ) + ] + + assert len(responses) == 1 + assert responses[0].partial is False + assert responses[0].content.parts[0].function_call.args == {"city": "Paris"} + + def _text_stream_chunks(text_fragments, finish_reason="stop"): stream = [ ModelResponseStream( @@ -7509,7 +7566,10 @@ async def test_streaming_buffers_hold_fragments_instead_of_growing_copies( LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True ) try: - # Suspends on the first partial text response, with both buffers filled. + # FunctionChunks now yield progressive partials. Drain those first so + # the next item is the first text partial, with both buffers filled. + for _ in arg_fragments: + await responses.__anext__() await responses.__anext__() buffers = responses.ag_frame.f_locals assert buffers["text_parts"] == text_fragments[:1] From 57a6f9bb1821d5e8db6c3f65e341c4eb873d0f94 Mon Sep 17 00:00:00 2001 From: Chaitanya Laxman Date: Mon, 7 Sep 2026 11:03:34 +0400 Subject: [PATCH 2/2] fix(models): emit LiteLlm function-call arg deltas not prefixes partial_args.string_value is a per-event delta in interactions_utils and Gemini. Joining the buffer on every FunctionChunk duplicated JSON for consumers that concatenate string_value. Fixes #5342 --- src/google/adk/models/lite_llm.py | 11 +++-- tests/unittests/models/test_litellm.py | 60 +++++++++++++------------- 2 files changed, 34 insertions(+), 37 deletions(-) diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index b454ab06630..aa116fe9180 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -2576,19 +2576,17 @@ def _message_to_generate_content_response( def _function_chunk_partial_response( fc_state: dict[str, Any], *, + args_delta: str | None, model_version: str | None, ) -> LlmResponse: - """Builds a partial LlmResponse from the current FunctionChunk buffer.""" - accumulated_args = "".join(fc_state["args_parts"]) + """Builds a partial LlmResponse for one FunctionChunk args delta.""" function_call = types.FunctionCall( id=fc_state["id"], name=fc_state["name"] or None, will_continue=True, ) - if accumulated_args: - function_call.partial_args = [ - types.PartialArg(string_value=accumulated_args) - ] + if args_delta: + function_call.partial_args = [types.PartialArg(string_value=args_delta)] return LlmResponse( content=types.Content( role="model", @@ -3421,6 +3419,7 @@ def _reset_stream_buffers() -> None: if is_feature_enabled(FeatureName.PROGRESSIVE_SSE_STREAMING): yield _function_chunk_partial_response( function_calls[index], + args_delta=chunk.args, model_version=part.model, ) elif isinstance(chunk, TextChunk): diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index 55c46812f40..281d15f342f 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -4710,7 +4710,7 @@ async def test_completion_additional_args(mock_completion, mock_client): LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True ) ] - assert len(responses) == 6 + assert [r for r in responses if not r.partial] mock_completion.assert_called_once() _, kwargs = mock_completion.call_args @@ -4738,7 +4738,7 @@ async def test_completion_with_drop_params(mock_completion, mock_client): LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True ) ] - assert len(responses) == 6 + assert [r for r in responses if not r.partial] mock_completion.assert_called_once() @@ -4802,12 +4802,13 @@ async def test_generate_content_async_stream_tool_call_includes_aggregated_text( mock_completion.return_value = iter(STREAMING_MODEL_RESPONSE) - responses = [ - response - async for response in lite_llm_instance.generate_content_async( - LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True - ) - ] + with temporary_feature_override(FeatureName.PROGRESSIVE_SSE_STREAMING, True): + responses = [ + response + async for response in lite_llm_instance.generate_content_async( + LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True + ) + ] assert len(responses) == 6 assert responses[0].content.role == "model" assert responses[0].content.parts[0].text == "zero, " @@ -5085,7 +5086,7 @@ async def test_generate_content_async_stream_with_reasoning_tokens( LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True ) ] - assert len(responses) == 6 + assert responses[-1].partial is False assert responses[0].content.role == "model" assert responses[0].content.parts[0].text == "zero, " assert responses[1].content.role == "model" @@ -5157,7 +5158,7 @@ async def test_generate_content_async_stream_with_usage_metadata( LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True ) ] - assert len(responses) == 6 + assert responses[-1].partial is False assert responses[-1].usage_metadata.prompt_token_count == 10 assert responses[-1].usage_metadata.candidates_token_count == 5 assert responses[-1].usage_metadata.total_token_count == 15 @@ -5197,7 +5198,7 @@ async def test_generate_content_async_stream_with_bedrock_cache_tokens( LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True ) ] - assert len(responses) == 6 + assert responses[-1].partial is False assert responses[-1].usage_metadata.prompt_token_count == 10 assert responses[-1].usage_metadata.candidates_token_count == 5 assert responses[-1].usage_metadata.total_token_count == 15 @@ -5366,8 +5367,8 @@ async def test_generate_content_async_stream_with_empty_chunk( ) ] - assert len(responses) == 3 final_response = responses[-1] + assert final_response.partial is False assert final_response.content.role == "model" # Crucially, assert that only ONE tool call was generated, @@ -5420,8 +5421,8 @@ async def test_streaming_tool_call_truncated_by_max_tokens( ) ] - assert len(responses) == 2 error_response = responses[-1] + assert not error_response.partial assert error_response.error_code == types.FinishReason.MAX_TOKENS assert error_response.finish_reason == types.FinishReason.MAX_TOKENS assert "truncated" in error_response.error_message @@ -5468,8 +5469,8 @@ async def test_streaming_tool_call_complete_with_length_finish_reason( ) ] - assert len(responses) == 2 final_response = responses[-1] + assert final_response.partial is False assert final_response.content.role == "model" assert len(final_response.content.parts) == 1 @@ -5521,8 +5522,8 @@ async def test_streaming_tool_call_malformed_arguments_returns_empty( ) ] - assert len(responses) == 2 final_response = responses[-1] + assert final_response.partial is False assert final_response.content.role == "model" function_call = final_response.content.parts[0].function_call assert function_call.name == "test_function" @@ -7359,7 +7360,6 @@ async def test_streaming_tool_call_args_assembled_from_many_fragments( ) ] - assert len(responses) == len(fragments) + 1 function_call = responses[-1].content.parts[0].function_call assert function_call.name == "my_func" assert function_call.id == "call_xyz" @@ -7445,7 +7445,6 @@ async def test_streaming_tool_call_brace_in_string_does_not_falsely_complete( ) ] - assert len(responses) == len(function_chunks) + 1 parts = responses[-1].content.parts assert len(parts) == 2 args_by_name = {p.function_call.name: p.function_call.args for p in parts} @@ -7462,22 +7461,21 @@ async def test_streaming_function_chunks_yield_partials_when_progressive_sse_on( _stream_chunks_from_function_chunks(_function_chunks_for_args(fragments)) ) - responses = [ - r - async for r in lite_llm_instance.generate_content_async( - LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True - ) - ] + with temporary_feature_override(FeatureName.PROGRESSIVE_SSE_STREAMING, True): + responses = [ + r + async for r in lite_llm_instance.generate_content_async( + LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True + ) + ] partials = [r for r in responses if r.partial] assert len(partials) == len(fragments) - accumulated = "" for fragment, partial in zip(fragments, partials): assert partial.get_function_calls() function_call = partial.get_function_calls()[0] assert function_call.will_continue is True - accumulated += fragment - assert function_call.partial_args[0].string_value == accumulated + assert function_call.partial_args[0].string_value == fragment assert responses[-1].partial is False assert responses[-1].content.parts[0].function_call.args == {"city": "Paris"} @@ -7566,11 +7564,11 @@ async def test_streaming_buffers_hold_fragments_instead_of_growing_copies( LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True ) try: - # FunctionChunks now yield progressive partials. Drain those first so - # the next item is the first text partial, with both buffers filled. - for _ in arg_fragments: - await responses.__anext__() - await responses.__anext__() + # Drain until the first text partial, so both buffers are filled. + while True: + partial = await responses.__anext__() + if partial.content and any(p.text for p in partial.content.parts or []): + break buffers = responses.ag_frame.f_locals assert buffers["text_parts"] == text_fragments[:1] assert buffers["function_calls"][0]["args_parts"] == arg_fragments