fix(google_genai): include thinking tokens in completion usage

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

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