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 02f6197e..03b2cf47 100644 --- a/py/src/braintrust/integrations/google_genai/test_google_genai.py +++ b/py/src/braintrust/integrations/google_genai/test_google_genai.py @@ -1656,6 +1656,8 @@ def test_interactions_create_and_get(memory_logger): response.usage.total_output_tokens + response.usage.total_thought_tokens ) assert create_span["metrics"]["tokens"] == response.usage.total_tokens + assert "tool_use_tokens" not in create_span["metrics"] + assert create_span["metadata"]["total_tool_use_tokens"] == response.usage.total_tool_use_tokens assert get_span["input"]["id"] == response.id assert get_span["metadata"]["interaction_id"] == response.id diff --git a/py/src/braintrust/integrations/google_genai/tracing.py b/py/src/braintrust/integrations/google_genai/tracing.py index 82ce43f1..54349e52 100644 --- a/py/src/braintrust/integrations/google_genai/tracing.py +++ b/py/src/braintrust/integrations/google_genai/tracing.py @@ -527,8 +527,6 @@ def _extract_interaction_usage_metrics(usage: Any, metrics: dict[str, Any]) -> N metrics["prompt_cached_tokens"] = usage.total_cached_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 def _extract_interaction_text(outputs: list[dict[str, Any]]) -> str | None: @@ -599,6 +597,9 @@ def _extract_interaction_metadata(response: "Interaction") -> dict[str, Any]: "role": getattr(response, "role", None), "response_mime_type": getattr(response, "response_mime_type", None), "response_modalities": _materialize_interaction_value(getattr(response, "response_modalities", None)), + "total_tool_use_tokens": ( + usage_serialized.get("total_tool_use_tokens") if isinstance(usage_serialized, dict) else None + ), "usage_by_modality": usage_by_modality, } )