fix(google_genai): attach grounding_metadata to streaming responses#155
Open
Abhijeet Prasad (AbhiPrasad) wants to merge 1 commit intomainfrom
Open
fix(google_genai): attach grounding_metadata to streaming responses#155Abhijeet Prasad (AbhiPrasad) wants to merge 1 commit intomainfrom
Abhijeet Prasad (AbhiPrasad) wants to merge 1 commit intomainfrom
Conversation
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.
resolves #153
AI Summary
Issue
Google Search grounding metadata (
grounding_chunks,grounding_supports,search_entry_point,web_search_queries) was silently dropped from span output when using streaming responses.Root Cause
In
py/src/braintrust/integrations/google_genai/tracing.py, the_aggregate_generate_content_chunks()function manually builds the aggregated output dict for streaming responses. It copiedfinish_reasonandsafety_ratingsfrom the last candidate, but notgrounding_metadata.For non-streaming responses, the raw
GenerateContentResponsePydantic model was returned directly and serialized viamodel_dump(exclude_none=True), which naturally includedgrounding_metadata. So only the streaming path was broken.Fix
Added extraction of
grounding_metadatafrom the last chunk's candidate in_aggregate_generate_content_chunks(), serializing it viamodel_dump(exclude_none=True)to match the non-streaming output format:Tests
Added 4 VCR-backed tests covering all transport modes with real Google API cassettes:
test_google_search_grounding[sync]— non-streaming synctest_google_search_grounding[stream]— streaming synctest_google_search_grounding_async[async]— non-streaming asynctest_google_search_grounding_async[async_stream]— streaming asyncEach test verifies that
grounding_metadataon the first candidate containsweb_search_queries,grounding_chunks, andgrounding_supports.Red/Green Evidence
streamandasync_streamtests failed withAssertionError: Expected grounding_metadata on first candidate, got keys: ['content', 'finish_reason', 'safety_ratings']