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"):