From cb7b217db0bcc57516b7a591c4d5debb37f999ed Mon Sep 17 00:00:00 2001 From: ekeith <55766816+evanmkeith@users.noreply.github.com> Date: Mon, 27 Jul 2026 14:19:07 -0700 Subject: [PATCH] ix(google_genai): include thinking tokens in completion usage ## Summary - Fold Gemini `thoughts_token_count` into Google GenAI `completion_tokens` - Preserve `completion_reasoning_tokens` as the separate reasoning detail metric - Tighten the reasoning regression test to assert span metrics match provider usage metadata ## Testing - `nox -s "test_google_genai(latest)" -- -k test_reasoning` - `nox -s "test_google_genai(latest)"` - `nox -s "test_google_genai(1.75.0)" -- -k test_reasoning` - `nox -s "test_google_genai(1.30.0)" -- -k test_reasoning` - `git diff --check` --- .../integrations/google_genai/test_google_genai.py | 6 +++++- py/src/braintrust/integrations/google_genai/tracing.py | 7 ++++++- 2 files changed, 11 insertions(+), 2 deletions(-) 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..b36cf421 100644 --- a/py/src/braintrust/integrations/google_genai/test_google_genai.py +++ b/py/src/braintrust/integrations/google_genai/test_google_genai.py @@ -1030,10 +1030,14 @@ def test_reasoning(memory_logger): assert len(spans) == 2 first_span, second_span = spans - for span in spans: + for span, response in ((first_span, first_response), (second_span, follow_up_response)): assert span["metadata"]["model"] == REASONING_MODEL assert span["input"]["config"]["thinking_config"]["include_thoughts"] is True assert span["metrics"]["completion_reasoning_tokens"] > 0 + assert response.usage_metadata + assert span["metrics"]["completion_tokens"] == ( + response.usage_metadata.candidates_token_count + response.usage_metadata.thoughts_token_count + ) assert span["output"] assert first_prompt in str(first_span["input"]) diff --git a/py/src/braintrust/integrations/google_genai/tracing.py b/py/src/braintrust/integrations/google_genai/tracing.py index f0c07bfb..e9172576 100644 --- a/py/src/braintrust/integrations/google_genai/tracing.py +++ b/py/src/braintrust/integrations/google_genai/tracing.py @@ -335,8 +335,13 @@ def _extract_usage_metadata_metrics( ) -> None: if hasattr(usage_metadata, "prompt_token_count"): metrics["prompt_tokens"] = usage_metadata.prompt_token_count + thoughts_token_count = getattr(usage_metadata, "thoughts_token_count", None) if hasattr(usage_metadata, "candidates_token_count"): - metrics["completion_tokens"] = usage_metadata.candidates_token_count + candidates_token_count = usage_metadata.candidates_token_count + if candidates_token_count is None and thoughts_token_count is None: + metrics["completion_tokens"] = None + else: + 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"):