From ff43ceaaee5288efcdd2cc5667fc463adc336763 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 4 Aug 2026 12:31:44 -0400 Subject: [PATCH] fix(adk): align Google usage token totals ADK traces omitted tool-use prompt tokens from prompt totals and reasoning tokens from completion totals. Customers using Google Search or reasoning models therefore saw token breakdowns that disagreed with Google's reported total, making usage analysis and provider reconciliation misleading. Keep the normalized breakdown consistent with the provider total: prompt + tool-use prompt + candidates + thoughts = total Retain reasoning and tool-use details so customers can still inspect where the aggregate usage came from. Cover both supported ADK versions with real Google Search recordings. --- .../test_adk_usage_metadata_metrics.yaml | 133 ++++++++++++++++ .../test_adk_usage_metadata_metrics.yaml | 147 ++++++++++++++++++ .../braintrust/integrations/adk/test_adk.py | 49 ++++++ py/src/braintrust/integrations/adk/tracing.py | 39 ++--- .../integrations/google_genai/tracing.py | 64 ++------ py/src/braintrust/integrations/test_utils.py | 12 ++ py/src/braintrust/integrations/utils.py | 68 ++++++++ 7 files changed, 433 insertions(+), 79 deletions(-) create mode 100644 py/src/braintrust/integrations/adk/cassettes/1.14.1/test_adk_usage_metadata_metrics.yaml create mode 100644 py/src/braintrust/integrations/adk/cassettes/latest/test_adk_usage_metadata_metrics.yaml diff --git a/py/src/braintrust/integrations/adk/cassettes/1.14.1/test_adk_usage_metadata_metrics.yaml b/py/src/braintrust/integrations/adk/cassettes/1.14.1/test_adk_usage_metadata_metrics.yaml new file mode 100644 index 000000000..bf068b1d6 --- /dev/null +++ b/py/src/braintrust/integrations/adk/cassettes/1.14.1/test_adk_usage_metadata_metrics.yaml @@ -0,0 +1,133 @@ +interactions: +- request: + body: '{"contents": [{"parts": [{"text": "What is the current population of Tokyo, + Japan? Answer in one sentence."}], "role": "user"}], "systemInstruction": {"parts": + [{"text": "Use Google Search to answer the user''s question accurately and concisely.\n\nYou + are an agent. Your internal name is \"usage_metadata_agent\"."}], "role": "user"}, + "tools": [{"googleSearch": {}}], "generationConfig": {}}' + headers: + Content-Type: + - application/json + user-agent: + - google-genai-sdk/1.75.0 gl-python/3.14.6 google-adk/1.14.1 gl-python/3.14.6 + x-goog-api-client: + - google-genai-sdk/1.75.0 gl-python/3.14.6 google-adk/1.14.1 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\": \"As of 2026, the estimated total population + of Tokyo is 14,270,748.\"\n }\n ],\n \"role\": \"model\"\n + \ },\n \"finishReason\": \"STOP\",\n \"index\": 0,\n \"groundingMetadata\": + {\n \"searchEntryPoint\": {\n \"renderedContent\": \"\\u003cstyle\\u003e\\n.container + {\\n align-items: center;\\n border-radius: 8px;\\n display: flex;\\n font-family: + Google Sans, Roboto, sans-serif;\\n font-size: 14px;\\n line-height: 20px;\\n + \ padding: 8px 12px;\\n}\\n.chip {\\n display: inline-block;\\n border: + solid 1px;\\n border-radius: 16px;\\n min-width: 14px;\\n padding: 5px + 16px;\\n text-align: center;\\n user-select: none;\\n margin: 0 8px;\\n + \ -webkit-tap-highlight-color: transparent;\\n}\\n.carousel {\\n overflow: + auto;\\n scrollbar-width: none;\\n white-space: nowrap;\\n margin-right: + -12px;\\n}\\n.headline {\\n display: flex;\\n margin-right: 4px;\\n}\\n.gradient-container + {\\n position: relative;\\n}\\n.gradient {\\n position: absolute;\\n transform: + translate(3px, -9px);\\n height: 36px;\\n width: 9px;\\n}\\n@media (prefers-color-scheme: + light) {\\n .container {\\n background-color: #fafafa;\\n box-shadow: + 0 0 0 1px #0000000f;\\n }\\n .headline-label {\\n color: #1f1f1f;\\n + \ }\\n .chip {\\n background-color: #ffffff;\\n border-color: #d2d2d2;\\n + \ color: #5e5e5e;\\n text-decoration: none;\\n }\\n .chip:hover {\\n + \ background-color: #f2f2f2;\\n }\\n .chip:focus {\\n background-color: + #f2f2f2;\\n }\\n .chip:active {\\n background-color: #d8d8d8;\\n border-color: + #b6b6b6;\\n }\\n .logo-dark {\\n display: none;\\n }\\n .gradient {\\n + \ background: linear-gradient(90deg, #fafafa 15%, #fafafa00 100%);\\n }\\n}\\n@media + (prefers-color-scheme: dark) {\\n .container {\\n background-color: #1f1f1f;\\n + \ box-shadow: 0 0 0 1px #ffffff26;\\n }\\n .headline-label {\\n color: + #fff;\\n }\\n .chip {\\n background-color: #2c2c2c;\\n border-color: + #3c4043;\\n color: #fff;\\n text-decoration: none;\\n }\\n .chip:hover + {\\n background-color: #353536;\\n }\\n .chip:focus {\\n background-color: + #353536;\\n }\\n .chip:active {\\n background-color: #464849;\\n border-color: + #53575b;\\n }\\n .logo-light {\\n display: none;\\n }\\n .gradient + {\\n background: linear-gradient(90deg, #1f1f1f 15%, #1f1f1f00 100%);\\n + \ }\\n}\\n\\u003c/style\\u003e\\n\\u003cdiv class=\\\"container\\\"\\u003e\\n + \ \\u003cdiv class=\\\"headline\\\"\\u003e\\n \\u003csvg class=\\\"logo-light\\\" + width=\\\"18\\\" height=\\\"18\\\" viewBox=\\\"9 9 35 35\\\" fill=\\\"none\\\" + xmlns=\\\"http://www.w3.org/2000/svg\\\"\\u003e\\n \\u003cpath fill-rule=\\\"evenodd\\\" + clip-rule=\\\"evenodd\\\" d=\\\"M42.8622 27.0064C42.8622 25.7839 42.7525 24.6084 + 42.5487 23.4799H26.3109V30.1568H35.5897C35.1821 32.3041 33.9596 34.1222 32.1258 + 35.3448V39.6864H37.7213C40.9814 36.677 42.8622 32.2571 42.8622 27.0064V27.0064Z\\\" + fill=\\\"#4285F4\\\"/\\u003e\\n \\u003cpath fill-rule=\\\"evenodd\\\" + clip-rule=\\\"evenodd\\\" d=\\\"M26.3109 43.8555C30.9659 43.8555 34.8687 42.3195 + 37.7213 39.6863L32.1258 35.3447C30.5898 36.3792 28.6306 37.0061 26.3109 37.0061C21.8282 + 37.0061 18.0195 33.9811 16.6559 29.906H10.9194V34.3573C13.7563 39.9841 19.5712 + 43.8555 26.3109 43.8555V43.8555Z\\\" fill=\\\"#34A853\\\"/\\u003e\\n \\u003cpath + fill-rule=\\\"evenodd\\\" clip-rule=\\\"evenodd\\\" d=\\\"M16.6559 29.8904C16.3111 + 28.8559 16.1074 27.7588 16.1074 26.6146C16.1074 25.4704 16.3111 24.3733 16.6559 + 23.3388V18.8875H10.9194C9.74388 21.2072 9.06992 23.8247 9.06992 26.6146C9.06992 + 29.4045 9.74388 32.022 10.9194 34.3417L15.3864 30.8621L16.6559 29.8904V29.8904Z\\\" + fill=\\\"#FBBC05\\\"/\\u003e\\n \\u003cpath fill-rule=\\\"evenodd\\\" + clip-rule=\\\"evenodd\\\" d=\\\"M26.3109 16.2386C28.85 16.2386 31.107 17.1164 + 32.9095 18.8091L37.8466 13.8719C34.853 11.082 30.9659 9.3736 26.3109 9.3736C19.5712 + 9.3736 13.7563 13.245 10.9194 18.8875L16.6559 23.3388C18.0195 19.2636 21.8282 + 16.2386 26.3109 16.2386V16.2386Z\\\" fill=\\\"#EA4335\\\"/\\u003e\\n \\u003c/svg\\u003e\\n + \ \\u003csvg class=\\\"logo-dark\\\" width=\\\"18\\\" height=\\\"18\\\" + viewBox=\\\"0 0 48 48\\\" xmlns=\\\"http://www.w3.org/2000/svg\\\"\\u003e\\n + \ \\u003ccircle cx=\\\"24\\\" cy=\\\"23\\\" fill=\\\"#FFF\\\" r=\\\"22\\\"/\\u003e\\n + \ \\u003cpath d=\\\"M33.76 34.26c2.75-2.56 4.49-6.37 4.49-11.26 0-.89-.08-1.84-.29-3H24.01v5.99h8.03c-.4 + 2.02-1.5 3.56-3.07 4.56v.75l3.91 2.97h.88z\\\" fill=\\\"#4285F4\\\"/\\u003e\\n + \ \\u003cpath d=\\\"M15.58 25.77A8.845 8.845 0 0 0 24 31.86c1.92 0 3.62-.46 + 4.97-1.31l4.79 3.71C31.14 36.7 27.65 38 24 38c-5.93 0-11.01-3.4-13.45-8.36l.17-1.01 + 4.06-2.85h.8z\\\" fill=\\\"#34A853\\\"/\\u003e\\n \\u003cpath d=\\\"M15.59 + 20.21a8.864 8.864 0 0 0 0 5.58l-5.03 3.86c-.98-2-1.53-4.25-1.53-6.64 0-2.39.55-4.64 + 1.53-6.64l1-.22 3.81 2.98.22 1.08z\\\" fill=\\\"#FBBC05\\\"/\\u003e\\n \\u003cpath + d=\\\"M24 14.14c2.11 0 4.02.75 5.52 1.98l4.36-4.36C31.22 9.43 27.81 8 24 8c-5.93 + 0-11.01 3.4-13.45 8.36l5.03 3.85A8.86 8.86 0 0 1 24 14.14z\\\" fill=\\\"#EA4335\\\"/\\u003e\\n + \ \\u003c/svg\\u003e\\n \\u003cdiv class=\\\"gradient-container\\\"\\u003e\\u003cdiv + class=\\\"gradient\\\"\\u003e\\u003c/div\\u003e\\u003c/div\\u003e\\n \\u003c/div\\u003e\\n + \ \\u003cdiv class=\\\"carousel\\\"\\u003e\\n \\u003ca class=\\\"chip\\\" + href=\\\"https://vertexaisearch.cloud.google.com/grounding-api-redirect/AUZIYQH9hELTUaoQVeZdpYUpyo2tui_DlXnI201Tl6CbR859pFwuXQXu9JMa6Ejb9nAs7U8l0gU0sGaV9jF3syljqRMrpOd0_xxZ8Rz4izTjYtjVyWHajQXjbafTrDXlY2mW-CMpUeEm5Wg0q63vt6C4pSR5PEDbueY7gXfgStqeNhtLgaYeNaS6WvgCa1hAoUWWZOQMeKChF_006JwKqoWXuWISSkomnQ==\\\"\\u003ecurrent + population of Tokyo, Japan\\u003c/a\\u003e\\n \\u003c/div\\u003e\\n\\u003c/div\\u003e\\n\"\n + \ },\n \"groundingChunks\": [\n {\n \"web\": + {\n \"uri\": \"https://vertexaisearch.cloud.google.com/grounding-api-redirect/AUZIYQEpi7ONwA9RPjVU3PSHeR1t9uzsnQEIRM4YkGBYScGS_rM2QOBJVYBMgVwlDCatTBo8TbZds6lbK2jnOLtT41q5botXEYDZavDBzTCjnhjhkeSzOcyqukPUFfUk6M4maEBDIurGF0FCSoCt_vzk\",\n + \ \"title\": \"wikipedia.org\"\n }\n }\n ],\n + \ \"groundingSupports\": [\n {\n \"segment\": {\n + \ \"endIndex\": 66,\n \"text\": \"As of 2026, the + estimated total population of Tokyo is 14,270,748.\"\n },\n \"groundingChunkIndices\": + [\n 0\n ]\n }\n ],\n \"webSearchQueries\": + [\n \"current population of Tokyo, Japan\"\n ]\n }\n + \ }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": 48,\n \"candidatesTokenCount\": + 44,\n \"totalTokenCount\": 760,\n \"promptTokensDetails\": [\n {\n + \ \"modality\": \"TEXT\",\n \"tokenCount\": 48\n }\n ],\n + \ \"toolUsePromptTokenCount\": 100,\n \"toolUsePromptTokensDetails\": + [\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 100\n + \ }\n ],\n \"thoughtsTokenCount\": 568,\n \"serviceTier\": \"standard\"\n + \ },\n \"modelVersion\": \"gemini-2.5-flash\",\n \"responseId\": \"AhByarLrMtzH-8YP3uDUcQ\"\n}\n" + headers: + Alt-Svc: + - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 + Content-Type: + - application/json; charset=UTF-8 + Date: + - Tue, 04 Aug 2026 16:15:02 GMT + Server: + - scaffolding on HTTPServer2 + Server-Timing: + - gfet4t7; dur=4018 + 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' + content-length: + - '6819' + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/adk/cassettes/latest/test_adk_usage_metadata_metrics.yaml b/py/src/braintrust/integrations/adk/cassettes/latest/test_adk_usage_metadata_metrics.yaml new file mode 100644 index 000000000..f6b8d0ba3 --- /dev/null +++ b/py/src/braintrust/integrations/adk/cassettes/latest/test_adk_usage_metadata_metrics.yaml @@ -0,0 +1,147 @@ +interactions: +- request: + body: '{"contents": [{"parts": [{"text": "What is the current population of Tokyo, + Japan? Answer in one sentence."}], "role": "user"}], "systemInstruction": {"parts": + [{"text": "Use Google Search to answer the user''s question accurately and concisely.\n\nYou + are an agent. Your internal name is \"usage_metadata_agent\"."}], "role": "user"}, + "tools": [{"googleSearch": {}}], "generationConfig": {}}' + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '391' + content-type: + - application/json + host: + - generativelanguage.googleapis.com + user-agent: + - google-genai-sdk/2.16.0 gl-python/3.14.6 google-adk/2.5.0 gl-python/3.14.6 + x-goog-api-client: + - google-genai-sdk/2.16.0 gl-python/3.14.6 google-adk/2.5.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\": \"As of 2026, the estimated total population + of Tokyo Metropolis is approximately 14.27 million people.\"\n }\n + \ ],\n \"role\": \"model\"\n },\n \"finishReason\": + \"STOP\",\n \"index\": 0,\n \"groundingMetadata\": {\n \"searchEntryPoint\": + {\n \"renderedContent\": \"\\u003cstyle\\u003e\\n.container {\\n + \ align-items: center;\\n border-radius: 8px;\\n display: flex;\\n font-family: + Google Sans, Roboto, sans-serif;\\n font-size: 14px;\\n line-height: 20px;\\n + \ padding: 8px 12px;\\n}\\n.chip {\\n display: inline-block;\\n border: + solid 1px;\\n border-radius: 16px;\\n min-width: 14px;\\n padding: 5px + 16px;\\n text-align: center;\\n user-select: none;\\n margin: 0 8px;\\n + \ -webkit-tap-highlight-color: transparent;\\n}\\n.carousel {\\n overflow: + auto;\\n scrollbar-width: none;\\n white-space: nowrap;\\n margin-right: + -12px;\\n}\\n.headline {\\n display: flex;\\n margin-right: 4px;\\n}\\n.gradient-container + {\\n position: relative;\\n}\\n.gradient {\\n position: absolute;\\n transform: + translate(3px, -9px);\\n height: 36px;\\n width: 9px;\\n}\\n@media (prefers-color-scheme: + light) {\\n .container {\\n background-color: #fafafa;\\n box-shadow: + 0 0 0 1px #0000000f;\\n }\\n .headline-label {\\n color: #1f1f1f;\\n + \ }\\n .chip {\\n background-color: #ffffff;\\n border-color: #d2d2d2;\\n + \ color: #5e5e5e;\\n text-decoration: none;\\n }\\n .chip:hover {\\n + \ background-color: #f2f2f2;\\n }\\n .chip:focus {\\n background-color: + #f2f2f2;\\n }\\n .chip:active {\\n background-color: #d8d8d8;\\n border-color: + #b6b6b6;\\n }\\n .logo-dark {\\n display: none;\\n }\\n .gradient {\\n + \ background: linear-gradient(90deg, #fafafa 15%, #fafafa00 100%);\\n }\\n}\\n@media + (prefers-color-scheme: dark) {\\n .container {\\n background-color: #1f1f1f;\\n + \ box-shadow: 0 0 0 1px #ffffff26;\\n }\\n .headline-label {\\n color: + #fff;\\n }\\n .chip {\\n background-color: #2c2c2c;\\n border-color: + #3c4043;\\n color: #fff;\\n text-decoration: none;\\n }\\n .chip:hover + {\\n background-color: #353536;\\n }\\n .chip:focus {\\n background-color: + #353536;\\n }\\n .chip:active {\\n background-color: #464849;\\n border-color: + #53575b;\\n }\\n .logo-light {\\n display: none;\\n }\\n .gradient + {\\n background: linear-gradient(90deg, #1f1f1f 15%, #1f1f1f00 100%);\\n + \ }\\n}\\n\\u003c/style\\u003e\\n\\u003cdiv class=\\\"container\\\"\\u003e\\n + \ \\u003cdiv class=\\\"headline\\\"\\u003e\\n \\u003csvg class=\\\"logo-light\\\" + width=\\\"18\\\" height=\\\"18\\\" viewBox=\\\"9 9 35 35\\\" fill=\\\"none\\\" + xmlns=\\\"http://www.w3.org/2000/svg\\\"\\u003e\\n \\u003cpath fill-rule=\\\"evenodd\\\" + clip-rule=\\\"evenodd\\\" d=\\\"M42.8622 27.0064C42.8622 25.7839 42.7525 24.6084 + 42.5487 23.4799H26.3109V30.1568H35.5897C35.1821 32.3041 33.9596 34.1222 32.1258 + 35.3448V39.6864H37.7213C40.9814 36.677 42.8622 32.2571 42.8622 27.0064V27.0064Z\\\" + fill=\\\"#4285F4\\\"/\\u003e\\n \\u003cpath fill-rule=\\\"evenodd\\\" + clip-rule=\\\"evenodd\\\" d=\\\"M26.3109 43.8555C30.9659 43.8555 34.8687 42.3195 + 37.7213 39.6863L32.1258 35.3447C30.5898 36.3792 28.6306 37.0061 26.3109 37.0061C21.8282 + 37.0061 18.0195 33.9811 16.6559 29.906H10.9194V34.3573C13.7563 39.9841 19.5712 + 43.8555 26.3109 43.8555V43.8555Z\\\" fill=\\\"#34A853\\\"/\\u003e\\n \\u003cpath + fill-rule=\\\"evenodd\\\" clip-rule=\\\"evenodd\\\" d=\\\"M16.6559 29.8904C16.3111 + 28.8559 16.1074 27.7588 16.1074 26.6146C16.1074 25.4704 16.3111 24.3733 16.6559 + 23.3388V18.8875H10.9194C9.74388 21.2072 9.06992 23.8247 9.06992 26.6146C9.06992 + 29.4045 9.74388 32.022 10.9194 34.3417L15.3864 30.8621L16.6559 29.8904V29.8904Z\\\" + fill=\\\"#FBBC05\\\"/\\u003e\\n \\u003cpath fill-rule=\\\"evenodd\\\" + clip-rule=\\\"evenodd\\\" d=\\\"M26.3109 16.2386C28.85 16.2386 31.107 17.1164 + 32.9095 18.8091L37.8466 13.8719C34.853 11.082 30.9659 9.3736 26.3109 9.3736C19.5712 + 9.3736 13.7563 13.245 10.9194 18.8875L16.6559 23.3388C18.0195 19.2636 21.8282 + 16.2386 26.3109 16.2386V16.2386Z\\\" fill=\\\"#EA4335\\\"/\\u003e\\n \\u003c/svg\\u003e\\n + \ \\u003csvg class=\\\"logo-dark\\\" width=\\\"18\\\" height=\\\"18\\\" + viewBox=\\\"0 0 48 48\\\" xmlns=\\\"http://www.w3.org/2000/svg\\\"\\u003e\\n + \ \\u003ccircle cx=\\\"24\\\" cy=\\\"23\\\" fill=\\\"#FFF\\\" r=\\\"22\\\"/\\u003e\\n + \ \\u003cpath d=\\\"M33.76 34.26c2.75-2.56 4.49-6.37 4.49-11.26 0-.89-.08-1.84-.29-3H24.01v5.99h8.03c-.4 + 2.02-1.5 3.56-3.07 4.56v.75l3.91 2.97h.88z\\\" fill=\\\"#4285F4\\\"/\\u003e\\n + \ \\u003cpath d=\\\"M15.58 25.77A8.845 8.845 0 0 0 24 31.86c1.92 0 3.62-.46 + 4.97-1.31l4.79 3.71C31.14 36.7 27.65 38 24 38c-5.93 0-11.01-3.4-13.45-8.36l.17-1.01 + 4.06-2.85h.8z\\\" fill=\\\"#34A853\\\"/\\u003e\\n \\u003cpath d=\\\"M15.59 + 20.21a8.864 8.864 0 0 0 0 5.58l-5.03 3.86c-.98-2-1.53-4.25-1.53-6.64 0-2.39.55-4.64 + 1.53-6.64l1-.22 3.81 2.98.22 1.08z\\\" fill=\\\"#FBBC05\\\"/\\u003e\\n \\u003cpath + d=\\\"M24 14.14c2.11 0 4.02.75 5.52 1.98l4.36-4.36C31.22 9.43 27.81 8 24 8c-5.93 + 0-11.01 3.4-13.45 8.36l5.03 3.85A8.86 8.86 0 0 1 24 14.14z\\\" fill=\\\"#EA4335\\\"/\\u003e\\n + \ \\u003c/svg\\u003e\\n \\u003cdiv class=\\\"gradient-container\\\"\\u003e\\u003cdiv + class=\\\"gradient\\\"\\u003e\\u003c/div\\u003e\\u003c/div\\u003e\\n \\u003c/div\\u003e\\n + \ \\u003cdiv class=\\\"carousel\\\"\\u003e\\n \\u003ca class=\\\"chip\\\" + href=\\\"https://vertexaisearch.cloud.google.com/grounding-api-redirect/AUZIYQEyAtM7AiI-CeckMls3bGFtqzFIOnWGYIU9ym_iDr_pDA_Ez7MankcyamycqgFJQumfImXm1CRM2UPKsHPHIht16fbIQGEgSic-skhj9Dvae78NNk9m2QLzmIFeG7XAzQMgRewbrgZaQ1dMjPDLkFwEmW9XrtLrBC_behQhjte1djV_jU__ooOAHLN9J07GTKxgwYDS67DaHEi5-mCwY2dPFuGG\\\"\\u003ecurrent + population of Tokyo Japan\\u003c/a\\u003e\\n \\u003c/div\\u003e\\n\\u003c/div\\u003e\\n\"\n + \ },\n \"groundingChunks\": [\n {\n \"web\": + {\n \"uri\": \"https://vertexaisearch.cloud.google.com/grounding-api-redirect/AUZIYQHH0P3wtfpg-SgANWeTlB4QM0CjV1wVsQq3Zej435oFAi10fr8kPEP4epaq6SDdd5tMp0-Zf9ecLGZpqoAgEQdad7SbQWKVTscNIaaz_C2anJt-sL1l2ttjV5wo8ZI2fn5ZZL3m8jpk3T8Ishchw3O-\",\n + \ \"title\": \"tokyo.lg.jp\"\n }\n },\n {\n + \ \"web\": {\n \"uri\": \"https://vertexaisearch.cloud.google.com/grounding-api-redirect/AUZIYQFHIZptFH4erONxyTvzxa8bEnBAEOskMWR11FqwCqK9vz9daQzVb94V-Eh6GBAublbi41Bk2V_SqZ3sygCepFOOiNH8T5WS5CFV0PiJ8726pY5kMPG4AUIE6uMF7L6dY4M425bjXN3i-ZuZsaFa\",\n + \ \"title\": \"wikipedia.org\"\n }\n }\n ],\n + \ \"groundingSupports\": [\n {\n \"segment\": {\n + \ \"endIndex\": 101,\n \"text\": \"As of 2026, the + estimated total population of Tokyo Metropolis is approximately 14.27 million + people.\"\n },\n \"groundingChunkIndices\": [\n 0,\n + \ 1\n ]\n }\n ],\n \"webSearchQueries\": + [\n \"current population of Tokyo Japan\"\n ]\n }\n }\n + \ ],\n \"usageMetadata\": {\n \"promptTokenCount\": 48,\n \"candidatesTokenCount\": + 42,\n \"totalTokenCount\": 671,\n \"promptTokensDetails\": [\n {\n + \ \"modality\": \"TEXT\",\n \"tokenCount\": 48\n }\n ],\n + \ \"toolUsePromptTokenCount\": 93,\n \"toolUsePromptTokensDetails\": + [\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 93\n + \ }\n ],\n \"thoughtsTokenCount\": 488,\n \"serviceTier\": \"standard\"\n + \ },\n \"modelVersion\": \"gemini-2.5-flash\",\n \"responseId\": \"nQ9yaujeL_-QjrEPhqmhgQI\"\n}\n" + headers: + alt-svc: + - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 + content-length: + - '7241' + content-type: + - application/json; charset=UTF-8 + date: + - Tue, 04 Aug 2026 16:13:21 GMT + server: + - scaffolding on HTTPServer2 + server-timing: + - gfet4t7; dur=4002 + 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/adk/test_adk.py b/py/src/braintrust/integrations/adk/test_adk.py index 263d1c331..84d57955b 100644 --- a/py/src/braintrust/integrations/adk/test_adk.py +++ b/py/src/braintrust/integrations/adk/test_adk.py @@ -15,6 +15,7 @@ from google.adk.agents import LlmAgent, ParallelAgent, SequentialAgent from google.adk.runners import Runner from google.adk.sessions import InMemorySessionService +from google.adk.tools import google_search from google.genai import types from pydantic import BaseModel, Field @@ -619,6 +620,54 @@ async def test_adk_binary_data_attachment_conversion(memory_logger): assert "89504e47" not in llm_str.lower(), "Raw binary data (hex) should not be in LLM span input" +@pytest.mark.vcr +@pytest.mark.asyncio +async def test_adk_usage_metadata_metrics(memory_logger): + """Google Search usage includes tool prompts and reasoning in normalized totals.""" + assert not memory_logger.pop() + + agent = Agent( + name="usage_metadata_agent", + model="gemini-2.5-flash", + instruction="Use Google Search to answer the user's question accurately and concisely.", + tools=[google_search], + ) + app_name = "usage_metadata_app" + user_id = "test-user" + session_id = "test-session-usage-metadata" + runner = await _create_runner(agent, app_name=app_name, user_id=user_id, session_id=session_id) + user_msg = types.Content( + role="user", + parts=[types.Part(text="What is the current population of Tokyo, Japan? Answer in one sentence.")], + ) + + events = [event async for event in runner.run_async(user_id=user_id, session_id=session_id, new_message=user_msg)] + usage_metadata = next( + (event.usage_metadata for event in reversed(events) if getattr(event, "usage_metadata", None) is not None), + None, + ) + assert usage_metadata is not None + assert usage_metadata.tool_use_prompt_token_count > 0 + assert usage_metadata.thoughts_token_count > 0 + + spans = memory_logger.pop() + llm_spans = [row for row in spans if row["span_attributes"].get("type") == "llm"] + assert len(llm_spans) == 1 + llm_span = llm_spans[0] + metrics = llm_span["metrics"] + + assert metrics["prompt_tokens"] == (usage_metadata.prompt_token_count + usage_metadata.tool_use_prompt_token_count) + assert metrics["completion_tokens"] == ( + usage_metadata.candidates_token_count + usage_metadata.thoughts_token_count + ) + assert metrics["completion_reasoning_tokens"] == usage_metadata.thoughts_token_count + assert metrics["tokens"] == usage_metadata.total_token_count + assert metrics["tokens"] == metrics["prompt_tokens"] + metrics["completion_tokens"] + assert llm_span["metadata"]["usage_by_modality"]["tool_use_prompt_tokens_details"] == [ + detail.model_dump(exclude_none=True) for detail in usage_metadata.tool_use_prompt_tokens_details + ] + + @pytest.mark.vcr @pytest.mark.asyncio async def test_adk_captures_metrics(memory_logger): diff --git a/py/src/braintrust/integrations/adk/tracing.py b/py/src/braintrust/integrations/adk/tracing.py index 9460ff791..e5e18872b 100644 --- a/py/src/braintrust/integrations/adk/tracing.py +++ b/py/src/braintrust/integrations/adk/tracing.py @@ -11,7 +11,11 @@ from typing import Any from braintrust.bt_json import bt_safe_deep_copy -from braintrust.integrations.utils import _materialize_attachment +from braintrust.integrations.utils import ( + _extract_google_usage_metadata_metrics, + _extract_google_usage_metadata_provider_metadata, + _materialize_attachment, +) from braintrust.logger import start_span as _bt_start_span @@ -172,7 +176,7 @@ def _capture_config(config: Any) -> dict[str, Any] | Any: return captured or config -def _extract_metrics(response: Any) -> dict[str, float] | None: +def _extract_metrics(response: Any) -> dict[str, Any] | None: """Extract token usage metrics from Google GenAI response.""" if not response: return None @@ -181,26 +185,7 @@ def _extract_metrics(response: Any) -> dict[str, float] | None: if not usage_metadata: return None - metrics: dict[str, float] = {} - - # Core token counts - if hasattr(usage_metadata, "prompt_token_count") and usage_metadata.prompt_token_count is not None: - metrics["prompt_tokens"] = float(usage_metadata.prompt_token_count) - - if hasattr(usage_metadata, "candidates_token_count") and usage_metadata.candidates_token_count is not None: - metrics["completion_tokens"] = float(usage_metadata.candidates_token_count) - - if hasattr(usage_metadata, "total_token_count") and usage_metadata.total_token_count is not None: - metrics["tokens"] = float(usage_metadata.total_token_count) - - # Cached token metrics - if hasattr(usage_metadata, "cached_content_token_count") and usage_metadata.cached_content_token_count is not None: - metrics["prompt_cached_tokens"] = float(usage_metadata.cached_content_token_count) - - # Reasoning token metrics (thoughts_token_count) - if hasattr(usage_metadata, "thoughts_token_count") and usage_metadata.thoughts_token_count is not None: - metrics["completion_reasoning_tokens"] = float(usage_metadata.thoughts_token_count) - + metrics = _extract_google_usage_metadata_metrics(usage_metadata) return metrics if metrics else None @@ -474,8 +459,14 @@ async def _trace(): span_attributes={"llm_call_type": call_type}, ) - # Log output and metrics (span.log will handle serialization) - llm_span.log(output=output, metrics=metrics) + # Log output, metrics, and provider-specific usage details. + llm_span.log( + output=output, + metrics=metrics, + metadata=_extract_google_usage_metadata_provider_metadata( + getattr(last_event, "usage_metadata", None) + ), + ) async with aclosing(_trace()) as agen: async for event in agen: diff --git a/py/src/braintrust/integrations/google_genai/tracing.py b/py/src/braintrust/integrations/google_genai/tracing.py index 54349e52a..c29fab40d 100644 --- a/py/src/braintrust/integrations/google_genai/tracing.py +++ b/py/src/braintrust/integrations/google_genai/tracing.py @@ -8,7 +8,11 @@ from collections.abc import Awaitable, Callable, Iterable from typing import TYPE_CHECKING, Any, TypeAlias -from braintrust.integrations.utils import _materialize_attachment +from braintrust.integrations.utils import ( + _extract_google_usage_metadata_metrics, + _extract_google_usage_metadata_provider_metadata, + _materialize_attachment, +) from braintrust.logger import ( NOOP_SPAN, _state, @@ -331,61 +335,11 @@ 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: - 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 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 + metadata = _extract_google_usage_metadata_provider_metadata(usage_metadata) + return _materialize_interaction_value(metadata) if metadata is not None else None def _extract_generate_content_metrics(response: "GenerateContentResponse", start: float) -> dict[str, Any]: @@ -397,7 +351,7 @@ def _extract_generate_content_metrics(response: "GenerateContentResponse", start ) if hasattr(response, "usage_metadata") and response.usage_metadata: - _extract_usage_metadata_metrics(response.usage_metadata, metrics) + metrics.update(_extract_google_usage_metadata_metrics(response.usage_metadata)) return clean_nones(dict(metrics)) @@ -756,7 +710,7 @@ def _aggregate_generate_content_chunks( if usage_metadata: aggregated["usage_metadata"] = usage_metadata - _extract_usage_metadata_metrics(usage_metadata, metrics) + metrics.update(_extract_google_usage_metadata_metrics(usage_metadata)) if text: aggregated["text"] = text diff --git a/py/src/braintrust/integrations/test_utils.py b/py/src/braintrust/integrations/test_utils.py index 51db7af3f..1ee1c9bc2 100644 --- a/py/src/braintrust/integrations/test_utils.py +++ b/py/src/braintrust/integrations/test_utils.py @@ -15,6 +15,8 @@ _attachment_filename_for_mime_type, _camel_to_snake, _extract_audio_output, + _extract_google_usage_metadata_metrics, + _extract_google_usage_metadata_provider_metadata, _infer_audio_mime_type, _is_supported_metric_value, _log_and_end_span, @@ -231,6 +233,16 @@ def test_try_to_dict_returns_original_when_no_conversion_is_possible(): assert result is obj +def test_extract_google_usage_metadata_preserves_zero_and_omits_unreported_values(): + class UsageMetadata: + prompt_token_count = 0 + + usage_metadata = UsageMetadata() + + assert _extract_google_usage_metadata_metrics(usage_metadata) == {"prompt_tokens": 0} + assert _extract_google_usage_metadata_provider_metadata(usage_metadata) is None + + def test_parse_openai_usage_metrics_handles_nested_token_details(): usage = { "prompt_tokens": 10, diff --git a/py/src/braintrust/integrations/utils.py b/py/src/braintrust/integrations/utils.py index d4257df65..723016d50 100644 --- a/py/src/braintrust/integrations/utils.py +++ b/py/src/braintrust/integrations/utils.py @@ -529,6 +529,74 @@ def _prettify_response_params(params: dict[str, Any], *, drop_not_given: bool = return ret +def _extract_google_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_google_usage_metadata_metrics(usage_metadata: Any) -> dict[str, Any]: + """Normalize Google GenerateContent usage metadata into Braintrust metrics.""" + if usage_metadata is None: + return {} + + metrics: dict[str, Any] = {} + 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) + + total_token_count = getattr(usage_metadata, "total_token_count", None) + if total_token_count is not None: + metrics["tokens"] = total_token_count + + cached_content_token_count = getattr(usage_metadata, "cached_content_token_count", None) + if cached_content_token_count is not None: + metrics["prompt_cached_tokens"] = cached_content_token_count + + if thoughts_token_count is not None: + metrics["completion_reasoning_tokens"] = thoughts_token_count + + prompt_audio_tokens = _extract_google_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_google_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_google_modality_token_count(candidates_tokens_details, "IMAGE") + if completion_image_tokens is not None: + metrics["completion_image_tokens"] = completion_image_tokens + + return metrics + + +def _extract_google_usage_metadata_provider_metadata(usage_metadata: Any) -> dict[str, Any] | None: + """Preserve Google usage breakdowns that do not map to standard metrics.""" + if usage_metadata is None: + return 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] = details + return {"usage_by_modality": usage_by_modality} if usage_by_modality else None + + def _parse_openai_usage_metrics( usage: Any, *,