diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index e3560a795d0..aa116fe9180 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,30 @@ 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 for one FunctionChunk args delta.""" + function_call = types.FunctionCall( + id=fc_state["id"], + name=fc_state["name"] or None, + will_continue=True, + ) + if args_delta: + function_call.partial_args = [types.PartialArg(string_value=args_delta)] + 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 +3416,12 @@ 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], + args_delta=chunk.args, + 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..281d15f342f 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 [r for r in responses if not r.partial] 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 [r for r in responses if not r.partial] mock_completion.assert_called_once() @@ -4800,13 +4802,14 @@ 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 - ) - ] - assert len(responses) == 4 + 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, " assert responses[0].model_version == "test_model" @@ -4816,16 +4819,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 +5086,25 @@ async def test_generate_content_async_stream_with_reasoning_tokens( LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True ) ] - assert len(responses) == 4 + 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" 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 +5158,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 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 + assert responses[-1].usage_metadata.cached_content_token_count == 8 + assert responses[-1].usage_metadata.thoughts_token_count == 5 @pytest.mark.asyncio @@ -5191,12 +5198,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 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 + 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 +5367,8 @@ async def test_generate_content_async_stream_with_empty_chunk( ) ] - assert len(responses) == 1 - final_response = responses[0] + 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, @@ -5414,8 +5421,8 @@ async def test_streaming_tool_call_truncated_by_max_tokens( ) ] - assert len(responses) == 1 - error_response = responses[0] + 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 @@ -5462,8 +5469,8 @@ async def test_streaming_tool_call_complete_with_length_finish_reason( ) ] - assert len(responses) == 1 - final_response = responses[0] + final_response = responses[-1] + assert final_response.partial is False assert final_response.content.role == "model" assert len(final_response.content.parts) == 1 @@ -5515,8 +5522,8 @@ async def test_streaming_tool_call_malformed_arguments_returns_empty( ) ] - assert len(responses) == 1 - final_response = responses[0] + 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" @@ -7353,8 +7360,7 @@ async def test_streaming_tool_call_args_assembled_from_many_fragments( ) ] - assert len(responses) == 1 - function_call = responses[0].content.parts[0].function_call + 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,63 @@ async def test_streaming_tool_call_brace_in_string_does_not_falsely_complete( ) ] - assert len(responses) == 1 - parts = responses[0].content.parts + 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)) + ) + + 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) + 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 + 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"} + + +@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,8 +7564,11 @@ 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. - 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