fix(google-genai): retain URL context metadata in streams - #840
Draft
Kayvan Zahiri (Kayvan-Zahiri) wants to merge 1 commit into
Draft
Kayvan Zahiri (Kayvan-Zahiri) wants to merge 1 commit into
Kayvan Zahiri (Kayvan-Zahiri) wants to merge 1 commit into
Conversation
|
|
||
|
|
||
| @pytest.mark.parametrize("include_url_metadata", [True, False]) | ||
| def test_stream_preserves_url_context_metadata(include_url_metadata): |
There was a problem hiding this comment.
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Streaming
generate_content_stream()andagenerate_content_stream()reconstruct candidate output but omiturl_context_metadata, losing the retrieved URLs and retrieval statuses that non-streaming responses retain. Preserve that field alongsidegrounding_metadatain the shared stream aggregator.Fixes #774.
The regression uses real
google.genai.types.GenerateContentResponseobjects with two chunks, successful and failed URL retrievals, and a no-metadata case. Before the fix, the metadata-present case fails withKeyError: '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.tomland replay-only existing cassettes (CI=true, with the correspondingBRAINTRUST_TEST_PACKAGE_VERSION):test_google_genai.pysuite on 2.25.0: 43 passed, 2 skipped.test_google_genai.pysuite on 1.75.0: 41 passed, 4 skipped.test_google_genai.pysuite on 1.30.0: 35 passed, 10 skipped.pylint --errors-onlypassed for both changed files.Tests and dependency installs used an isolated local virtual environment.
make fixupcould not findmiseon this machine, so itspre-commit run --all-filescommand was run directly with that environment. The older-version subprocess tests required an explicit sourcePYTHONPATHbecause the local editable install was not resolved there. No dependency, cassette, or runtime environment changes are included in the PR.