diff --git a/py/src/braintrust/integrations/google_genai/cassettes/latest/test_url_context_metadata_async_stream.yaml b/py/src/braintrust/integrations/google_genai/cassettes/latest/test_url_context_metadata_async_stream.yaml new file mode 100644 index 000000000..9c9be9d4a --- /dev/null +++ b/py/src/braintrust/integrations/google_genai/cassettes/latest/test_url_context_metadata_async_stream.yaml @@ -0,0 +1,161 @@ +interactions: +- request: + body: '{"contents": [{"parts": [{"text": "Summarize the article at https://ai.google.dev/gemini-api/docs/url-context"}], + "role": "user"}], "tools": [{"urlContext": {}}], "generationConfig": {"maxOutputTokens": + 300}}' + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '208' + content-type: + - application/json + host: + - generativelanguage.googleapis.com + user-agent: + - google-genai-sdk/2.28.0 gl-python/3.14.7 + x-goog-api-client: + - google-genai-sdk/2.28.0 gl-python/3.14.7 + method: POST + uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse + response: + body: + string: "data: {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \" \\n\\nThe + URL\"}],\"role\": \"model\"},\"index\": 0,\"groundingMetadata\": {},\"urlContextMetadata\": + {\"urlMetadata\": [{\"retrievedUrl\": \"https://ai.google.dev/gemini-api/docs/url-context\",\"urlRetrievalStatus\": + \"URL_RETRIEVAL_STATUS_SUCCESS\"}]}}],\"usageMetadata\": {\"promptTokenCount\": + 24,\"candidatesTokenCount\": 27,\"totalTokenCount\": 4029,\"promptTokensDetails\": + [{\"modality\": \"TEXT\",\"tokenCount\": 24}],\"toolUsePromptTokenCount\": + 3942,\"toolUsePromptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": + 3942}],\"thoughtsTokenCount\": 36,\"serviceTier\": \"standard\"},\"modelVersion\": + \"gemini-2.5-flash\",\"responseId\": \"RMbHapKgA8SMjrEP_JjKSQ\"}\r\n\r\ndata: + {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \" context tool in + the Gemini API allows models to access and analyze content from specified + URLs to\"}],\"role\": \"model\"},\"index\": 0,\"groundingMetadata\": {}}],\"usageMetadata\": + {\"promptTokenCount\": 24,\"candidatesTokenCount\": 44,\"totalTokenCount\": + 4046,\"promptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": 24}],\"toolUsePromptTokenCount\": + 3942,\"toolUsePromptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": + 3942}],\"thoughtsTokenCount\": 36,\"serviceTier\": \"standard\"},\"modelVersion\": + \"gemini-2.5-flash\",\"responseId\": \"RMbHapKgA8SMjrEP_JjKSQ\"}\r\n\r\ndata: + {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \" enhance their responses. + This tool is useful for tasks such as extracting data, comparing documents, + synthesizing content, and analyzing code or documentation. It works by first + attempting to fetch content from an internal index cache\"}],\"role\": \"model\"},\"index\": + 0,\"groundingMetadata\": {}}],\"usageMetadata\": {\"promptTokenCount\": 24,\"candidatesTokenCount\": + 84,\"totalTokenCount\": 4086,\"promptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": + 24}],\"toolUsePromptTokenCount\": 3942,\"toolUsePromptTokensDetails\": [{\"modality\": + \"TEXT\",\"tokenCount\": 3942}],\"thoughtsTokenCount\": 36,\"serviceTier\": + \"standard\"},\"modelVersion\": \"gemini-2.5-flash\",\"responseId\": \"RMbHapKgA8SMjrEP_JjKSQ\"}\r\n\r\ndata: + {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \" and falling back + to a live fetch if the URL is not cached. The URL context tool can be combined + with other tools, like Google Search, for more powerful workflows. Responses + from the model using\"}],\"role\": \"model\"},\"index\": 0,\"groundingMetadata\": + {}}],\"usageMetadata\": {\"promptTokenCount\": 24,\"candidatesTokenCount\": + 123,\"totalTokenCount\": 4125,\"promptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": + 24}],\"toolUsePromptTokenCount\": 3942,\"toolUsePromptTokensDetails\": [{\"modality\": + \"TEXT\",\"tokenCount\": 3942}],\"thoughtsTokenCount\": 36,\"serviceTier\": + \"standard\"},\"modelVersion\": \"gemini-2.5-flash\",\"responseId\": \"RMbHapKgA8SMjrEP_JjKSQ\"}\r\n\r\ndata: + {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \" this tool include + inline URL citation annotations, linking parts of the text to their source + URLs. There are limitations, including a maximum of 20 URLs per request, a + 34MB content size limit per URL\"}],\"role\": \"model\"},\"index\": 0,\"groundingMetadata\": + {}}],\"usageMetadata\": {\"promptTokenCount\": 24,\"candidatesTokenCount\": + 166,\"totalTokenCount\": 4168,\"promptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": + 24}],\"toolUsePromptTokenCount\": 3942,\"toolUsePromptTokensDetails\": [{\"modality\": + \"TEXT\",\"tokenCount\": 3942}],\"thoughtsTokenCount\": 36,\"serviceTier\": + \"standard\"},\"modelVersion\": \"gemini-2.5-flash\",\"responseId\": \"RMbHapKgA8SMjrEP_JjKSQ\"}\r\n\r\ndata: + {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \", and the requirement + for publicly accessible URLs. It supports various content types such as text, + images, and PDFs, but does not support paywalled content, YouTube videos, + Google Workspace files, or other video/audio files\"}],\"role\": \"model\"},\"index\": + 0,\"groundingMetadata\": {}}],\"usageMetadata\": {\"promptTokenCount\": 24,\"candidatesTokenCount\": + 210,\"totalTokenCount\": 4212,\"promptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": + 24}],\"toolUsePromptTokenCount\": 3942,\"toolUsePromptTokensDetails\": [{\"modality\": + \"TEXT\",\"tokenCount\": 3942}],\"thoughtsTokenCount\": 36,\"serviceTier\": + \"standard\"},\"modelVersion\": \"gemini-2.5-flash\",\"responseId\": \"RMbHapKgA8SMjrEP_JjKSQ\"}\r\n\r\ndata: + {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \".\"}],\"role\": + \"model\"},\"index\": 0,\"groundingMetadata\": {}}],\"usageMetadata\": {\"promptTokenCount\": + 24,\"candidatesTokenCount\": 211,\"totalTokenCount\": 4213,\"promptTokensDetails\": + [{\"modality\": \"TEXT\",\"tokenCount\": 24}],\"toolUsePromptTokenCount\": + 3942,\"toolUsePromptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": + 3942}],\"thoughtsTokenCount\": 36,\"serviceTier\": \"standard\"},\"modelVersion\": + \"gemini-2.5-flash\",\"responseId\": \"RMbHapKgA8SMjrEP_JjKSQ\"}\r\n\r\ndata: + {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \"The URL context + tool within the Gemini API enables models to retrieve and analyze information + from provided URLs, improving the quality of their responses. This functionality + is beneficial for tasks like data extraction,\"}],\"role\": \"model\"},\"index\": + 0,\"groundingMetadata\": {}}],\"usageMetadata\": {\"promptTokenCount\": 24,\"candidatesTokenCount\": + 247,\"totalTokenCount\": 4249,\"promptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": + 24}],\"toolUsePromptTokenCount\": 3942,\"toolUsePromptTokensDetails\": [{\"modality\": + \"TEXT\",\"tokenCount\": 3942}],\"thoughtsTokenCount\": 36,\"serviceTier\": + \"standard\"},\"modelVersion\": \"gemini-2.5-flash\",\"responseId\": \"RMbHapKgA8SMjrEP_JjKSQ\"}\r\n\r\ndata: + {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \" document comparison, + content synthesis, and code analysis. The tool utilizes a two-step retrieval + process, checking an internal index cache before performing a live fetch\"}],\"role\": + \"model\"},\"finishReason\": \"MAX_TOKENS\",\"index\": 0,\"groundingMetadata\": + {\"groundingChunks\": [{\"web\": {\"uri\": \"https://ai.google.dev/gemini-api/docs/url-context\",\"title\": + \"URL context - Interactions API | Google AI for Developers\"}}],\"groundingSupports\": + [{\"segment\": {\"startIndex\": 3,\"endIndex\": 132,\"text\": \"The URL context + tool in the Gemini API allows models to access and analyze content from specified + URLs to enhance their responses\"},\"groundingChunkIndices\": [0]},{\"segment\": + {\"startIndex\": 134,\"endIndex\": 267,\"text\": \"This tool is useful for + tasks such as extracting data, comparing documents, synthesizing content, + and analyzing code or documentation\"},\"groundingChunkIndices\": [0]},{\"segment\": + {\"startIndex\": 269,\"endIndex\": 401,\"text\": \"It works by first attempting + to fetch content from an internal index cache and falling back to a live fetch + if the URL is not cached\"},\"groundingChunkIndices\": [0]},{\"segment\": + {\"startIndex\": 403,\"endIndex\": 505,\"text\": \"The URL context tool can + be combined with other tools, like Google Search, for more powerful workflows\"},\"groundingChunkIndices\": + [0]},{\"segment\": {\"startIndex\": 507,\"endIndex\": 635,\"text\": \"Responses + from the model using this tool include inline URL citation annotations, linking + parts of the text to their source URLs\"},\"groundingChunkIndices\": [0]},{\"segment\": + {\"startIndex\": 637,\"endIndex\": 787,\"text\": \"There are limitations, + including a maximum of 20 URLs per request, a 34MB content size limit per + URL, and the requirement for publicly accessible URLs\"},\"groundingChunkIndices\": + [0]},{\"segment\": {\"startIndex\": 789,\"endIndex\": 961,\"text\": \"It supports + various content types such as text, images, and PDFs, but does not support + paywalled content, YouTube videos, Google Workspace files, or other video/audio + files\"},\"groundingChunkIndices\": [0]},{\"segment\": {\"startIndex\": 962,\"endIndex\": + 1116,\"text\": \"The URL context tool within the Gemini API enables models + to retrieve and analyze information from provided URLs, improving the quality + of their responses\"},\"groundingChunkIndices\": [0]},{\"segment\": {\"startIndex\": + 1118,\"endIndex\": 1240,\"text\": \"This functionality is beneficial for tasks + like data extraction, document comparison, content synthesis, and code analysis\"},\"groundingChunkIndices\": + [0]}]}}],\"usageMetadata\": {\"promptTokenCount\": 24,\"candidatesTokenCount\": + 277,\"totalTokenCount\": 4279,\"promptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": + 24}],\"toolUsePromptTokenCount\": 3942,\"toolUsePromptTokensDetails\": [{\"modality\": + \"TEXT\",\"tokenCount\": 3942}],\"thoughtsTokenCount\": 36,\"serviceTier\": + \"standard\"},\"modelVersion\": \"gemini-2.5-flash\",\"responseId\": \"RMbHapKgA8SMjrEP_JjKSQ\"}\r\n\r\n" + headers: + alt-svc: + - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 + content-disposition: + - attachment + content-type: + - text/event-stream + date: + - Thu, 08 Oct 2026 16:35:17 GMT + server: + - scaffolding on HTTPServer2 + server-timing: + - gfet4t7; dur=1239 + transfer-encoding: + - chunked + vary: + - Origin + - X-Origin + - Referer + x-content-type-options: + - nosniff + x-frame-options: + - SAMEORIGIN + x-xss-protection: + - '0' + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/google_genai/cassettes/latest/test_url_context_metadata_stream.yaml b/py/src/braintrust/integrations/google_genai/cassettes/latest/test_url_context_metadata_stream.yaml new file mode 100644 index 000000000..8c0345267 --- /dev/null +++ b/py/src/braintrust/integrations/google_genai/cassettes/latest/test_url_context_metadata_stream.yaml @@ -0,0 +1,152 @@ +interactions: +- request: + body: '{"contents": [{"parts": [{"text": "Summarize the article at https://ai.google.dev/gemini-api/docs/url-context"}], + "role": "user"}], "tools": [{"urlContext": {}}], "generationConfig": {"maxOutputTokens": + 300}}' + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '208' + content-type: + - application/json + host: + - generativelanguage.googleapis.com + user-agent: + - google-genai-sdk/2.28.0 gl-python/3.14.7 + x-goog-api-client: + - google-genai-sdk/2.28.0 gl-python/3.14.7 + method: POST + uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse + response: + body: + string: "data: {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \"The + URL context tool in the Gemini API allows models to access and understand + content from provided\"}],\"role\": \"model\"},\"index\": 0,\"groundingMetadata\": + {},\"urlContextMetadata\": {\"urlMetadata\": [{\"retrievedUrl\": \"https://ai.google.dev/gemini-api/docs/url-context\",\"urlRetrievalStatus\": + \"URL_RETRIEVAL_STATUS_SUCCESS\"}]}}],\"usageMetadata\": {\"promptTokenCount\": + 24,\"candidatesTokenCount\": 42,\"totalTokenCount\": 4044,\"promptTokensDetails\": + [{\"modality\": \"TEXT\",\"tokenCount\": 24}],\"toolUsePromptTokenCount\": + 3942,\"toolUsePromptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": + 3942}],\"thoughtsTokenCount\": 36,\"serviceTier\": \"standard\"},\"modelVersion\": + \"gemini-2.5-flash\",\"responseId\": \"QcbHavmNCP-_39IPgoDIkA8\"}\r\n\r\ndata: + {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \" URLs to enhance + their responses. This tool is useful for tasks such as extracting data, comparing + documents, synthesizing content, and analyzing code or documentation. It operates + by first checking an internal cache for the URL content and,\"}],\"role\": + \"model\"},\"index\": 0,\"groundingMetadata\": {}}],\"usageMetadata\": {\"promptTokenCount\": + 24,\"candidatesTokenCount\": 85,\"totalTokenCount\": 4087,\"promptTokensDetails\": + [{\"modality\": \"TEXT\",\"tokenCount\": 24}],\"toolUsePromptTokenCount\": + 3942,\"toolUsePromptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": + 3942}],\"thoughtsTokenCount\": 36,\"serviceTier\": \"standard\"},\"modelVersion\": + \"gemini-2.5-flash\",\"responseId\": \"QcbHavmNCP-_39IPgoDIkA8\"}\r\n\r\ndata: + {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \" if not found, performing + a live fetch. The tool can be combined with other tools, including Google + Search, for more powerful workflows. Responses include `url_citation` annotations + to link text\"}],\"role\": \"model\"},\"index\": 0,\"groundingMetadata\": + {}}],\"usageMetadata\": {\"promptTokenCount\": 24,\"candidatesTokenCount\": + 123,\"totalTokenCount\": 4125,\"promptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": + 24}],\"toolUsePromptTokenCount\": 3942,\"toolUsePromptTokensDetails\": [{\"modality\": + \"TEXT\",\"tokenCount\": 3942}],\"thoughtsTokenCount\": 36,\"serviceTier\": + \"standard\"},\"modelVersion\": \"gemini-2.5-flash\",\"responseId\": \"QcbHavmNCP-_39IPgoDIkA8\"}\r\n\r\ndata: + {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \" to its source, + and safety checks are performed on URLs. There are limitations, such as a + maximum of 20 URLs per request, a 34MB content size limit per URL, and the + requirement for URLs\"}],\"role\": \"model\"},\"index\": 0,\"groundingMetadata\": + {}}],\"usageMetadata\": {\"promptTokenCount\": 24,\"candidatesTokenCount\": + 167,\"totalTokenCount\": 4169,\"promptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": + 24}],\"toolUsePromptTokenCount\": 3942,\"toolUsePromptTokensDetails\": [{\"modality\": + \"TEXT\",\"tokenCount\": 3942}],\"thoughtsTokenCount\": 36,\"serviceTier\": + \"standard\"},\"modelVersion\": \"gemini-2.5-flash\",\"responseId\": \"QcbHavmNCP-_39IPgoDIkA8\"}\r\n\r\ndata: + {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \" to be publicly + accessible. It supports various content types like text, image, and PDF, but + not paywalled content, YouTube videos, or Google Workspace files.\"}],\"role\": + \"model\"},\"index\": 0,\"groundingMetadata\": {}}],\"usageMetadata\": {\"promptTokenCount\": + 24,\"candidatesTokenCount\": 199,\"totalTokenCount\": 4201,\"promptTokensDetails\": + [{\"modality\": \"TEXT\",\"tokenCount\": 24}],\"toolUsePromptTokenCount\": + 3942,\"toolUsePromptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": + 3942}],\"thoughtsTokenCount\": 36,\"serviceTier\": \"standard\"},\"modelVersion\": + \"gemini-2.5-flash\",\"responseId\": \"QcbHavmNCP-_39IPgoDIkA8\"}\r\n\r\ndata: + {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \"The URL context + tool in the Gemini API\"}],\"role\": \"model\"},\"index\": 0,\"groundingMetadata\": + {}}],\"usageMetadata\": {\"promptTokenCount\": 24,\"candidatesTokenCount\": + 207,\"totalTokenCount\": 4209,\"promptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": + 24}],\"toolUsePromptTokenCount\": 3942,\"toolUsePromptTokensDetails\": [{\"modality\": + \"TEXT\",\"tokenCount\": 3942}],\"thoughtsTokenCount\": 36,\"serviceTier\": + \"standard\"},\"modelVersion\": \"gemini-2.5-flash\",\"responseId\": \"QcbHavmNCP-_39IPgoDIkA8\"}\r\n\r\ndata: + {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \" enables models + to use content from provided URLs to improve their responses. It helps with + tasks such as data extraction, document comparison, content synthesis, and + code analysis. The tool uses a two-step retrieval process, first checking + an internal\"}],\"role\": \"model\"},\"index\": 0,\"groundingMetadata\": {}}],\"usageMetadata\": + {\"promptTokenCount\": 24,\"candidatesTokenCount\": 254,\"totalTokenCount\": + 4256,\"promptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": 24}],\"toolUsePromptTokenCount\": + 3942,\"toolUsePromptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": + 3942}],\"thoughtsTokenCount\": 36,\"serviceTier\": \"standard\"},\"modelVersion\": + \"gemini-2.5-flash\",\"responseId\": \"QcbHavmNCP-_39IPgoDIkA8\"}\r\n\r\ndata: + {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \" cache and then + performing a live fetch if needed. It can be combined with other tools like + Google Search and includes `url_citation` annotations in responses\"}],\"role\": + \"model\"},\"finishReason\": \"MAX_TOKENS\",\"index\": 0,\"groundingMetadata\": + {\"groundingChunks\": [{\"web\": {\"uri\": \"https://ai.google.dev/gemini-api/docs/url-context\",\"title\": + \"URL context - Interactions API | Google AI for Developers\"}}],\"groundingSupports\": + [{\"segment\": {\"startIndex\": 134,\"endIndex\": 267,\"text\": \"This tool + is useful for tasks such as extracting data, comparing documents, synthesizing + content, and analyzing code or documentation\"},\"groundingChunkIndices\": + [0]},{\"segment\": {\"startIndex\": 269,\"endIndex\": 379,\"text\": \"It operates + by first checking an internal cache for the URL content and, if not found, + performing a live fetch\"},\"groundingChunkIndices\": [0]},{\"segment\": {\"startIndex\": + 381,\"endIndex\": 476,\"text\": \"The tool can be combined with other tools, + including Google Search, for more powerful workflows\"},\"groundingChunkIndices\": + [0]},{\"segment\": {\"startIndex\": 478,\"endIndex\": 590,\"text\": \"Responses + include `url_citation` annotations to link text to its source, and safety + checks are performed on URLs\"},\"groundingChunkIndices\": [0]},{\"segment\": + {\"startIndex\": 592,\"endIndex\": 746,\"text\": \"There are limitations, + such as a maximum of 20 URLs per request, a 34MB content size limit per URL, + and the requirement for URLs to be publicly accessible\"},\"groundingChunkIndices\": + [0]},{\"segment\": {\"startIndex\": 748,\"endIndex\": 877,\"text\": \"It supports + various content types like text, image, and PDF, but not paywalled content, + YouTube videos, or Google Workspace files\"},\"groundingChunkIndices\": [0]},{\"segment\": + {\"startIndex\": 994,\"endIndex\": 1096,\"text\": \"It helps with tasks such + as data extraction, document comparison, content synthesis, and code analysis\"},\"groundingChunkIndices\": + [0]},{\"segment\": {\"startIndex\": 1098,\"endIndex\": 1217,\"text\": \"The + tool uses a two-step retrieval process, first checking an internal cache and + then performing a live fetch if needed\"},\"groundingChunkIndices\": [0]}]}}],\"usageMetadata\": + {\"promptTokenCount\": 24,\"candidatesTokenCount\": 284,\"totalTokenCount\": + 4286,\"promptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": 24}],\"toolUsePromptTokenCount\": + 3942,\"toolUsePromptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": + 3942}],\"thoughtsTokenCount\": 36,\"serviceTier\": \"standard\"},\"modelVersion\": + \"gemini-2.5-flash\",\"responseId\": \"QcbHavmNCP-_39IPgoDIkA8\"}\r\n\r\n" + headers: + alt-svc: + - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 + content-disposition: + - attachment + content-type: + - text/event-stream + date: + - Thu, 08 Oct 2026 16:35:14 GMT + server: + - scaffolding on HTTPServer2 + server-timing: + - gfet4t7; dur=1557 + transfer-encoding: + - chunked + vary: + - Origin + - X-Origin + - Referer + x-content-type-options: + - nosniff + x-frame-options: + - SAMEORIGIN + x-xss-protection: + - '0' + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/google_genai/test_google_genai.py b/py/src/braintrust/integrations/google_genai/test_google_genai.py index a760c34c3..a40a76485 100644 --- a/py/src/braintrust/integrations/google_genai/test_google_genai.py +++ b/py/src/braintrust/integrations/google_genai/test_google_genai.py @@ -1278,6 +1278,66 @@ def _assert_grounding_metadata(span_output): ) +def _assert_url_context_metadata(span_output): + candidates = span_output.get("candidates", []) + assert candidates, "Expected candidates in span output" + + url_context = candidates[0].get("url_context_metadata") + assert url_context is not None, f"Expected url_context_metadata, got keys: {list(candidates[0])}" + url_metadata = url_context.get("url_metadata") + assert url_metadata, "Expected retrieved URL details in url_context_metadata" + assert any(item.get("retrieved_url") for item in url_metadata) + + +@LATEST_ONLY +@pytest.mark.vcr +def test_url_context_metadata_stream(memory_logger): + assert not memory_logger.pop() + + client = Client() + stream = client.models.generate_content_stream( + model=GROUNDING_MODEL, + contents="Summarize the article at https://ai.google.dev/gemini-api/docs/url-context", + config=types.GenerateContentConfig( + tools=[types.Tool(url_context=types.UrlContext())], + max_output_tokens=300, + ), + ) + text = "".join(chunk.text or "" for chunk in stream) + + assert text + spans = memory_logger.pop() + assert len(spans) == 1 + assert spans[0]["output"] + _assert_url_context_metadata(spans[0]["output"]) + + +@LATEST_ONLY +@pytest.mark.vcr +@pytest.mark.asyncio +async def test_url_context_metadata_async_stream(memory_logger): + assert not memory_logger.pop() + + client = Client() + stream = await client.aio.models.generate_content_stream( + model=GROUNDING_MODEL, + contents="Summarize the article at https://ai.google.dev/gemini-api/docs/url-context", + config=types.GenerateContentConfig( + tools=[types.Tool(url_context=types.UrlContext())], + max_output_tokens=300, + ), + ) + text = "" + async for chunk in stream: + text += chunk.text or "" + + assert text + spans = memory_logger.pop() + assert len(spans) == 1 + assert spans[0]["output"] + _assert_url_context_metadata(spans[0]["output"]) + + # Test: Google Search Grounding (Sync) @pytest.mark.vcr @pytest.mark.parametrize( @@ -1497,6 +1557,9 @@ def test_interactions_create_and_get(memory_logger): assert create_span["metadata"]["model"] == INTERACTIONS_MODEL assert create_span["metadata"]["provider"] == "google" assert create_span["metadata"]["interaction_id"] == response.id + assert create_span["metadata"]["status"] == response.status + assert "response_mime_type" not in create_span["metadata"] + assert "response_modalities" not in create_span["metadata"] assert get_span["metadata"]["provider"] == "google" assert create_span["output"]["status"] == "completed" assert "Paris" in create_span["output"]["text"] diff --git a/py/src/braintrust/integrations/google_genai/tracing.py b/py/src/braintrust/integrations/google_genai/tracing.py index c29fab40d..467c0fa41 100644 --- a/py/src/braintrust/integrations/google_genai/tracing.py +++ b/py/src/braintrust/integrations/google_genai/tracing.py @@ -549,8 +549,8 @@ def _extract_interaction_metadata(response: "Interaction") -> dict[str, Any]: "interaction_id": getattr(response, "id", None), "previous_interaction_id": getattr(response, "previous_interaction_id", None), "role": getattr(response, "role", None), - "response_mime_type": getattr(response, "response_mime_type", None), - "response_modalities": _materialize_interaction_value(getattr(response, "response_modalities", None)), + "status": getattr(response, "status", None), + "continuation_token": getattr(response, "continuation_token", None), "total_tool_use_tokens": ( usage_serialized.get("total_tool_use_tokens") if isinstance(usage_serialized, dict) else None ), @@ -659,6 +659,7 @@ def _aggregate_generate_content_chunks( other_parts = [] usage_metadata = None last_response = None + candidate_fields: list[dict[str, Any]] = [] for chunk in chunks: last_response = chunk @@ -667,7 +668,17 @@ def _aggregate_generate_content_chunks( usage_metadata = chunk.usage_metadata if hasattr(chunk, "candidates") and chunk.candidates: - for candidate in chunk.candidates: + for candidate_index, candidate in enumerate(chunk.candidates): + while len(candidate_fields) <= candidate_index: + candidate_fields.append({}) + serialized_candidate = _materialize_interaction_value( + candidate.model_dump(exclude_none=True, exclude={"content"}) + ) + if isinstance(serialized_candidate, dict): + candidate_fields[candidate_index].update( + {key: value for key, value in serialized_candidate.items() if value is not None} + ) + if hasattr(candidate, "content") and candidate.content: if hasattr(candidate.content, "parts") and candidate.content.parts: for part in candidate.content.parts: @@ -694,16 +705,9 @@ def _aggregate_generate_content_chunks( if parts and last_response and hasattr(last_response, "candidates"): candidates = [] - for candidate in last_response.candidates: - candidate_dict = {"content": {"parts": parts, "role": "model"}} - - if hasattr(candidate, "finish_reason"): - candidate_dict["finish_reason"] = candidate.finish_reason - if hasattr(candidate, "safety_ratings"): - candidate_dict["safety_ratings"] = candidate.safety_ratings - if hasattr(candidate, "grounding_metadata") and candidate.grounding_metadata: - candidate_dict["grounding_metadata"] = candidate.grounding_metadata - + for candidate_index, _candidate in enumerate(last_response.candidates): + candidate_dict = dict(candidate_fields[candidate_index]) + candidate_dict["content"] = {"parts": parts, "role": "model"} candidates.append(candidate_dict) aggregated["candidates"] = candidates