Skip to content

fix(google-genai): retain URL context metadata in streams - #840

Draft
Kayvan Zahiri (Kayvan-Zahiri) wants to merge 1 commit into
braintrustdata:mainfrom
Kayvan-Zahiri:codex/fix-genai-url-context
Draft

Kayvan Zahiri (Kayvan-Zahiri) wants to merge 1 commit into
braintrustdata:mainfrom
Kayvan-Zahiri:codex/fix-genai-url-context

Conversation

@Kayvan-Zahiri

Copy link
Copy Markdown

Streaming generate_content_stream() and agenerate_content_stream() reconstruct candidate output but omit url_context_metadata, losing the retrieved URLs and retrieval statuses that non-streaming responses retain. Preserve that field alongside grounding_metadata in the shared stream aggregator.

Fixes #774.

The regression uses real google.genai.types.GenerateContentResponse objects with two chunks, successful and failed URL retrievals, and a no-metadata case. Before the fix, the metadata-present case fails with KeyError: 'url_context_metadata'; the no-metadata case passes. The payload follows Google's documented response structure. This is synthetic typed SDK payload coverage: existing cassettes contain no URL-context response, and provider credentials were unavailable to record a new one.

Validation on Python 3.13.4, using all three Google SDK pins from py/pyproject.toml and replay-only existing cassettes (CI=true, with the corresponding BRAINTRUST_TEST_PACKAGE_VERSION):

  • Full test_google_genai.py suite on 2.25.0: 43 passed, 2 skipped.
  • Full test_google_genai.py suite on 1.75.0: 41 passed, 4 skipped.
  • Full test_google_genai.py suite on 1.30.0: 35 passed, 10 skipped.
  • All repository pre-commit hooks passed; targeted pylint --errors-only passed for both changed files.

Tests and dependency installs used an isolated local virtual environment. make fixup could not find mise on this machine, so its pre-commit run --all-files command was run directly with that environment. The older-version subprocess tests required an explicit source PYTHONPATH because the local editable install was not resolved there. No dependency, cassette, or runtime environment changes are included in the PR.



@pytest.mark.parametrize("include_url_metadata", [True, False])
def test_stream_preserves_url_context_metadata(include_url_metadata):

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

generally we prefer tests that don't use mocks/fakes but test directly against the llm call itself, via https://github.com/braintrustdata/braintrust-sdk-python/blob/main/docs/vcr-testing.md. Are you able to record cassettes?

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[bot] Google GenAI: streaming responses drop url_context_metadata that non-streaming responses preserve

2 participants