diff --git a/src/google/adk/utils/streaming_utils.py b/src/google/adk/utils/streaming_utils.py index 1e8ac05a..fd5fd4ad 100644 --- a/src/google/adk/utils/streaming_utils.py +++ b/src/google/adk/utils/streaming_utils.py @@ -35,7 +35,9 @@ class StreamingResponseAggregator: def __init__(self) -> None: self._text: list[str] = [] self._thought_text: list[str] = [] - self._usage_metadata = None + self._usage_metadata: Optional[ + types.GenerateContentResponseUsageMetadata + ] = None self._grounding_metadata: Optional[types.GroundingMetadata] = None self._citation_metadata: Optional[types.CitationMetadata] = None self._response = None @@ -264,7 +266,10 @@ class StreamingResponseAggregator: # results = [] self._response = response llm_response = LlmResponse.create(response) - self._usage_metadata = llm_response.usage_metadata + # Usage is typically reported on a single chunk; keep the last reported + # value rather than letting a usage-less trailing chunk erase it. + if llm_response.usage_metadata: + self._usage_metadata = llm_response.usage_metadata if llm_response.grounding_metadata: self._grounding_metadata = llm_response.grounding_metadata if llm_response.citation_metadata: diff --git a/tests/unittests/utils/test_streaming_utils.py b/tests/unittests/utils/test_streaming_utils.py index ab03bc3c..a2dd0dae 100644 --- a/tests/unittests/utils/test_streaming_utils.py +++ b/tests/unittests/utils/test_streaming_utils.py @@ -392,6 +392,102 @@ class TestStreamingResponseAggregator: else: await run_test() + @pytest.mark.asyncio + @pytest.mark.parametrize("use_progressive_sse", [False, True]) + async def test_close_preserves_usage_metadata_from_earlier_chunk( + self, use_progressive_sse + ): + """A later chunk without usage must not erase an earlier chunk's counts. + + Providers typically report token usage on a single chunk; the trailing + chunks of the same turn carry none. The aggregated response is the one + that gets persisted, so it must retain the counts it already saw. + """ + with temporary_feature_override( + FeatureName.PROGRESSIVE_SSE_STREAMING, use_progressive_sse + ): + aggregator = streaming_utils.StreamingResponseAggregator() + # First chunk carries the token counts. + response1 = types.GenerateContentResponse( + candidates=[ + types.Candidate( + content=types.Content(parts=[types.Part(text="Hello ")]), + ) + ], + usage_metadata=types.GenerateContentResponseUsageMetadata( + prompt_token_count=10, + candidates_token_count=5, + total_token_count=15, + ), + ) + # Second chunk carries none. + response2 = types.GenerateContentResponse( + candidates=[ + types.Candidate( + content=types.Content(parts=[types.Part(text="World!")]), + 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 closed_response.usage_metadata is not None + assert closed_response.usage_metadata.prompt_token_count == 10 + assert closed_response.usage_metadata.candidates_token_count == 5 + assert closed_response.usage_metadata.total_token_count == 15 + + @pytest.mark.asyncio + @pytest.mark.parametrize("use_progressive_sse", [False, True]) + async def test_close_uses_latest_reported_usage_metadata( + self, use_progressive_sse + ): + """When several chunks report usage, the most recent one wins.""" + with temporary_feature_override( + FeatureName.PROGRESSIVE_SSE_STREAMING, use_progressive_sse + ): + aggregator = streaming_utils.StreamingResponseAggregator() + response1 = types.GenerateContentResponse( + candidates=[ + types.Candidate( + content=types.Content(parts=[types.Part(text="Hello ")]), + ) + ], + usage_metadata=types.GenerateContentResponseUsageMetadata( + prompt_token_count=10, + candidates_token_count=5, + total_token_count=15, + ), + ) + response2 = types.GenerateContentResponse( + candidates=[ + types.Candidate( + content=types.Content(parts=[types.Part(text="World!")]), + finish_reason=types.FinishReason.STOP, + ) + ], + usage_metadata=types.GenerateContentResponseUsageMetadata( + prompt_token_count=10, + candidates_token_count=9, + total_token_count=19, + ), + ) + + 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.usage_metadata is not None + assert closed_response.usage_metadata.total_token_count == 19 + @pytest.mark.asyncio @pytest.mark.parametrize("use_progressive_sse", [False, True]) async def test_close_propagates_model_version(self, use_progressive_sse):