fix(google): keep streamed citations adjacent to text

CHANGELOG.md

@@ -12,6 +12,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Bug fixes
 
+* `ChatGoogle()` and `ChatVertex()` now emit grounded citations immediately after
+  the answer text they support while streaming, before the related web search or
+  URL-fetch activity. This keeps streamed citations attached to their answer
+  text instead of separating them into a later activity segment.
 * OpenAI-based providers now support the OpenAI 3 SDK and its native `httpx2`
   clients. Custom clients are typed and documented as `httpx2`; legacy `httpx`
   clients remain supported at runtime during migration. (#387)

chatlas/_provider_google.py

@@ -493,19 +493,33 @@ def stream_content(self, chunk, completion, turns=()) -> list[Content]:
                     else:
                         part_contents.append(ContentText.model_construct(text=text))
 
-        grounding_metadata = getattr(candidate, "grounding_metadata", None)
-        url_context_metadata = getattr(candidate, "url_context_metadata", None)
-
-        grounding_contents: list[Content] = []
-        if grounding_metadata is not None:
-            gm_dict = grounding_metadata.model_dump()
-            grounding_contents.extend(google_search_contents(gm_dict))
-            grounding_contents.extend(google_grounding_citations(gm_dict))
-        if url_context_metadata is not None:
-            uc_dict = url_context_metadata.model_dump()
-            grounding_contents.extend(google_url_context_contents(uc_dict))
-
-        return part_contents + grounding_contents
+        if not any(
+            getattr(candidate, "finish_reason", None) is not None
+            for candidate in candidates
+        ):
+            return part_contents
+        if completion is None:
+            return part_contents
+
+        activity_contents: list[Content] = []
+        citation_contents: list[Content] = []
+
+        # URL context can arrive before later answer text. Rebuild annotations
+        # from the merged completion so they cannot split a grounded span.
+        for merged_candidate in completion.get("candidates") or []:
+            grounding_metadata = merged_candidate.get("grounding_metadata")
+            if grounding_metadata:
+                citation_contents.extend(
+                    google_grounding_citations(grounding_metadata)
+                )
+                activity_contents.extend(google_search_contents(grounding_metadata))
+            url_context_metadata = merged_candidate.get("url_context_metadata")
+            if url_context_metadata:
+                activity_contents.extend(
+                    google_url_context_contents(url_context_metadata)
+                )
+
+        return part_contents + citation_contents + activity_contents
 
     def stream_merge_chunks(self, completion, chunk):
         chunkd = chunk.model_dump()

tests/test_provider_google.py

@@ -14,13 +14,22 @@
     Content,
     ContentCitation,
     ContentText,
+    ContentToolRequestFetch,
     ContentToolRequestSearch,
     ContentToolResponseFetch,
     ContentToolResponseSearch,
     WebSource,
 )
 from google.genai.errors import APIError
-from google.genai.types import GroundingMetadataDict
+from google.genai.types import (
+    FinishReason as GoogleFinishReason,
+)
+from google.genai.types import (
+    GroundingMetadataDict,
+    UrlContextMetadata,
+    UrlMetadata,
+    UrlRetrievalStatus,
+)
 from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential
 
 from .conftest import (
@@ -348,6 +357,20 @@ def _plain_chunk(text: str, index: int | None = 0):
     )
 
 
+def _url_context_chunk(text: str, url: str, index: int | None = 0):
+    chunk = _plain_chunk(text, index)
+    assert chunk.candidates
+    chunk.candidates[0].url_context_metadata = UrlContextMetadata(
+        url_metadata=[
+            UrlMetadata(
+                retrieved_url=url,
+                url_retrieval_status=UrlRetrievalStatus.URL_RETRIEVAL_STATUS_SUCCESS,
+            )
+        ]
+    )
+    return chunk
+
+
 def test_google_grounding_metadata_matches_streamed_content():
     """Streaming and final-turn citations agree for Gemini's real response shape.
 
@@ -363,6 +386,8 @@ def test_google_grounding_metadata_matches_streamed_content():
         _plain_chunk("ggplot2 1.0.0 "),
         _grounding_chunk("was released on 2014-05-21.", "https://a.com", "A", "q1"),
     ]
+    assert chunks[-1].candidates
+    chunks[-1].candidates[0].finish_reason = GoogleFinishReason.STOP
 
     streamed: list[Content] = []
     completion = None
@@ -387,6 +412,69 @@ def test_google_grounding_metadata_matches_streamed_content():
     ]
 
 
+def test_google_stream_keeps_citations_next_to_text():
+    provider = GoogleProvider(
+        model="gemini-2.5-flash",
+        api_key="dummy",
+        name="Google/Gemini",
+        kwargs=None,
+    )
+    chunk = _grounding_chunk(
+        "A grounded answer.",
+        "https://a.com",
+        "A",
+        "grounded answer source",
+    )
+    assert chunk.candidates
+    chunk.candidates[0].finish_reason = GoogleFinishReason.STOP
+
+    completion = provider.stream_merge_chunks(None, chunk)
+    contents = provider.stream_content(chunk, completion)
+
+    assert [type(content) for content in contents] == [
+        ContentText,
+        ContentCitation,
+        ContentToolRequestSearch,
+        ContentToolResponseSearch,
+    ]
+
+
+def test_google_stream_defers_early_fetch_until_citations():
+    provider = GoogleProvider(
+        model="gemini-2.5-flash",
+        api_key="dummy",
+        name="Google/Gemini",
+        kwargs=None,
+    )
+    chunks = [
+        _url_context_chunk("A grounded answer was ", "https://a.com"),
+        _grounding_chunk(
+            "released today.",
+            "https://a.com",
+            "A",
+            "grounded answer source",
+        ),
+    ]
+    assert chunks[-1].candidates
+    chunks[-1].candidates[0].finish_reason = GoogleFinishReason.STOP
+
+    contents: list[Content] = []
+    completion = None
+    for chunk in chunks:
+        completion = provider.stream_merge_chunks(completion, chunk)
+        contents.extend(provider.stream_content(chunk, completion))
+
+    assert [type(content) for content in contents] == [
+        ContentText,
+        ContentText,
+        ContentCitation,
+        ContentToolRequestSearch,
+        ContentToolResponseSearch,
+        ContentToolRequestFetch,
+        ContentToolResponseFetch,
+    ]
+
+
 def test_google_late_grounding_metadata_not_dropped():
     """Metadata on a later candidate must survive into the final turn.