diff --git a/py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py b/py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py index a4015d6a..5051d60a 100644 --- a/py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py +++ b/py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py @@ -5,6 +5,7 @@ import inspect import os import time +from dataclasses import replace import pydantic_ai import pytest @@ -15,7 +16,7 @@ from packaging.version import Version from pydantic import BaseModel from pydantic_ai import Agent, ModelSettings -from pydantic_ai.messages import ModelRequest, UserPromptPart +from pydantic_ai.messages import ModelRequest, ModelResponse, UserPromptPart from pydantic_ai.usage import UsageLimits @@ -150,8 +151,6 @@ def test_extract_response_metrics_leaf_fields(): went away are read back. `RequestUsage.details` is `dict[str, int]`, so attribute-style lookups on it (the prior implementation) silently no-op. """ - from types import SimpleNamespace - from braintrust.integrations.pydantic_ai.tracing import _extract_response_metrics from pydantic_ai.usage import RequestUsage @@ -164,7 +163,10 @@ def test_extract_response_metrics_leaf_fields(): output_audio_tokens=11, details={"reasoning_tokens": 13, "cached_tokens": 17}, ) - response = SimpleNamespace(usage=usage) + # Exercise the legacy details fallback when normalized reasoning usage is unavailable. + if hasattr(usage, "output_reasoning_tokens"): + setattr(usage, "output_reasoning_tokens", None) + response = ModelResponse(parts=[], usage=usage) maybe_metrics = _extract_response_metrics(response, start_time=1.0, end_time=2.0) assert maybe_metrics is not None @@ -1822,16 +1824,13 @@ def calc(a: int, b: int) -> int: async def test_agent_run_anthropic_reasoning_tokens(memory_logger): """Real Anthropic extended-thinking response through pydantic-ai's own usage mapping. - pydantic-ai stashes the reasoning count under ``details["thinking_tokens"]`` for - Anthropic (vs ``reasoning_tokens`` for OpenAI); the extractor must surface it as - ``completion_reasoning_tokens``. This exercises the real provider mapping via a - checked-in cassette; ``test_reasoning_tokens_extraction_provider_keys`` is the - supplemental synthetic coverage. + Verify the normalized usage field works even without the legacy details keys. """ if os.environ.get("BRAINTRUST_TEST_PACKAGE_VERSION") != "latest": pytest.skip("Anthropic extended-thinking usage requires the latest pydantic-ai cassette") import anthropic + from braintrust.integrations.pydantic_ai.tracing import _extract_response_metrics from pydantic_ai.models.anthropic import AnthropicModel, AnthropicModelSettings from pydantic_ai.providers.anthropic import AnthropicProvider @@ -1856,8 +1855,27 @@ async def test_agent_run_anthropic_reasoning_tokens(memory_logger): chat_span = next((s for s in spans if "chat" in s["span_attributes"]["name"]), None) assert chat_span is not None, "chat span not found" assert chat_span["metadata"]["provider"] == "anthropic" + + response = result.all_messages()[-1] + assert isinstance(response, ModelResponse) + reasoning_tokens = response.usage.output_reasoning_tokens + assert reasoning_tokens > 0 + # Reasoning is a subset of completion usage, not an additional token class. # pylint: disable=unsupported-membership-test,unsubscriptable-object - assert chat_span["metrics"]["completion_reasoning_tokens"] > 0 + assert chat_span["metrics"]["completion_reasoning_tokens"] == reasoning_tokens + assert reasoning_tokens <= chat_span["metrics"]["completion_tokens"] + assert chat_span["metrics"]["completion_tokens"] == response.usage.output_tokens + for span in spans: + if span["span_attributes"]["type"] == SpanTypeAttribute.TASK: + assert "completion_reasoning_tokens" not in span["metrics"] + # pylint: enable=unsupported-membership-test,unsubscriptable-object + + # Reuse the real provider usage without its legacy, provider-specific details. + normalized_response = replace(response, usage=replace(response.usage, details={})) + metrics = _extract_response_metrics(normalized_response, start_time=1.0, end_time=2.0) + assert metrics is not None + # pylint: disable=unsupported-membership-test,unsubscriptable-object + assert metrics["completion_reasoning_tokens"] == reasoning_tokens # pylint: enable=unsupported-membership-test,unsubscriptable-object @@ -1884,32 +1902,31 @@ async def test_agent_run_openai_web_search_usage(memory_logger): @pytest.mark.parametrize( - "details_key", + "details_key, reasoning_tokens, expected", [ - "thinking_tokens", # Anthropic - "thoughts_tokens", # Google + pytest.param("thinking_tokens", None, 128, id="anthropic-fallback"), + pytest.param("thoughts_tokens", None, 128, id="google-fallback"), + pytest.param("reasoning_tokens", 64, 64, id="normalized-precedence"), + pytest.param("reasoning_tokens", 0, 0, id="normalized-zero"), ], ) -def test_reasoning_tokens_extraction_provider_keys(details_key): - """pydantic_ai stashes the reasoning-token count under a provider-specific key: - OpenAI "reasoning_tokens", Anthropic "thinking_tokens", Google "thoughts_tokens". - All three must surface as `completion_reasoning_tokens` (previously only OpenAI's - key was read, silently dropping Anthropic/Google reasoning). OpenAI's key is - covered by `test_extract_response_metrics_leaf_fields`. - """ - from types import SimpleNamespace - +def test_reasoning_tokens_extraction(details_key, reasoning_tokens, expected): + """Prefer normalized usage, including zero, with legacy provider-key fallbacks.""" from braintrust.integrations.pydantic_ai.tracing import _extract_response_metrics from pydantic_ai.usage import RequestUsage - usage = RequestUsage(input_tokens=10, output_tokens=20, details={details_key: 128}) - response = SimpleNamespace(parts=[], usage=usage) + usage = RequestUsage(input_tokens=10, output_tokens=200, details={details_key: 128}) + if reasoning_tokens is not None and not hasattr(usage, "output_reasoning_tokens"): + pytest.skip("Normalized reasoning tokens require newer pydantic-ai") + if hasattr(usage, "output_reasoning_tokens"): + setattr(usage, "output_reasoning_tokens", reasoning_tokens) + response = ModelResponse(parts=[], usage=usage) metrics = _extract_response_metrics(response, start_time=1.0, end_time=2.0) assert metrics is not None # pylint: disable=unsupported-membership-test,unsubscriptable-object - assert metrics["completion_reasoning_tokens"] == 128.0 + assert metrics["completion_reasoning_tokens"] == expected # pylint: enable=unsupported-membership-test,unsubscriptable-object diff --git a/py/src/braintrust/integrations/pydantic_ai/tracing.py b/py/src/braintrust/integrations/pydantic_ai/tracing.py index 9eca9918..d03c3378 100644 --- a/py/src/braintrust/integrations/pydantic_ai/tracing.py +++ b/py/src/braintrust/integrations/pydantic_ai/tracing.py @@ -1421,26 +1421,24 @@ def _extract_response_metrics( if hasattr(usage, "output_audio_tokens") and usage.output_audio_tokens is not None: metrics["completion_audio_tokens"] = float(usage.output_audio_tokens) - # RequestUsage.details is a dict; providers stash reasoning_tokens/cached_tokens here. - # The reasoning-token count lives under a provider-specific key: OpenAI uses - # "reasoning_tokens", Anthropic "thinking_tokens", Google "thoughts_tokens". All three - # are the same normalized quantity (a subset of the output tokens), so read whichever - # is present. + # Prefer normalized reasoning usage; older versions use provider-specific details. + reasoning = getattr(usage, "output_reasoning_tokens", None) details = getattr(usage, "details", None) if isinstance(details, dict): - reasoning = next( - ( - details[key] - for key in ("reasoning_tokens", "thinking_tokens", "thoughts_tokens") - if details.get(key) is not None - ), - None, - ) - if reasoning is not None: - metrics["completion_reasoning_tokens"] = float(reasoning) + if reasoning is None: + reasoning = next( + ( + details[key] + for key in ("reasoning_tokens", "thinking_tokens", "thoughts_tokens") + if details.get(key) is not None + ), + None, + ) cached = details.get("cached_tokens") if cached is not None: metrics["prompt_cached_tokens"] = float(cached) + if reasoning is not None: + metrics["completion_reasoning_tokens"] = float(reasoning) return metrics if metrics else None