From d54baf9554f33ce911878ceeb5eee5e4e02bb03a Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Wed, 29 Jul 2026 15:36:59 -0400 Subject: [PATCH] fix(google_genai): align usage metadata metrics Include tool-use prompts and reasoning in normalized token totals, map audio and image modality usage, and preserve provider-specific cache and tool breakdowns.\n\nAdd sanitized VCR coverage for audio input/output, generated images, cached content, and streaming usage metadata. --- .agents/skills/sdk-integrations/SKILL.md | 4 +- .../latest/test_audio_input_usage.yaml | 66 +++++ .../test_cached_content_usage_metadata.yaml | 277 ++++++++++++++++++ .../latest/test_generated_audio_usage.yaml | 65 ++++ .../latest/test_generated_image_usage.yaml | 64 ++++ .../google_genai/test_google_genai.py | 232 +++++++++++++-- .../integrations/google_genai/tracing.py | 110 +++++-- 7 files changed, 757 insertions(+), 61 deletions(-) create mode 100644 py/src/braintrust/integrations/google_genai/cassettes/latest/test_audio_input_usage.yaml create mode 100644 py/src/braintrust/integrations/google_genai/cassettes/latest/test_cached_content_usage_metadata.yaml create mode 100644 py/src/braintrust/integrations/google_genai/cassettes/latest/test_generated_audio_usage.yaml create mode 100644 py/src/braintrust/integrations/google_genai/cassettes/latest/test_generated_image_usage.yaml diff --git a/.agents/skills/sdk-integrations/SKILL.md b/.agents/skills/sdk-integrations/SKILL.md index 793750ce..36d9cd82 100644 --- a/.agents/skills/sdk-integrations/SKILL.md +++ b/.agents/skills/sdk-integrations/SKILL.md @@ -202,6 +202,8 @@ Only emit spec-listed keys: - `prompt_audio_tokens`, `completion_audio_tokens`, `completion_image_tokens` — audio/image models when provider reports - `start`, `end` — standard span timing +Map provider modality breakdowns only to standardized metrics: input audio → `prompt_audio_tokens`, output audio → `completion_audio_tokens`, and output image → `completion_image_tokens`. These are subsets of prompt/completion usage, not additional tokens. Preserve unstandardized modality breakdowns as deliberately named provider metadata rather than inventing metric keys. + For streaming, still produce one span per API call with accumulated `input`/`output`. Capture usage from stream metadata (e.g. OpenAI `stream_options.include_usage`) when surfaced. For reasoning models, capture the full output structure (reasoning summaries + message blocks) and include prior reasoning in `input` for multi-turn calls. @@ -261,7 +263,7 @@ Assert on emitted spans (not just provider return values): For streaming, assert both the provider iterator/async-iterator still works AND the final span has aggregated `output` + stream-specific `metrics`. -Cassettes live in `integrations//cassettes//` (e.g. `cassettes/latest/`, `cassettes/0.48.0/`). Nox sets `BRAINTRUST_TEST_PACKAGE_VERSION` so cassettes land correctly. Do not add per-test `vcr_cassette_dir` / `cassette_library_dir` fixtures — `integrations/conftest.py` handles it. Re-record only when behavior intentionally changed. Sanitize cassettes when the provider returns binary bodies. +Cassettes live in `integrations//cassettes//` (e.g. `cassettes/latest/`, `cassettes/0.48.0/`). Nox sets `BRAINTRUST_TEST_PACKAGE_VERSION` so cassettes land correctly. Do not add per-test `vcr_cassette_dir` / `cassette_library_dir` fixtures — `integrations/conftest.py` handles it. Re-record only when behavior intentionally changed. Sanitize binary media in both request and response bodies so checked-in cassettes do not retain large base64 payloads. Confirm the exact session name from `noxfile.py` — don't assume it matches the folder. diff --git a/py/src/braintrust/integrations/google_genai/cassettes/latest/test_audio_input_usage.yaml b/py/src/braintrust/integrations/google_genai/cassettes/latest/test_audio_input_usage.yaml new file mode 100644 index 00000000..093afb87 --- /dev/null +++ b/py/src/braintrust/integrations/google_genai/cassettes/latest/test_audio_input_usage.yaml @@ -0,0 +1,66 @@ +interactions: +- request: + body: '{"contents": [{"parts": [{"inlineData": {"data": "UklGRiQAAABXQVZFZm10IBAAAAABAAEAgD4AAAB9AAACABAAZGF0YQAAAAA=", + "mimeType": "audio/wav"}}, {"text": "Describe this audio in five words or less."}], + "role": "user"}], "generationConfig": {"maxOutputTokens": 100}}' + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '42928' + content-type: + - application/json + host: + - generativelanguage.googleapis.com + user-agent: + - google-genai-sdk/2.14.0 gl-python/3.14.6 + x-goog-api-client: + - google-genai-sdk/2.14.0 gl-python/3.14.6 + method: POST + uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent + response: + body: + string: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": + [\n {\n \"text\": \"Relaxing acoustic guitar music.\"\n + \ }\n ],\n \"role\": \"model\"\n },\n \"finishReason\": + \"STOP\",\n \"index\": 0\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": + 43,\n \"candidatesTokenCount\": 6,\n \"totalTokenCount\": 49,\n \"promptTokensDetails\": + [\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 10\n + \ },\n {\n \"modality\": \"AUDIO\",\n \"tokenCount\": + 33\n }\n ],\n \"serviceTier\": \"standard\"\n },\n \"modelVersion\": + \"gemini-2.5-flash-lite\",\n \"responseId\": \"N1NqapHyDayQ-8YPkLLHgAo\"\n}\n" + headers: + alt-svc: + - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 + content-length: + - '649' + content-type: + - application/json; charset=UTF-8 + date: + - Wed, 29 Jul 2026 19:23:36 GMT + server: + - scaffolding on HTTPServer2 + server-timing: + - gfet4t7; dur=1155 + transfer-encoding: + - chunked + vary: + - Origin + - X-Origin + - Referer + x-content-type-options: + - nosniff + x-frame-options: + - SAMEORIGIN + x-gemini-service-tier: + - standard + x-xss-protection: + - '0' + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/google_genai/cassettes/latest/test_cached_content_usage_metadata.yaml b/py/src/braintrust/integrations/google_genai/cassettes/latest/test_cached_content_usage_metadata.yaml new file mode 100644 index 00000000..c64d9ddf --- /dev/null +++ b/py/src/braintrust/integrations/google_genai/cassettes/latest/test_cached_content_usage_metadata.yaml @@ -0,0 +1,277 @@ +interactions: +- request: + body: '{"model": "models/gemini-2.5-flash", "ttl": "300s", "displayName": "braintrust-google-genai-vcr-test", + "contents": [{"parts": [{"text": "The sky is blue because molecules scatter + blue light. The sky is blue because molecules scatter blue light. The sky is + blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. The sky is blue because molecules + scatter blue light. The sky is blue because molecules scatter blue light. The + sky is blue because molecules scatter blue light. "}], "role": "user"}]}' + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '8259' + content-type: + - application/json + host: + - generativelanguage.googleapis.com + user-agent: + - google-genai-sdk/2.14.0 gl-python/3.14.6 + x-goog-api-client: + - google-genai-sdk/2.14.0 gl-python/3.14.6 + method: POST + uri: https://generativelanguage.googleapis.com/v1beta/cachedContents + response: + body: + string: "{\n \"name\": \"cachedContents/aixcv88yts91zirycnmno31aitdexjt7uc6b19mw\",\n + \ \"model\": \"models/gemini-2.5-flash\",\n \"createTime\": \"2026-07-29T19:15:22.193758Z\",\n + \ \"updateTime\": \"2026-07-29T19:15:22.193758Z\",\n \"expireTime\": \"2026-07-29T19:20:21.884570801Z\",\n + \ \"displayName\": \"braintrust-google-genai-vcr-test\",\n \"usageMetadata\": + {\n \"totalTokenCount\": 1502\n }\n}\n" + headers: + alt-svc: + - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 + content-length: + - '361' + content-type: + - application/json; charset=UTF-8 + date: + - Wed, 29 Jul 2026 19:15:22 GMT + server: + - scaffolding on HTTPServer2 + server-timing: + - gfet4t7; dur=591 + 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 +- request: + body: '{"contents": [{"parts": [{"text": "In one sentence, why is the sky blue?"}], + "role": "user"}], "cachedContent": "cachedContents/aixcv88yts91zirycnmno31aitdexjt7uc6b19mw", + "generationConfig": {"maxOutputTokens": 100}}' + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '216' + content-type: + - application/json + host: + - generativelanguage.googleapis.com + user-agent: + - google-genai-sdk/2.14.0 gl-python/3.14.6 + x-goog-api-client: + - google-genai-sdk/2.14.0 gl-python/3.14.6 + method: POST + uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent + response: + body: + string: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": + [\n {\n \"text\": \"The sky is blue because molecules + scatter blue light.\"\n }\n ],\n \"role\": \"model\"\n + \ },\n \"finishReason\": \"STOP\",\n \"index\": 0\n }\n ],\n + \ \"usageMetadata\": {\n \"promptTokenCount\": 1513,\n \"candidatesTokenCount\": + 10,\n \"totalTokenCount\": 1595,\n \"cachedContentTokenCount\": 1502,\n + \ \"promptTokensDetails\": [\n {\n \"modality\": \"TEXT\",\n + \ \"tokenCount\": 1513\n }\n ],\n \"cacheTokensDetails\": + [\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 1502\n + \ }\n ],\n \"thoughtsTokenCount\": 72,\n \"serviceTier\": \"standard\"\n + \ },\n \"modelVersion\": \"gemini-2.5-flash\",\n \"responseId\": \"SlFqaqzJGYGr1MkP4cu4yQ0\"\n}\n" + headers: + alt-svc: + - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 + content-length: + - '775' + content-type: + - application/json; charset=UTF-8 + date: + - Wed, 29 Jul 2026 19:15:23 GMT + server: + - scaffolding on HTTPServer2 + server-timing: + - gfet4t7; dur=888 + transfer-encoding: + - chunked + vary: + - Origin + - X-Origin + - Referer + x-content-type-options: + - nosniff + x-frame-options: + - SAMEORIGIN + x-gemini-service-tier: + - standard + x-xss-protection: + - '0' + status: + code: 200 + message: OK +- request: + body: '' + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-type: + - application/json + host: + - generativelanguage.googleapis.com + user-agent: + - google-genai-sdk/2.14.0 gl-python/3.14.6 + x-goog-api-client: + - google-genai-sdk/2.14.0 gl-python/3.14.6 + method: DELETE + uri: https://generativelanguage.googleapis.com/v1beta/cachedContents/aixcv88yts91zirycnmno31aitdexjt7uc6b19mw + response: + body: + string: '{} + + ' + headers: + alt-svc: + - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 + content-length: + - '3' + content-type: + - application/json; charset=UTF-8 + date: + - Wed, 29 Jul 2026 19:15:23 GMT + server: + - scaffolding on HTTPServer2 + server-timing: + - gfet4t7; dur=120 + 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_generated_audio_usage.yaml b/py/src/braintrust/integrations/google_genai/cassettes/latest/test_generated_audio_usage.yaml new file mode 100644 index 00000000..5ffc30d3 --- /dev/null +++ b/py/src/braintrust/integrations/google_genai/cassettes/latest/test_generated_audio_usage.yaml @@ -0,0 +1,65 @@ +interactions: +- request: + body: '{"contents": [{"parts": [{"text": "Read this sentence aloud in a warm voice: + Hello and welcome to Braintrust."}], "role": "user"}], "generationConfig": {"responseModalities": + ["AUDIO"], "speechConfig": {"voice_config": {"prebuilt_voice_config": {"voice_name": + "Kore"}}}}}' + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '271' + content-type: + - application/json + host: + - generativelanguage.googleapis.com + user-agent: + - google-genai-sdk/2.14.0 gl-python/3.14.6 + x-goog-api-client: + - google-genai-sdk/2.14.0 gl-python/3.14.6 + method: POST + uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-preview-tts:generateContent + response: + body: + string: '{"candidates": [{"content": {"parts": [{"inlineData": {"mimeType": + "audio/L16;codec=pcm;rate=24000", "data": "UklGRiQAAABXQVZFZm10IBAAAAABAAEAgD4AAAB9AAACABAAZGF0YQAAAAA="}}], + "role": "model"}, "finishReason": "STOP", "index": 0}], "usageMetadata": {"promptTokenCount": + 18, "candidatesTokenCount": 79, "totalTokenCount": 97, "promptTokensDetails": + [{"modality": "TEXT", "tokenCount": 18}], "candidatesTokensDetails": [{"modality": + "AUDIO", "tokenCount": 79}], "serviceTier": "standard"}, "modelVersion": "gemini-2.5-flash-preview-tts", + "responseId": "YlFqarfkLsLJ-8YPwIKYQA"}' + headers: + alt-svc: + - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 + content-length: + - '203712' + content-type: + - application/json; charset=UTF-8 + date: + - Wed, 29 Jul 2026 19:15:49 GMT + server: + - scaffolding on HTTPServer2 + server-timing: + - gfet4t7; dur=2784 + transfer-encoding: + - chunked + vary: + - Origin + - X-Origin + - Referer + x-content-type-options: + - nosniff + x-frame-options: + - SAMEORIGIN + x-gemini-service-tier: + - standard + x-xss-protection: + - '0' + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/google_genai/cassettes/latest/test_generated_image_usage.yaml b/py/src/braintrust/integrations/google_genai/cassettes/latest/test_generated_image_usage.yaml new file mode 100644 index 00000000..79929d40 --- /dev/null +++ b/py/src/braintrust/integrations/google_genai/cassettes/latest/test_generated_image_usage.yaml @@ -0,0 +1,64 @@ +interactions: +- request: + body: '{"contents": [{"parts": [{"text": "Generate a simple blue square icon."}], + "role": "user"}], "generationConfig": {"responseModalities": ["TEXT", "IMAGE"]}}' + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '155' + content-type: + - application/json + host: + - generativelanguage.googleapis.com + user-agent: + - google-genai-sdk/2.14.0 gl-python/3.14.6 + x-goog-api-client: + - google-genai-sdk/2.14.0 gl-python/3.14.6 + method: POST + uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image:generateContent + response: + body: + string: '{"candidates": [{"content": {"parts": [{"text": "Okay, here is your + simple blue square icon: "}, {"inlineData": {"mimeType": "image/png", "data": + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="}}], + "role": "model"}, "finishReason": "STOP", "index": 0}], "usageMetadata": {"promptTokenCount": + 8, "candidatesTokenCount": 1301, "totalTokenCount": 1309, "promptTokensDetails": + [{"modality": "TEXT", "tokenCount": 8}], "candidatesTokensDetails": [{"modality": + "IMAGE", "tokenCount": 1290}], "serviceTier": "standard"}, "modelVersion": + "gemini-2.5-flash-image", "responseId": "p1Bqau2WMYy0jrEPpYy64A0"}' + headers: + alt-svc: + - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 + content-length: + - '1031610' + content-type: + - application/json; charset=UTF-8 + date: + - Wed, 29 Jul 2026 19:12:43 GMT + server: + - scaffolding on HTTPServer2 + server-timing: + - gfet4t7; dur=3829 + transfer-encoding: + - chunked + vary: + - Origin + - X-Origin + - Referer + x-content-type-options: + - nosniff + x-frame-options: + - SAMEORIGIN + x-gemini-service-tier: + - standard + 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 60c626dd..02f6197e 100644 --- a/py/src/braintrust/integrations/google_genai/test_google_genai.py +++ b/py/src/braintrust/integrations/google_genai/test_google_genai.py @@ -34,57 +34,69 @@ REASONING_MODEL = "gemini-2.5-flash" TOOL_MODEL = "gemini-2.5-flash" if os.environ.get("BRAINTRUST_TEST_PACKAGE_VERSION") == "latest" else MODEL INTERACTIONS_MODEL = "gemini-2.5-flash" +GENERATED_AUDIO_MODEL = "gemini-2.5-flash-preview-tts" +GENERATED_IMAGE_MODEL = "gemini-2.5-flash-image" FIXTURES_DIR = Path(__file__).parent.parent.parent / "fixtures" TINY_PNG_BASE64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg==" +TINY_WAV_BASE64 = "UklGRiQAAABXQVZFZm10IBAAAAABAAEAgD4AAAB9AAACABAAZGF0YQAAAAA=" +LATEST_ONLY = pytest.mark.skipif( + os.environ.get("BRAINTRUST_TEST_PACKAGE_VERSION") != "latest", + reason="model is only covered by the latest google-genai matrix entry", +) -def _sanitize_generate_images_body(value): +def _sanitize_media_body(value): if isinstance(value, dict): - return { + sanitized = { key: ( - TINY_PNG_BASE64 - if key == "bytesBase64Encoded" and isinstance(val, str) - else _sanitize_generate_images_body(val) + TINY_PNG_BASE64 if key == "bytesBase64Encoded" and isinstance(val, str) else _sanitize_media_body(val) ) for key, val in value.items() } + inline_data = sanitized.get("inlineData") + if isinstance(inline_data, dict) and isinstance(inline_data.get("data"), str): + mime_type = inline_data.get("mimeType") + if isinstance(mime_type, str) and mime_type.startswith("image/"): + inline_data["data"] = TINY_PNG_BASE64 + elif isinstance(mime_type, str) and mime_type.startswith("audio/"): + inline_data["data"] = TINY_WAV_BASE64 + return sanitized if isinstance(value, list): - return [_sanitize_generate_images_body(item) for item in value] + return [_sanitize_media_body(item) for item in value] return value -def _sanitize_generate_images_response(response): - body = response.get("body", {}) - payload = body.get("string") +def _sanitize_media_payload(payload): if not payload: - return response + return payload is_bytes = isinstance(payload, bytes) - is_gzipped = False - + is_gzipped = is_bytes and payload[:2] == b"\x1f\x8b" + raw_payload = gzip.decompress(payload) if is_gzipped else payload if is_bytes: - raw_payload = payload - if raw_payload[:2] == b"\x1f\x8b": - raw_payload = gzip.decompress(raw_payload) - is_gzipped = True - payload = raw_payload.decode("utf-8") + raw_payload = raw_payload.decode("utf-8") try: - parsed = json.loads(payload) + parsed = json.loads(raw_payload) except Exception: - return response + return payload - sanitized = _sanitize_generate_images_body(parsed) + sanitized = _sanitize_media_body(parsed) if sanitized == parsed: - return response + return payload sanitized_payload = json.dumps(sanitized) - if is_bytes: - body["string"] = ( - gzip.compress(sanitized_payload.encode("utf-8")) if is_gzipped else sanitized_payload.encode("utf-8") - ) - else: - body["string"] = sanitized_payload + if not is_bytes: + return sanitized_payload + sanitized_bytes = sanitized_payload.encode("utf-8") + return gzip.compress(sanitized_bytes) if is_gzipped else sanitized_bytes + + +def _sanitize_media_response(response): + body = response.get("body", {}) + payload = body.get("string") + if payload: + body["string"] = _sanitize_media_payload(payload) return response @@ -96,10 +108,11 @@ def vcr_config(): def before_record_request(request): # Normalize HTTP method to uppercase for consistency (Google API quirk) request.method = request.method.upper() + request.body = _sanitize_media_payload(request.body) return request def before_record_response(response): - return _sanitize_generate_images_response(response) + return _sanitize_media_response(response) return { "record_mode": record_mode, @@ -171,6 +184,18 @@ def _assert_binary_not_logged(span, binary_data): assert binary_data[:8].hex() not in span_str +def _modality_token_count(details, modality): + return sum( + detail.token_count or 0 + for detail in details or [] + if getattr(detail.modality, "value", detail.modality) == modality + ) + + +def _serialized_modality_details(details): + return [detail.model_dump(exclude_none=True) for detail in details or []] + + # Test 1: Basic Completion (Sync) @pytest.mark.vcr @pytest.mark.parametrize( @@ -509,6 +534,33 @@ def test_document_input(memory_logger): _assert_metrics_are_valid(span["metrics"], start, end) +@LATEST_ONLY +@pytest.mark.vcr +def test_audio_input_usage(memory_logger): + """Map Google prompt-side audio modality details to prompt_audio_tokens.""" + assert not memory_logger.pop() + + audio_data = (FIXTURES_DIR / "test_audio.wav").read_bytes() + client = Client() + response = client.models.generate_content( + model=MODEL, + contents=[ + types.Part.from_bytes(data=audio_data, mime_type="audio/wav"), + types.Part.from_text(text="Describe this audio in five words or less."), + ], + config=types.GenerateContentConfig(max_output_tokens=100), + ) + + assert response.text + assert response.usage_metadata + expected_audio_tokens = _modality_token_count(response.usage_metadata.prompt_tokens_details, "AUDIO") + assert expected_audio_tokens > 0 + + spans = memory_logger.pop() + assert len(spans) == 1 + assert spans[0]["metrics"]["prompt_audio_tokens"] == expected_audio_tokens + + @pytest.mark.vcr def test_image_input_wrapped_in_content(memory_logger): """Verify binary Parts inside a Content wrapper are traced as attachments. @@ -1030,16 +1082,74 @@ def test_reasoning(memory_logger): assert len(spans) == 2 first_span, second_span = spans - for span in spans: + for span, response in zip(spans, (first_response, follow_up_response), strict=True): + usage_metadata = response.usage_metadata + assert usage_metadata is not None assert span["metadata"]["model"] == REASONING_MODEL assert span["input"]["config"]["thinking_config"]["include_thoughts"] is True - assert span["metrics"]["completion_reasoning_tokens"] > 0 + assert span["metrics"]["completion_reasoning_tokens"] == usage_metadata.thoughts_token_count + assert span["metrics"]["completion_tokens"] == ( + usage_metadata.candidates_token_count + usage_metadata.thoughts_token_count + ) assert span["output"] assert first_prompt in str(first_span["input"]) assert follow_up_prompt in str(second_span["input"]) +@LATEST_ONLY +@pytest.mark.vcr +def test_generated_audio_usage(memory_logger): + """Map Google candidate-side audio modality details to completion_audio_tokens.""" + assert not memory_logger.pop() + + client = Client() + response = client.models.generate_content( + model=GENERATED_AUDIO_MODEL, + contents="Read this sentence aloud in a warm voice: Hello and welcome to Braintrust.", + config=types.GenerateContentConfig( + response_modalities=["AUDIO"], + speech_config=types.SpeechConfig( + voice_config=types.VoiceConfig( + prebuilt_voice_config=types.PrebuiltVoiceConfig(voice_name="Kore"), + ) + ), + ), + ) + + assert response.candidates + assert response.usage_metadata + expected_audio_tokens = _modality_token_count(response.usage_metadata.candidates_tokens_details, "AUDIO") + assert expected_audio_tokens > 0 + + spans = memory_logger.pop() + assert len(spans) == 1 + assert spans[0]["metrics"]["completion_audio_tokens"] == expected_audio_tokens + + +@LATEST_ONLY +@pytest.mark.vcr +def test_generated_image_usage(memory_logger): + """Map Google candidate-side image modality details to completion_image_tokens.""" + assert not memory_logger.pop() + + client = Client() + response = client.models.generate_content( + model=GENERATED_IMAGE_MODEL, + contents="Generate a simple blue square icon.", + config=types.GenerateContentConfig(response_modalities=["TEXT", "IMAGE"]), + ) + + assert response.candidates + assert response.usage_metadata + expected_image_tokens = _modality_token_count(response.usage_metadata.candidates_tokens_details, "IMAGE") + assert expected_image_tokens > 0 + + spans = memory_logger.pop() + assert len(spans) == 1 + assert spans[0]["metrics"]["completion_image_tokens"] == expected_image_tokens + + def test_attachment_in_config(memory_logger): """Test that attachments in config are preserved through serialization.""" from braintrust.bt_json import bt_safe_deep_copy @@ -1331,6 +1441,7 @@ def test_google_search_grounding(memory_logger, mode): client = Client() start = time.time() + usage_metadata = None if mode == "sync": response = client.models.generate_content( model=GROUNDING_MODEL, @@ -1341,6 +1452,7 @@ def test_google_search_grounding(memory_logger, mode): ), ) text = response.text + usage_metadata = response.usage_metadata elif mode == "stream": stream = client.models.generate_content_stream( model=GROUNDING_MODEL, @@ -1354,6 +1466,8 @@ def test_google_search_grounding(memory_logger, mode): for chunk in stream: if chunk.text: text += chunk.text + if chunk.usage_metadata: + usage_metadata = chunk.usage_metadata end = time.time() @@ -1369,11 +1483,56 @@ def test_google_search_grounding(memory_logger, mode): assert "population" in str(span["input"]).lower() or "Tokyo" in str(span["input"]) assert span["output"] _assert_metrics_are_valid(span["metrics"], start, end) + assert usage_metadata is not None + assert span["metrics"]["prompt_tokens"] == usage_metadata.prompt_token_count + ( + usage_metadata.tool_use_prompt_token_count or 0 + ) + assert span["metrics"]["completion_tokens"] == usage_metadata.candidates_token_count + ( + usage_metadata.thoughts_token_count or 0 + ) + assert span["metrics"]["tokens"] == (span["metrics"]["prompt_tokens"] + span["metrics"]["completion_tokens"]) + tool_use_details = _serialized_modality_details(usage_metadata.tool_use_prompt_tokens_details) + if tool_use_details: + assert span["metadata"]["usage_by_modality"]["tool_use_prompt_tokens_details"] == tool_use_details # Verify grounding metadata is captured _assert_grounding_metadata(span["output"]) +@LATEST_ONLY +@pytest.mark.vcr +def test_cached_content_usage_metadata(memory_logger): + """Preserve Google's cache modality details as provider-specific metadata.""" + assert not memory_logger.pop() + + client = Client() + cached_content = client.caches.create( + model=REASONING_MODEL, + config=types.CreateCachedContentConfig( + display_name="braintrust-google-genai-vcr-test", + contents="The sky is blue because molecules scatter blue light. " * 150, + ttl="300s", + ), + ) + try: + response = client.models.generate_content( + model=REASONING_MODEL, + contents="In one sentence, why is the sky blue?", + config=types.GenerateContentConfig(cached_content=cached_content.name, max_output_tokens=100), + ) + finally: + client.caches.delete(name=cached_content.name) + + assert response.text + assert response.usage_metadata + cache_details = _serialized_modality_details(response.usage_metadata.cache_tokens_details) + assert cache_details + + spans = memory_logger.pop() + assert len(spans) == 1 + assert spans[0]["metadata"]["usage_by_modality"]["cache_tokens_details"] == cache_details + + # Test: Google Search Grounding (Async) @pytest.mark.vcr @pytest.mark.asyncio @@ -1491,8 +1650,12 @@ def test_interactions_create_and_get(memory_logger): assert get_span["metadata"]["provider"] == "google" assert create_span["output"]["status"] == "completed" assert "Paris" in create_span["output"]["text"] - assert create_span["metrics"]["prompt_tokens"] > 0 - assert create_span["metrics"]["completion_tokens"] > 0 + assert create_span["metrics"]["prompt_tokens"] == response.usage.total_input_tokens + assert create_span["metrics"]["completion_reasoning_tokens"] == response.usage.total_thought_tokens + assert create_span["metrics"]["completion_tokens"] == ( + response.usage.total_output_tokens + response.usage.total_thought_tokens + ) + assert create_span["metrics"]["tokens"] == response.usage.total_tokens assert get_span["input"]["id"] == response.id assert get_span["metadata"]["interaction_id"] == response.id @@ -1524,6 +1687,11 @@ def test_interactions_create_stream(memory_logger): assert create_span["output"]["status"] == "completed" assert create_span["output"]["text"] assert create_span["metrics"]["time_to_first_token"] >= 0 + final_interaction = next(event.interaction for event in reversed(events) if getattr(event, "interaction", None)) + assert create_span["metrics"]["completion_reasoning_tokens"] == final_interaction.usage.total_thought_tokens + assert create_span["metrics"]["completion_tokens"] == ( + final_interaction.usage.total_output_tokens + final_interaction.usage.total_thought_tokens + ) assert any( event_type in create_span["metadata"]["stream_event_types"] for event_type in ("content.start", "step.start") ) diff --git a/py/src/braintrust/integrations/google_genai/tracing.py b/py/src/braintrust/integrations/google_genai/tracing.py index f0c07bfb..82ce43f1 100644 --- a/py/src/braintrust/integrations/google_genai/tracing.py +++ b/py/src/braintrust/integrations/google_genai/tracing.py @@ -6,7 +6,7 @@ import logging import time from collections.abc import Awaitable, Callable, Iterable -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, TypeAlias from braintrust.integrations.utils import _materialize_attachment from braintrust.logger import ( @@ -63,6 +63,7 @@ def start_span(*args, **kwargs): "mcp_server_tool_result", "file_search_result", } +_LoggedResult: TypeAlias = tuple[Any, dict[str, Any]] | tuple[Any, dict[str, Any], dict[str, Any] | None] @dataclasses.dataclass @@ -330,19 +331,61 @@ def _prepare_interaction_id_traced_call( # --------------------------------------------------------------------------- +def _extract_modality_token_count(details: Any, modality: str) -> int | float | None: + counts = [] + for detail in details or []: + detail_modality = getattr(detail, "modality", None) + detail_modality = getattr(detail_modality, "value", detail_modality) + token_count = getattr(detail, "token_count", None) + if isinstance(detail_modality, str) and detail_modality.upper() == modality and token_count is not None: + counts.append(token_count) + return sum(counts) if counts else None + + def _extract_usage_metadata_metrics( usage_metadata: "GenerateContentResponseUsageMetadata", metrics: dict[str, Any] ) -> None: - if hasattr(usage_metadata, "prompt_token_count"): - metrics["prompt_tokens"] = usage_metadata.prompt_token_count - if hasattr(usage_metadata, "candidates_token_count"): - metrics["completion_tokens"] = usage_metadata.candidates_token_count + prompt_token_count = getattr(usage_metadata, "prompt_token_count", None) + tool_use_prompt_token_count = getattr(usage_metadata, "tool_use_prompt_token_count", None) + if prompt_token_count is not None or tool_use_prompt_token_count is not None: + metrics["prompt_tokens"] = (prompt_token_count or 0) + (tool_use_prompt_token_count or 0) + + candidates_token_count = getattr(usage_metadata, "candidates_token_count", None) + thoughts_token_count = getattr(usage_metadata, "thoughts_token_count", None) + if candidates_token_count is not None or thoughts_token_count is not None: + metrics["completion_tokens"] = (candidates_token_count or 0) + (thoughts_token_count or 0) + if hasattr(usage_metadata, "total_token_count"): metrics["tokens"] = usage_metadata.total_token_count if hasattr(usage_metadata, "cached_content_token_count"): metrics["prompt_cached_tokens"] = usage_metadata.cached_content_token_count - if hasattr(usage_metadata, "thoughts_token_count"): - metrics["completion_reasoning_tokens"] = usage_metadata.thoughts_token_count + if thoughts_token_count is not None: + metrics["completion_reasoning_tokens"] = thoughts_token_count + + prompt_audio_tokens = _extract_modality_token_count( + getattr(usage_metadata, "prompt_tokens_details", None), "AUDIO" + ) + if prompt_audio_tokens is not None: + metrics["prompt_audio_tokens"] = prompt_audio_tokens + + candidates_tokens_details = getattr(usage_metadata, "candidates_tokens_details", None) + completion_audio_tokens = _extract_modality_token_count(candidates_tokens_details, "AUDIO") + if completion_audio_tokens is not None: + metrics["completion_audio_tokens"] = completion_audio_tokens + completion_image_tokens = _extract_modality_token_count(candidates_tokens_details, "IMAGE") + if completion_image_tokens is not None: + metrics["completion_image_tokens"] = completion_image_tokens + + +def _extract_usage_metadata_provider_metadata( + usage_metadata: "GenerateContentResponseUsageMetadata", +) -> dict[str, Any] | None: + usage_by_modality = {} + for name in ("cache_tokens_details", "tool_use_prompt_tokens_details"): + details = getattr(usage_metadata, name, None) + if details: + usage_by_modality[name] = _materialize_interaction_value(details) + return {"usage_by_modality": usage_by_modality} if usage_by_modality else None def _extract_generate_content_metrics(response: "GenerateContentResponse", start: float) -> dict[str, Any]: @@ -472,14 +515,18 @@ def _extract_interaction_usage_metrics(usage: Any, metrics: dict[str, Any]) -> N if hasattr(usage, "total_input_tokens") and usage.total_input_tokens is not None: metrics["prompt_tokens"] = usage.total_input_tokens - if hasattr(usage, "total_output_tokens") and usage.total_output_tokens is not None: - metrics["completion_tokens"] = usage.total_output_tokens + + total_output_tokens = getattr(usage, "total_output_tokens", None) + total_thought_tokens = getattr(usage, "total_thought_tokens", None) + if total_output_tokens is not None or total_thought_tokens is not None: + metrics["completion_tokens"] = (total_output_tokens or 0) + (total_thought_tokens or 0) + if hasattr(usage, "total_tokens") and usage.total_tokens is not None: metrics["tokens"] = usage.total_tokens if hasattr(usage, "total_cached_tokens") and usage.total_cached_tokens is not None: metrics["prompt_cached_tokens"] = usage.total_cached_tokens - if hasattr(usage, "total_thought_tokens") and usage.total_thought_tokens is not None: - metrics["completion_reasoning_tokens"] = usage.total_thought_tokens + if total_thought_tokens is not None: + metrics["completion_reasoning_tokens"] = total_thought_tokens if hasattr(usage, "total_tool_use_tokens") and usage.total_tool_use_tokens is not None: metrics["tool_use_tokens"] = usage.total_tool_use_tokens @@ -568,8 +615,15 @@ def _extract_interaction_metrics(response: "Interaction", start: float) -> dict[ # --------------------------------------------------------------------------- -def _gc_process_result(result: "GenerateContentResponse", start: float) -> tuple[Any, dict[str, Any]]: - return result, _extract_generate_content_metrics(result, start) +def _gc_process_result( + result: "GenerateContentResponse", start: float +) -> tuple[Any, dict[str, Any], dict[str, Any] | None]: + usage_metadata = getattr(result, "usage_metadata", None) + return ( + result, + _extract_generate_content_metrics(result, start), + _extract_usage_metadata_provider_metadata(usage_metadata) if usage_metadata is not None else None, + ) def _embed_process_result(result: "EmbedContentResponse", start: float) -> tuple[Any, dict[str, Any]]: @@ -631,7 +685,7 @@ def _generic_process_result(result: Any, start: float) -> tuple[Any, dict[str, A def _aggregate_generate_content_chunks( chunks: "list[GenerateContentResponse]", start: float, first_token_time: float | None = None -) -> tuple[dict[str, Any], dict[str, Any]]: +) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any] | None]: end_time = time.time() metrics = dict( start=start, @@ -643,7 +697,7 @@ def _aggregate_generate_content_chunks( metrics["time_to_first_token"] = first_token_time - start if not chunks: - return {}, metrics + return {}, metrics, None text = "" thought_text = "" @@ -706,7 +760,11 @@ def _aggregate_generate_content_chunks( if text: aggregated["text"] = text - return aggregated, clean_nones(dict(metrics)) + return ( + aggregated, + clean_nones(dict(metrics)), + _extract_usage_metadata_provider_metadata(usage_metadata) if usage_metadata is not None else None, + ) def _is_interaction_content_event(event: Any) -> bool: @@ -1036,7 +1094,7 @@ def _run_traced_call( *, name: str, invoke: Callable[[], Any], - process_result: Callable[[Any, float], tuple[Any, dict[str, Any]] | tuple[Any, dict[str, Any], dict[str, Any]]], + process_result: Callable[[Any, float], _LoggedResult], prepare_call: Callable[ [Any, list[Any], dict[str, Any]], tuple[dict[str, Any], dict[str, Any]] ] = _prepare_traced_call, @@ -1060,7 +1118,7 @@ def _run_traced_call( span.log(output=output, metrics=metrics, metadata=metadata) parent_export = span.export() - if finalize_logged_output is not None and parent_export is not None and metrics is not None: + if finalize_logged_output is not None and parent_export is not None: finalize_logged_output(output, metrics, metadata, parent_export) return result @@ -1073,7 +1131,7 @@ async def _run_async_traced_call( *, name: str, invoke: Callable[[], Awaitable[Any]], - process_result: Callable[[Any, float], tuple[Any, dict[str, Any]] | tuple[Any, dict[str, Any], dict[str, Any]]], + process_result: Callable[[Any, float], _LoggedResult], prepare_call: Callable[ [Any, list[Any], dict[str, Any]], tuple[dict[str, Any], dict[str, Any]] ] = _prepare_traced_call, @@ -1097,7 +1155,7 @@ async def _run_async_traced_call( span.log(output=output, metrics=metrics, metadata=metadata) parent_export = span.export() - if finalize_logged_output is not None and parent_export is not None and metrics is not None: + if finalize_logged_output is not None and parent_export is not None: finalize_logged_output(output, metrics, metadata, parent_export) return result @@ -1150,9 +1208,7 @@ def _run_stream_traced_call( *, name: str, invoke: Callable[[], Any], - aggregate: Callable[ - [list[Any], float, float | None], tuple[Any, dict[str, Any]] | tuple[Any, dict[str, Any], dict[str, Any]] - ], + aggregate: Callable[[list[Any], float, float | None], _LoggedResult], span_type: SpanTypeAttribute = SpanTypeAttribute.LLM, first_token_predicate: Callable[[Any], bool] | None = None, prepare_call: Callable[ @@ -1194,7 +1250,7 @@ def stream_generator(): span.log(output=output, metrics=metrics, metadata=metadata) parent_export = span.export() - if finalize_logged_output is not None and parent_export is not None and metrics is not None: + if finalize_logged_output is not None and parent_export is not None: finalize_logged_output(output, metrics, metadata, parent_export) return output @@ -1209,9 +1265,7 @@ def _run_async_stream_traced_call( *, name: str, invoke: Callable[[], Awaitable[Any]], - aggregate: Callable[ - [list[Any], float, float | None], tuple[Any, dict[str, Any]] | tuple[Any, dict[str, Any], dict[str, Any]] - ], + aggregate: Callable[[list[Any], float, float | None], _LoggedResult], span_type: SpanTypeAttribute = SpanTypeAttribute.LLM, first_token_predicate: Callable[[Any], bool] | None = None, prepare_call: Callable[ @@ -1253,7 +1307,7 @@ async def stream_generator(): span.log(output=output, metrics=metrics, metadata=metadata) parent_export = span.export() - if finalize_logged_output is not None and parent_export is not None and metrics is not None: + if finalize_logged_output is not None and parent_export is not None: finalize_logged_output(output, metrics, metadata, parent_export) return stream_generator()