From 048a0202126b9984e9cb513ff4930b9160ce0122 Mon Sep 17 00:00:00 2001 From: Carson Date: Tue, 11 Aug 2026 21:11:04 -0500 Subject: [PATCH 1/5] fix(google): keep streamed citations next to answer text --- chatlas/_provider_google.py | 12 +++++++----- tests/test_provider_google.py | 24 ++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 5 deletions(-) diff --git a/chatlas/_provider_google.py b/chatlas/_provider_google.py index 0984a934..ccc1009d 100644 --- a/chatlas/_provider_google.py +++ b/chatlas/_provider_google.py @@ -496,16 +496,18 @@ def stream_content(self, chunk, completion, turns=()) -> list[Content]: grounding_metadata = getattr(candidate, "grounding_metadata", None) url_context_metadata = getattr(candidate, "url_context_metadata", None) - grounding_contents: list[Content] = [] + activity_contents: list[Content] = [] + citation_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)) + citation_contents.extend(google_grounding_citations(gm_dict)) + activity_contents.extend(google_search_contents(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)) + activity_contents.extend(google_url_context_contents(uc_dict)) - return part_contents + grounding_contents + 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..00f2f8d7 100644 --- a/tests/test_provider_google.py +++ b/tests/test_provider_google.py @@ -387,6 +387,30 @@ 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", + ) + + contents = provider.stream_content(chunk, completion=None) + + assert [type(content) for content in contents] == [ + ContentText, + ContentCitation, + ContentToolRequestSearch, + ContentToolResponseSearch, + ] + + def test_google_late_grounding_metadata_not_dropped(): """Metadata on a later candidate must survive into the final turn. From 6b2becf11dbac482b944b6c2fa0ae28ade70dc96 Mon Sep 17 00:00:00 2001 From: Carson Date: Tue, 11 Aug 2026 22:13:59 -0500 Subject: [PATCH 2/5] fix(google): defer streamed web activity --- chatlas/_provider_google.py | 30 +++++++++++----- tests/test_provider_google.py | 68 +++++++++++++++++++++++++++++++++-- 2 files changed, 87 insertions(+), 11 deletions(-) diff --git a/chatlas/_provider_google.py b/chatlas/_provider_google.py index ccc1009d..799b7dc9 100644 --- a/chatlas/_provider_google.py +++ b/chatlas/_provider_google.py @@ -493,19 +493,31 @@ 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) + 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] = [] - if grounding_metadata is not None: - gm_dict = grounding_metadata.model_dump() - citation_contents.extend(google_grounding_citations(gm_dict)) - activity_contents.extend(google_search_contents(gm_dict)) - if url_context_metadata is not None: - uc_dict = url_context_metadata.model_dump() - activity_contents.extend(google_url_context_contents(uc_dict)) + # 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 diff --git a/tests/test_provider_google.py b/tests/test_provider_google.py index 00f2f8d7..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 @@ -400,14 +425,53 @@ def test_google_stream_keeps_citations_next_to_text(): "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 = provider.stream_content(chunk, completion=None) + 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, ] From f3041c584053c256df89ba13c03185f7188435fa Mon Sep 17 00:00:00 2001 From: Carson Date: Tue, 11 Aug 2026 22:35:43 -0500 Subject: [PATCH 3/5] fix(deps): cap OpenAI before httpx2 migration --- pyproject.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 7c2d609f..92f082d8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,8 @@ dependencies = [ "jinja2", "orjson", "rich", - "openai", + # OpenAI 3 uses httpx2, while provider adapters still accept httpx clients. + "openai<3", "opentelemetry-api>=1.0", # `_typing_extensions.py` backfills from here on older interpreters; 3.13+ has # everything in `typing`. The floor is the release that added `TypeIs`. From 7ab454f259a1c30ad40ed9c5da562c2d42b6ded2 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 12 Aug 2026 17:29:22 +0000 Subject: [PATCH 4/5] chore: update provider types and pricing data --- chatlas/types/openai/_client.py | 8 ++++---- chatlas/types/openai/_client_azure.py | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/chatlas/types/openai/_client.py b/chatlas/types/openai/_client.py index cb985b69..507938e8 100644 --- a/chatlas/types/openai/_client.py +++ b/chatlas/types/openai/_client.py @@ -5,7 +5,7 @@ from typing import Awaitable, Callable, Mapping, Optional, TypedDict, Union -import httpx2 +import httpx import openai import openai._provider import openai.auth._workload @@ -19,12 +19,12 @@ class ChatClientArgs(TypedDict, total=False): project: str | None webhook_secret: str | None provider: openai._provider._Provider | None - base_url: str | httpx2.URL | None - websocket_base_url: str | httpx2.URL | None + base_url: str | httpx.URL | None + websocket_base_url: str | httpx.URL | None timeout: float | openai.Timeout | None | openai.NotGiven max_retries: int default_headers: Optional[Mapping[str, str]] default_query: Optional[Mapping[str, object]] - http_client: httpx2.Client | httpx2.AsyncClient | None + http_client: httpx.AsyncClient | None _strict_response_validation: bool _enforce_credentials: bool diff --git a/chatlas/types/openai/_client_azure.py b/chatlas/types/openai/_client_azure.py index e8e31cd6..d691d1bb 100644 --- a/chatlas/types/openai/_client_azure.py +++ b/chatlas/types/openai/_client_azure.py @@ -4,7 +4,7 @@ from typing import Awaitable, Callable, Mapping, Optional, TypedDict, Union -import httpx2 +import httpx import openai import openai.auth._workload @@ -21,11 +21,11 @@ class ChatAzureClientArgs(TypedDict, total=False): project: str | None webhook_secret: str | None base_url: str | None - websocket_base_url: str | httpx2.URL | None + websocket_base_url: str | httpx.URL | None timeout: float | openai.Timeout | None | openai.NotGiven max_retries: int default_headers: Optional[Mapping[str, str]] default_query: Optional[Mapping[str, object]] - http_client: httpx2.Client | httpx2.AsyncClient | None + http_client: httpx.AsyncClient | None _strict_response_validation: bool _enforce_credentials: bool From b06d710073a0124b48db415b6dc575bbfd3550df Mon Sep 17 00:00:00 2001 From: Carson Date: Thu, 13 Aug 2026 10:42:07 -0500 Subject: [PATCH 5/5] fix(deps): restore OpenAI 3 support --- CHANGELOG.md | 4 ++++ chatlas/types/openai/_client.py | 8 ++++---- chatlas/types/openai/_client_azure.py | 6 +++--- pyproject.toml | 3 +-- 4 files changed, 12 insertions(+), 9 deletions(-) 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/types/openai/_client.py b/chatlas/types/openai/_client.py index 507938e8..cb985b69 100644 --- a/chatlas/types/openai/_client.py +++ b/chatlas/types/openai/_client.py @@ -5,7 +5,7 @@ from typing import Awaitable, Callable, Mapping, Optional, TypedDict, Union -import httpx +import httpx2 import openai import openai._provider import openai.auth._workload @@ -19,12 +19,12 @@ class ChatClientArgs(TypedDict, total=False): project: str | None webhook_secret: str | None provider: openai._provider._Provider | None - base_url: str | httpx.URL | None - websocket_base_url: str | httpx.URL | None + base_url: str | httpx2.URL | None + websocket_base_url: str | httpx2.URL | None timeout: float | openai.Timeout | None | openai.NotGiven max_retries: int default_headers: Optional[Mapping[str, str]] default_query: Optional[Mapping[str, object]] - http_client: httpx.AsyncClient | None + http_client: httpx2.Client | httpx2.AsyncClient | None _strict_response_validation: bool _enforce_credentials: bool diff --git a/chatlas/types/openai/_client_azure.py b/chatlas/types/openai/_client_azure.py index d691d1bb..e8e31cd6 100644 --- a/chatlas/types/openai/_client_azure.py +++ b/chatlas/types/openai/_client_azure.py @@ -4,7 +4,7 @@ from typing import Awaitable, Callable, Mapping, Optional, TypedDict, Union -import httpx +import httpx2 import openai import openai.auth._workload @@ -21,11 +21,11 @@ class ChatAzureClientArgs(TypedDict, total=False): project: str | None webhook_secret: str | None base_url: str | None - websocket_base_url: str | httpx.URL | None + websocket_base_url: str | httpx2.URL | None timeout: float | openai.Timeout | None | openai.NotGiven max_retries: int default_headers: Optional[Mapping[str, str]] default_query: Optional[Mapping[str, object]] - http_client: httpx.AsyncClient | None + http_client: httpx2.Client | httpx2.AsyncClient | None _strict_response_validation: bool _enforce_credentials: bool diff --git a/pyproject.toml b/pyproject.toml index 92f082d8..7c2d609f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,8 +14,7 @@ dependencies = [ "jinja2", "orjson", "rich", - # OpenAI 3 uses httpx2, while provider adapters still accept httpx clients. - "openai<3", + "openai", "opentelemetry-api>=1.0", # `_typing_extensions.py` backfills from here on older interpreters; 3.13+ has # everything in `typing`. The floor is the release that added `TypeIs`.