diff --git a/CHANGELOG.md b/CHANGELOG.md index 98a277d7..15637e7c 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/chatlas/_provider_google.py b/chatlas/_provider_google.py index 0984a934..799b7dc9 100644 --- a/chatlas/_provider_google.py +++ b/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() diff --git a/tests/test_provider_google.py b/tests/test_provider_google.py index c65494d1..e7107710 100644 --- a/tests/test_provider_google.py +++ b/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.