Files
google--adk-python/tests/unittests/utils/test_streaming_utils.py
T
wehub-resource-sync ec2b666284
Continuous Integration / Pre-commit Linter (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.10) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.11) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.12) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.13) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.10) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.11) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.12) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.13) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.14) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Waiting to run
Copybara PR Handler / close-imported-pr (push) Waiting to run
chore: import upstream snapshot with attribution
2026-07-13 13:25:13 +08:00

719 lines
25 KiB
Python

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from google.adk.features._feature_registry import FeatureName
from google.adk.features._feature_registry import temporary_feature_override
from google.adk.flows.llm_flows.functions import AF_FUNCTION_CALL_ID_PREFIX
from google.adk.utils import streaming_utils
from google.genai import types
import pytest
class TestStreamingResponseAggregator:
@pytest.mark.asyncio
async def test_process_response_with_text(self):
aggregator = streaming_utils.StreamingResponseAggregator()
response = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(parts=[types.Part(text="Hello")])
)
]
)
results = []
async for r in aggregator.process_response(response):
results.append(r)
assert len(results) == 1
assert results[0].content.parts[0].text == "Hello"
assert results[0].partial
@pytest.mark.asyncio
async def test_process_response_with_thought(self):
aggregator = streaming_utils.StreamingResponseAggregator()
response = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
parts=[types.Part(text="Thinking...", thought=True)]
)
)
]
)
results = []
async for r in aggregator.process_response(response):
results.append(r)
assert len(results) == 1
assert results[0].content.parts[0].text == "Thinking..."
assert results[0].content.parts[0].thought
assert results[0].partial
@pytest.mark.asyncio
async def test_process_response_multiple(self):
aggregator = streaming_utils.StreamingResponseAggregator()
response1 = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(parts=[types.Part(text="Hello ")])
)
]
)
response2 = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(parts=[types.Part(text="World!")])
)
]
)
async for _ in aggregator.process_response(response1):
pass
results = []
async for r in aggregator.process_response(response2):
results.append(r)
assert len(results) == 1
assert results[0].content.parts[0].text == "World!"
closed_response = aggregator.close()
assert closed_response is not None
assert closed_response.content.parts[0].text == "Hello World!"
@pytest.mark.asyncio
async def test_process_response_interleaved_thought_and_text(self):
aggregator = streaming_utils.StreamingResponseAggregator()
response1 = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
parts=[types.Part(text="I am thinking...", thought=True)]
)
)
]
)
response2 = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
parts=[types.Part(text="Okay, I have a result.")]
)
)
]
)
response3 = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
parts=[types.Part(text=" The result is 42.")]
)
)
]
)
async for _ in aggregator.process_response(response1):
pass
async for _ in aggregator.process_response(response2):
pass
async for _ in aggregator.process_response(response3):
pass
closed_response = aggregator.close()
assert closed_response is not None
assert len(closed_response.content.parts) == 2
assert closed_response.content.parts[0].text == "I am thinking..."
assert closed_response.content.parts[0].thought
assert (
closed_response.content.parts[1].text
== "Okay, I have a result. The result is 42."
)
assert not closed_response.content.parts[1].thought
def test_close_with_no_responses(self):
aggregator = streaming_utils.StreamingResponseAggregator()
closed_response = aggregator.close()
assert closed_response is None
@pytest.mark.asyncio
async def test_close_with_finish_reason(self):
aggregator = streaming_utils.StreamingResponseAggregator()
response = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(parts=[types.Part(text="Hello")]),
finish_reason=types.FinishReason.STOP,
)
]
)
async for _ in aggregator.process_response(response):
pass
closed_response = aggregator.close()
assert closed_response is not None
assert closed_response.content.parts[0].text == "Hello"
assert closed_response.error_code is None
assert closed_response.error_message is None
@pytest.mark.asyncio
async def test_close_with_error(self):
aggregator = streaming_utils.StreamingResponseAggregator()
response = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(parts=[types.Part(text="Error")]),
finish_reason=types.FinishReason.RECITATION,
finish_message="Recitation error",
)
]
)
async for _ in aggregator.process_response(response):
pass
closed_response = aggregator.close()
assert closed_response is not None
assert closed_response.content.parts[0].text == "Error"
assert closed_response.error_code == types.FinishReason.RECITATION
assert closed_response.error_message == "Recitation error"
@pytest.mark.asyncio
@pytest.mark.parametrize("use_progressive_sse", [True, False])
async def test_empty_content_produces_empty_final_frame(
self, use_progressive_sse
):
"""A candidate with empty parts + STOP passes through without an error.
A terminal empty STOP chunk must not be classified as an error at the
streaming layer; consumers that batch parts across chunks rely on it
passing through cleanly.
"""
with temporary_feature_override(
FeatureName.PROGRESSIVE_SSE_STREAMING, use_progressive_sse
):
aggregator = streaming_utils.StreamingResponseAggregator()
response = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(parts=[]),
finish_reason=types.FinishReason.STOP,
)
]
)
results = []
async for r in aggregator.process_response(response):
results.append(r)
closed_response = aggregator.close()
assert len(results) == 1
assert results[0].content is not None
assert results[0].error_code is None
assert closed_response is not None
assert closed_response.partial is False
assert closed_response.content is None
assert closed_response.finish_reason == types.FinishReason.STOP
@pytest.mark.asyncio
@pytest.mark.parametrize("use_progressive_sse", [True, False])
async def test_prompt_feedback_block_returns_error_frame(
self, use_progressive_sse
):
"""A prompt-level safety block produces a final frame with the error code."""
with temporary_feature_override(
FeatureName.PROGRESSIVE_SSE_STREAMING, use_progressive_sse
):
aggregator = streaming_utils.StreamingResponseAggregator()
response = types.GenerateContentResponse(
prompt_feedback=types.GenerateContentResponsePromptFeedback(
block_reason=types.BlockedReason.SAFETY,
block_reason_message="Blocked by safety",
)
)
results = []
async for r in aggregator.process_response(response):
results.append(r)
closed_response = aggregator.close()
assert len(results) == 1
assert closed_response is not None
assert closed_response.partial is False
assert closed_response.error_code == types.BlockedReason.SAFETY
assert closed_response.error_message == "Blocked by safety"
assert closed_response.content is None
@pytest.mark.asyncio
@pytest.mark.parametrize("use_progressive_sse", [True, False])
async def test_pure_function_call_behavior_differs_by_mode(
self, use_progressive_sse
):
"""A pure function call yields the part in progressive mode and an empty frame otherwise."""
with temporary_feature_override(
FeatureName.PROGRESSIVE_SSE_STREAMING, use_progressive_sse
):
aggregator = streaming_utils.StreamingResponseAggregator()
response = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
parts=[
types.Part(
function_call=types.FunctionCall(
name="my_tool",
args={"x": 1},
)
)
]
),
finish_reason=types.FinishReason.STOP,
)
]
)
results = []
async for r in aggregator.process_response(response):
results.append(r)
closed_response = aggregator.close()
assert closed_response is not None
assert closed_response.partial is False
if use_progressive_sse:
assert closed_response.content is not None
assert len(closed_response.content.parts) == 1
assert closed_response.content.parts[0].function_call.name == "my_tool"
else:
assert closed_response.content is None
@pytest.mark.asyncio
@pytest.mark.parametrize(
"test_id, use_progressive_sse, metadata_type",
[
("grounding_default", False, "grounding"),
("grounding_progressive", True, "grounding"),
("citation_default", False, "citation"),
("citation_progressive", True, "citation"),
],
)
async def test_close_preserves_metadata(
self, test_id, use_progressive_sse, metadata_type
):
"""close() should carry metadata into the aggregated response."""
aggregator = streaming_utils.StreamingResponseAggregator()
metadata = None
response1 = None
response2 = None
if metadata_type == "grounding":
metadata = types.GroundingMetadata(
grounding_chunks=[
types.GroundingChunk(
retrieved_context=types.GroundingChunkRetrievedContext(
uri="https://example.com/doc1",
title="Source",
)
)
],
)
response1 = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(parts=[types.Part(text="Hello ")]),
grounding_metadata=metadata,
)
]
)
response2 = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(parts=[types.Part(text="World!")]),
finish_reason=types.FinishReason.STOP,
grounding_metadata=metadata,
)
]
)
elif metadata_type == "citation":
metadata = types.CitationMetadata(
citations=[
types.Citation(
start_index=0,
end_index=10,
uri="https://example.com/source",
title="Source",
)
]
)
response1 = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(parts=[types.Part(text="Cited text")]),
)
]
)
response2 = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(parts=[]),
finish_reason=types.FinishReason.STOP,
citation_metadata=metadata,
)
]
)
async def run_test():
async for _ in aggregator.process_response(response1):
pass
async for _ in aggregator.process_response(response2):
pass
closed_response = aggregator.close()
assert closed_response is not None
if use_progressive_sse:
assert closed_response.partial is False
if metadata_type == "grounding":
assert closed_response.grounding_metadata is not None
assert len(closed_response.grounding_metadata.grounding_chunks) == 1
elif metadata_type == "citation":
assert closed_response.citation_metadata is not None
assert len(closed_response.citation_metadata.citations) == 1
if use_progressive_sse:
with temporary_feature_override(
FeatureName.PROGRESSIVE_SSE_STREAMING, True
):
await run_test()
else:
await run_test()
@pytest.mark.asyncio
@pytest.mark.parametrize("use_progressive_sse", [False, True])
async def test_close_propagates_model_version(self, use_progressive_sse):
"""close() should carry model_version into the aggregated response."""
aggregator = streaming_utils.StreamingResponseAggregator()
response1 = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(parts=[types.Part(text="Hello ")]),
)
],
model_version="gemini-test-1.0",
)
response2 = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(parts=[types.Part(text="World!")]),
finish_reason=types.FinishReason.STOP,
)
],
model_version="gemini-test-1.0",
)
async def run_test():
async for _ in aggregator.process_response(response1):
pass
async for _ in aggregator.process_response(response2):
pass
closed_response = aggregator.close()
assert closed_response is not None
assert closed_response.model_version == "gemini-test-1.0"
if use_progressive_sse:
with temporary_feature_override(
FeatureName.PROGRESSIVE_SSE_STREAMING, True
):
await run_test()
else:
await run_test()
@pytest.mark.asyncio
async def test_non_progressive_merged_yield_propagates_model_version(self):
"""The mid-stream merged text event should carry model_version forward.
In non-progressive mode, when a new non-text response arrives after buffered
text, the aggregator yields a synthesized merged-text LlmResponse before
yielding the current partial. That merged event must preserve fields from
the source response (model_version, grounding_metadata, citation_metadata,
finish_reason).
"""
# PROGRESSIVE_SSE_STREAMING defaults to on; explicitly disable it to
# exercise the non-progressive merged-yield code path under test.
with temporary_feature_override(
FeatureName.PROGRESSIVE_SSE_STREAMING, False
):
aggregator = streaming_utils.StreamingResponseAggregator()
# First: buffer some text.
response1 = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
parts=[types.Part(text="Hello World!")]
),
)
],
model_version="gemini-test-2.0",
)
# Second: a response without text triggers the merged yield path.
response2 = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(parts=[]),
finish_reason=types.FinishReason.STOP,
)
],
model_version="gemini-test-2.0",
)
results = []
async for r in aggregator.process_response(response1):
results.append(r)
async for r in aggregator.process_response(response2):
results.append(r)
# The synthesized merged-text event should carry model_version.
merged_events = [
r
for r in results
if r.content
and r.content.parts
and r.content.parts[0].text == "Hello World!"
and not r.partial
]
assert merged_events, "expected a merged non-partial text event"
assert merged_events[0].model_version == "gemini-test-2.0"
class TestFunctionCallIdGeneration:
"""Tests for function call ID generation in streaming mode."""
@pytest.mark.asyncio
async def test_non_streaming_fc_generates_id_when_empty(self):
"""Non-streaming function call should get an adk-* ID if LLM didn't provide one."""
with temporary_feature_override(
FeatureName.PROGRESSIVE_SSE_STREAMING, True
):
aggregator = streaming_utils.StreamingResponseAggregator()
response = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
parts=[
types.Part(
function_call=types.FunctionCall(
name="my_tool",
args={"x": 1},
id=None, # No ID from LLM
)
)
]
),
finish_reason=types.FinishReason.STOP,
)
]
)
async for _ in aggregator.process_response(response):
pass
closed_response = aggregator.close()
assert closed_response is not None
fc = closed_response.content.parts[0].function_call
assert fc.id is not None
assert fc.id.startswith(AF_FUNCTION_CALL_ID_PREFIX)
@pytest.mark.asyncio
async def test_non_streaming_fc_preserves_llm_assigned_id(self):
"""Non-streaming function call should preserve ID if LLM provided one."""
with temporary_feature_override(
FeatureName.PROGRESSIVE_SSE_STREAMING, True
):
aggregator = streaming_utils.StreamingResponseAggregator()
response = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
parts=[
types.Part(
function_call=types.FunctionCall(
name="my_tool",
args={"x": 1},
id="llm-assigned-id",
)
)
]
),
finish_reason=types.FinishReason.STOP,
)
]
)
async for _ in aggregator.process_response(response):
pass
closed_response = aggregator.close()
assert closed_response is not None
fc = closed_response.content.parts[0].function_call
assert fc.id == "llm-assigned-id"
@pytest.mark.asyncio
async def test_streaming_fc_generates_consistent_id_across_chunks(self):
"""Streaming function call should have the same ID in partial and final responses."""
with temporary_feature_override(
FeatureName.PROGRESSIVE_SSE_STREAMING, True
):
aggregator = streaming_utils.StreamingResponseAggregator()
# First chunk: function call starts
response1 = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
parts=[
types.Part(
function_call=types.FunctionCall(
name="my_tool",
id=None,
partial_args=[
types.PartialArg(
json_path="$.x",
string_value="hello",
)
],
will_continue=True,
)
)
]
)
)
]
)
# Second chunk: function call continues
response2 = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
parts=[
types.Part(
function_call=types.FunctionCall(
name=None,
id=None,
partial_args=[
types.PartialArg(
json_path="$.x",
string_value=" world",
)
],
will_continue=False, # Complete
)
)
]
),
finish_reason=types.FinishReason.STOP,
)
]
)
partial_results = []
async for r in aggregator.process_response(response1):
partial_results.append(r)
async for r in aggregator.process_response(response2):
partial_results.append(r)
closed_response = aggregator.close()
assert closed_response is not None
final_fc = closed_response.content.parts[0].function_call
assert final_fc.id is not None
assert final_fc.id.startswith(AF_FUNCTION_CALL_ID_PREFIX)
assert final_fc.args == {"x": "hello world"}
# Verify partial and final events share the same ID
partial_fc = partial_results[0].content.parts[0].function_call
assert (
partial_fc.id == final_fc.id
), f"Partial FC ID ({partial_fc.id!r}) != Final FC ID ({final_fc.id!r})"
@pytest.mark.asyncio
async def test_multiple_streaming_fcs_get_different_ids(self):
"""Multiple function calls arriving in separate chunks should get different IDs."""
with temporary_feature_override(
FeatureName.PROGRESSIVE_SSE_STREAMING, True
):
aggregator = streaming_utils.StreamingResponseAggregator()
# First FC
response1 = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
parts=[
types.Part(
function_call=types.FunctionCall(
name="tool_a",
id=None,
partial_args=[
types.PartialArg(
json_path="$.a", string_value="val_a"
)
],
will_continue=False,
)
)
]
)
)
]
)
# Second FC
response2 = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
parts=[
types.Part(
function_call=types.FunctionCall(
name="tool_b",
id=None,
partial_args=[
types.PartialArg(
json_path="$.b", string_value="val_b"
)
],
will_continue=False,
)
)
]
),
finish_reason=types.FinishReason.STOP,
)
]
)
async for _ in aggregator.process_response(response1):
pass
async for _ in aggregator.process_response(response2):
pass
closed_response = aggregator.close()
assert closed_response is not None
assert len(closed_response.content.parts) == 2
fc_a = closed_response.content.parts[0].function_call
fc_b = closed_response.content.parts[1].function_call
assert fc_a.id is not None
assert fc_b.id is not None
assert fc_a.id.startswith(AF_FUNCTION_CALL_ID_PREFIX)
assert fc_b.id.startswith(AF_FUNCTION_CALL_ID_PREFIX)
assert fc_a.id != fc_b.id # Different IDs for different FCs