fix: propagate model_version and other metadata in streaming responses

google-genai-bot · copybara-github · commit 342b59d55c9b · 2026-06-08T18:39:22.000-07:00
PiperOrigin-RevId: 928881287
diff --git a/src/google/adk/utils/streaming_utils.py b/src/google/adk/utils/streaming_utils.py
@@ -336,6 +336,10 @@ async def process_response(
       yield LlmResponse(
           content=types.ModelContent(parts=parts),
           usage_metadata=llm_response.usage_metadata,
+          grounding_metadata=llm_response.grounding_metadata,
+          citation_metadata=llm_response.citation_metadata,
+          finish_reason=llm_response.finish_reason,
+          model_version=llm_response.model_version,
       )
       self._thought_text = ''
       self._text = ''
@@ -386,6 +390,7 @@ def close(self) -> Optional[LlmResponse]:
           usage_metadata=self._usage_metadata,
           finish_reason=finish_reason,
           partial=False,
+          model_version=self._response.model_version,
       )
 
     # ========== Non-Progressive SSE Streaming (old behavior) ==========
@@ -405,4 +410,5 @@ def close(self) -> Optional[LlmResponse]:
         usage_metadata=self._usage_metadata,
         finish_reason=finish_reason,
         partial=False,
+        model_version=self._response.model_version,
     )
diff --git a/tests/unittests/utils/test_streaming_utils.py b/tests/unittests/utils/test_streaming_utils.py
@@ -386,6 +386,103 @@ async def 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.
