Skip to content

fix(pydantic_ai): strip unnecessary wrapper tokens#321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Abhijeet Prasad (AbhiPrasad) merged 5 commits into from
Apr 20, 2026
Merged

fix(pydantic_ai): strip unnecessary wrapper tokens #321

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
6a6cb76
fix(pydantic_ai): strip token metrics from wrapper spans
colinbennettbrain Apr 16, 2026
61be582
fix(pydantic_ai): restore audio + dict-style usage details on leaf span
colinbennettbrain Apr 17, 2026
171897c
Merge branch 'main' into colin/pydantic-strip-wrapper-tokens
colinbennettbrain Apr 17, 2026
246fcc2
fix(pydantic-ai): narrow Optional metrics type for pylint
Apr 17, 2026
39c970a
fix(py): fix pylint in pydantic ai metrics test
colinbennettbrain Apr 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
173 changes: 116 additions & 57 deletions py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -67,6 +67,46 @@ def _assert_metrics_are_valid(metrics, start, end):
assert metrics["completion_tokens"] > 0


def test_extract_response_metrics_leaf_fields():
"""`_extract_response_metrics` is the only path that populates token metrics on
the leaf chat span post-#312/strip-wrapper-tokens. Verify the audio-token and
dict-style `details` fields that were dropped when `_extract_usage_metrics`
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

usage = RequestUsage(
input_tokens=100,
output_tokens=50,
cache_read_tokens=10,
cache_write_tokens=20,
input_audio_tokens=7,
output_audio_tokens=11,
details={"reasoning_tokens": 13, "cached_tokens": 17},
)
response = SimpleNamespace(usage=usage)

maybe_metrics = _extract_response_metrics(response, start_time=1.0, end_time=2.0)
assert maybe_metrics is not None

# pylint: disable=unsupported-membership-test,unsubscriptable-object
assert maybe_metrics["prompt_tokens"] == 100
assert maybe_metrics["completion_tokens"] == 50
assert maybe_metrics["tokens"] == 150 # total_tokens is a property: input + output
assert maybe_metrics["prompt_cache_creation_tokens"] == 20
assert maybe_metrics["prompt_audio_tokens"] == 7
assert maybe_metrics["completion_audio_tokens"] == 11
assert maybe_metrics["completion_reasoning_tokens"] == 13
# details["cached_tokens"] is read after cache_read_tokens, so it wins when both
# are populated; this preserves the pre-strip behavior of `_extract_usage_metrics`.
assert maybe_metrics["prompt_cached_tokens"] == 17
# pylint: enable=unsupported-membership-test,unsubscriptable-object


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_direct_model_request_creates_nested_chat_span_without_class_scan(memory_logger, direct):
Expand DownExpand Up@@ -140,23 +180,21 @@ async def test_agent_run_async(memory_logger):
assert chat_span["metadata"]["provider"] == "openai"
_assert_metrics_are_valid(chat_span["metrics"], start, end)

# Agent spans should have token metrics
assert "prompt_tokens" in agent_span["metrics"]
assert "completion_tokens" in agent_span["metrics"]
assert agent_span["metrics"]["prompt_tokens"] > 0
assert agent_span["metrics"]["completion_tokens"] > 0
# Regression: wrapper agent_run span must NOT log token metrics. The leaf chat
# span already logs them, and trace-tree rollup (self + descendants) plus any
# unfiltered sum over metrics would otherwise double-count tokens/cost at the
# parent regardless of span type.
for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
assert token_key not in agent_span["metrics"], (
f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
)

# Regression: no double-counting of cost/tokens. Experiment-level aggregations
# sum metrics across type='llm' spans, so a single agent turn must contribute
# its tokens exactly once. The wrapper agent_run span logs the same usage as
# the leaf chat span; only the leaf should be type=LLM.
# Only the leaf chat span should be type=LLM, and it must own the token totals.
llm_spans = [s for s in spans if s["span_attributes"]["type"] == SpanTypeAttribute.LLM]
assert len(llm_spans) == 1, f"expected exactly one LLM-typed span, got {len(llm_spans)}"
assert llm_spans[0]["span_id"] == chat_span["span_id"]
llm_prompt_tokens_sum = sum(s["metrics"].get("prompt_tokens", 0) for s in llm_spans)
llm_completion_tokens_sum = sum(s["metrics"].get("completion_tokens", 0) for s in llm_spans)
assert llm_prompt_tokens_sum == chat_span["metrics"]["prompt_tokens"]
assert llm_completion_tokens_sum == chat_span["metrics"]["completion_tokens"]
assert chat_span["metrics"]["prompt_tokens"] > 0
assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -245,9 +283,12 @@ def is_descendant(child_span, ancestor_id):
assert chat_span["metadata"]["provider"] == "openai"
_assert_metrics_are_valid(chat_span["metrics"], start, end)

# Agent spans should have token metrics
assert "prompt_tokens" in agent_sync_span["metrics"]
assert "completion_tokens" in agent_sync_span["metrics"]
# Wrapper agent_run_sync span must not log token metrics (would double-count at rollup).
assert "prompt_tokens" not in agent_sync_span["metrics"]
assert "completion_tokens" not in agent_sync_span["metrics"]
# Tokens live on the leaf chat span only.
assert chat_span["metrics"]["prompt_tokens"] > 0
assert chat_span["metrics"]["completion_tokens"] > 0


def test_agent_to_cli_sync(memory_logger, monkeypatch):
Expand DownExpand Up@@ -544,9 +585,12 @@ async def test_agent_run_stream(memory_logger):
print(f"span_parents: {chat_span['span_parents']}")
print(f"metrics: {chat_span['metrics']}")

# Agent spans should have token metrics
assert "prompt_tokens" in agent_span["metrics"]
assert "completion_tokens" in agent_span["metrics"]
# Wrapper stream span must not log token metrics (would double-count at rollup).
# time_to_first_token is asserted above; it's a non-summable timing metric and stays.
assert "prompt_tokens" not in agent_span["metrics"]
assert "completion_tokens" not in agent_span["metrics"]
assert chat_span["metrics"]["prompt_tokens"] > 0
assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -741,6 +785,15 @@ async def test_direct_model_request_stream(memory_logger, direct):
assert ttft <= duration, f"time_to_first_token ({ttft}s) should be <= duration ({duration}s)"
assert ttft < 3.0, f"time_to_first_token should be < 3s for API call, got {ttft}s"

# Regression: model_request_stream wrapper span (TASK) must NOT log token metrics.
# `_DirectStreamWrapper` here is constructed with span_type=TASK, so __aexit__
# routes to `_wrapper_span_metrics`. The leaf chat span (when present) is the only
# place tokens should be recorded.
for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
assert token_key not in direct_span["metrics"], (
f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
)

print(f"✓ Direct stream time_to_first_token: {ttft}s (duration: {duration}s)")


Expand DownExpand Up@@ -842,9 +895,11 @@ def is_descendant(child_span, ancestor_id):
assert chat_span["metadata"]["provider"] == "openai"
_assert_metrics_are_valid(chat_span["metrics"], start, end)

# Agent spans should have token metrics
assert "prompt_tokens" in agent_span["metrics"]
assert "completion_tokens" in agent_span["metrics"]
# Wrapper agent_run span must not log token metrics (would double-count at rollup).
assert "prompt_tokens" not in agent_span["metrics"]
assert "completion_tokens" not in agent_span["metrics"]
assert chat_span["metrics"]["prompt_tokens"] > 0
assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -1143,9 +1198,9 @@ def is_descendant(child_span, ancestor_id):
# Chat span may not have complete metrics since it's an intermediate span
assert "start" in chat_span["metrics"]

# Agent spans should have token metrics
assert "prompt_tokens" in agent_span["metrics"]
assert "completion_tokens" in agent_span["metrics"]
# Wrapper agent_run_stream_sync span must not log token metrics.
assert "prompt_tokens" not in agent_span["metrics"]
assert "completion_tokens" not in agent_span["metrics"]


@pytest.mark.vcr
Expand DownExpand Up@@ -1183,6 +1238,14 @@ async def test_agent_run_stream_events(memory_logger):
assert agent_span["metrics"]["event_count"] == event_count
_assert_metrics_are_valid(agent_span["metrics"], start, end)

# Regression: wrapper agent_run_stream_events span must NOT log token metrics.
# Trace-tree rollup over non-scorer spans plus experiment-level token sums would
# otherwise double-count against the leaf chat span.
for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
assert token_key not in agent_span["metrics"], (
f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
)


@pytest.mark.vcr
def test_direct_model_request_stream_sync(memory_logger, direct):
Expand DownExpand Up@@ -1223,6 +1286,15 @@ def test_direct_model_request_stream_sync(memory_logger, direct):
assert ttft <= duration, f"time_to_first_token ({ttft}s) should be <= duration ({duration}s)"
assert ttft < 3.0, f"time_to_first_token should be < 3s for API call, got {ttft}s"

# Regression: model_request_stream_sync wrapper span must NOT log token metrics.
# `_DirectStreamWrapperSync.__exit__` always uses `_wrapper_span_metrics`; sync
# direct streaming has no leaf chat span, so this is the only span we emit and
# token totals would otherwise be derived twice once a leaf is added.
for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
assert token_key not in span["metrics"], (
f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
)

print(f"✓ Direct sync stream time_to_first_token: {ttft}s (duration: {duration}s)")


Expand DownExpand Up@@ -2030,53 +2102,40 @@ def helper_tool() -> str:
def test_reasoning_tokens_extraction(memory_logger):
"""Test that reasoning tokens are extracted from model responses.

For reasoning models like o1/o3, usage.details.reasoning_tokens should be
captured in the metrics field.
For reasoning models like o1/o3, providers (e.g. OpenAI) populate
`usage.details["reasoning_tokens"]` (a dict entry, not an attribute) and
`_extract_response_metrics` must surface it as `completion_reasoning_tokens`.
"""
assert not memory_logger.pop()

# Mock a response that has reasoning tokens
from unittest.mock import MagicMock
from types import SimpleNamespace

# Create a mock response with reasoning tokens
mock_response = MagicMock()
mock_response.parts = [
MagicMock(
part_kind="thinking",
content="Let me think about this...",
),
MagicMock(
part_kind="text",
content="The answer is 42",
),
]
mock_response.usage = MagicMock()
mock_response.usage.input_tokens = 10
mock_response.usage.output_tokens = 20
mock_response.usage.total_tokens = 30
mock_response.usage.cache_read_tokens = 0
mock_response.usage.cache_write_tokens = 0
mock_response.usage.details = MagicMock()
mock_response.usage.details.reasoning_tokens = 128

# Test the metric extraction function directly
from braintrust.integrations.pydantic_ai.tracing import _extract_response_metrics
from pydantic_ai.usage import RequestUsage

# Use a real RequestUsage so `details` is the actual dict[str, int] shape
# pydantic_ai produces. A MagicMock would falsely satisfy attribute-style
# lookups and hide real-world dict-only behavior.
usage = RequestUsage(
input_tokens=10,
output_tokens=20,
cache_read_tokens=0,
cache_write_tokens=0,
details={"reasoning_tokens": 128},
)
response = SimpleNamespace(parts=[], usage=usage)

start_time = time.time()
end_time = start_time + 5.0

metrics = _extract_response_metrics(mock_response, start_time, end_time)
metrics = _extract_response_metrics(response, start_time, end_time)

# Verify all metrics are present
assert metrics is not None, "Should extract metrics"
# pylint: disable=unsupported-membership-test,unsubscriptable-object
assert "prompt_tokens" in metrics, "Should have prompt_tokens"
assert metrics["prompt_tokens"] == 10.0
assert "completion_tokens" in metrics, "Should have completion_tokens"
assert metrics["completion_tokens"] == 20.0
assert "tokens" in metrics, "Should have total tokens"
assert metrics["tokens"] == 30.0
assert "completion_reasoning_tokens" in metrics, "Should have completion_reasoning_tokens"
assert metrics["tokens"] == 30.0 # total_tokens is a property: input + output
assert metrics["completion_reasoning_tokens"] == 128.0, (
f"Expected 128.0, got {metrics['completion_reasoning_tokens']}"
)
Expand Down
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { // Add copy buttons to all 
 blocks
(function() {
function addCopyButtons() {
document.querySelectorAll('pre code').forEach(function(codeBlock) {
if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;
codeBlock.parentElement.setAttribute('data-copy-added', 'true');
var btn = document.createElement('button');
btn.textContent = 'Copy';
btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';
btn.onmouseover = function() { this.style.opacity = '1'; };
btn.onmouseout = function() { this.style.opacity = '0.7'; };
btn.onclick = function() {
navigator.clipboard.writeText(codeBlock.textContent).then(function() {
btn.textContent = 'Copied!';
setTimeout(function() { btn.textContent = 'Copy'; }, 1500);
});
};
codeBlock.parentElement.style.position = 'relative';
codeBlock.parentElement.appendChild(btn);
});
}
addCopyButtons();
// Re-run on dynamic content
var observer = new MutationObserver(addCopyButtons);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
fix(pydantic_ai): strip unnecessary wrapper tokens by colinbennettbrain · Pull Request #321 · braintrustdata/braintrust-sdk-python · GitHub
Skip to content
fix(pydantic_ai): strip unnecessary wrapper tokens#321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

Abhijeet Prasad (AbhiPrasad)
merged 5 commits into

from

Apr 20, 2026
 Merged


fix(pydantic_ai): strip unnecessary wrapper tokens
#321
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
6a6cb76

fix(pydantic_ai): strip token metrics from wrapper spans


colinbennettbrain Apr 16, 2026
61be582

fix(pydantic_ai): restore audio + dict-style usage details on leaf span


colinbennettbrain Apr 17, 2026
171897c

Merge branch 'main' into colin/pydantic-strip-wrapper-tokens


colinbennettbrain Apr 17, 2026
246fcc2

fix(pydantic-ai): narrow Optional metrics type for pylint

Apr 17, 2026
39c970a

fix(py): fix pylint in pydantic ai metrics test


colinbennettbrain Apr 17, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


173 changes: 116 additions & 57 deletions
py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -67,6 +67,46 @@ def _assert_metrics_are_valid(metrics, start, end):
 assert metrics["completion_tokens"] > 0


def test_extract_response_metrics_leaf_fields():
 """`_extract_response_metrics` is the only path that populates token metrics on
 the leaf chat span post-#312/strip-wrapper-tokens. Verify the audio-token and
 dict-style `details` fields that were dropped when `_extract_usage_metrics`
 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

 usage = RequestUsage(
 input_tokens=100,
 output_tokens=50,
 cache_read_tokens=10,
 cache_write_tokens=20,
 input_audio_tokens=7,
 output_audio_tokens=11,
 details={"reasoning_tokens": 13, "cached_tokens": 17},
 )
 response = SimpleNamespace(usage=usage)

 maybe_metrics = _extract_response_metrics(response, start_time=1.0, end_time=2.0)
 assert maybe_metrics is not None

 # pylint: disable=unsupported-membership-test,unsubscriptable-object
 assert maybe_metrics["prompt_tokens"] == 100
 assert maybe_metrics["completion_tokens"] == 50
 assert maybe_metrics["tokens"] == 150 # total_tokens is a property: input + output
 assert maybe_metrics["prompt_cache_creation_tokens"] == 20
 assert maybe_metrics["prompt_audio_tokens"] == 7
 assert maybe_metrics["completion_audio_tokens"] == 11
 assert maybe_metrics["completion_reasoning_tokens"] == 13
 # details["cached_tokens"] is read after cache_read_tokens, so it wins when both
 # are populated; this preserves the pre-strip behavior of `_extract_usage_metrics`.
 assert maybe_metrics["prompt_cached_tokens"] == 17
 # pylint: enable=unsupported-membership-test,unsubscriptable-object


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_direct_model_request_creates_nested_chat_span_without_class_scan(memory_logger, direct):
Expand DownExpand Up@@ -140,23 +180,21 @@ async def test_agent_run_async(memory_logger):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 assert agent_span["metrics"]["prompt_tokens"] > 0
 assert agent_span["metrics"]["completion_tokens"] > 0
 # Regression: wrapper agent_run span must NOT log token metrics. The leaf chat
 # span already logs them, and trace-tree rollup (self + descendants) plus any
 # unfiltered sum over metrics would otherwise double-count tokens/cost at the
 # parent regardless of span type.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in agent_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 # Regression: no double-counting of cost/tokens. Experiment-level aggregations
 # sum metrics across type='llm' spans, so a single agent turn must contribute
 # its tokens exactly once. The wrapper agent_run span logs the same usage as
 # the leaf chat span; only the leaf should be type=LLM.
 # Only the leaf chat span should be type=LLM, and it must own the token totals.
 llm_spans = [s for s in spans if s["span_attributes"]["type"] == SpanTypeAttribute.LLM]
 assert len(llm_spans) == 1, f"expected exactly one LLM-typed span, got {len(llm_spans)}"
 assert llm_spans[0]["span_id"] == chat_span["span_id"]
 llm_prompt_tokens_sum = sum(s["metrics"].get("prompt_tokens", 0) for s in llm_spans)
 llm_completion_tokens_sum = sum(s["metrics"].get("completion_tokens", 0) for s in llm_spans)
 assert llm_prompt_tokens_sum == chat_span["metrics"]["prompt_tokens"]
 assert llm_completion_tokens_sum == chat_span["metrics"]["completion_tokens"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -245,9 +283,12 @@ def is_descendant(child_span, ancestor_id):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_sync_span["metrics"]
 assert "completion_tokens" in agent_sync_span["metrics"]
 # Wrapper agent_run_sync span must not log token metrics (would double-count at rollup).
 assert "prompt_tokens" not in agent_sync_span["metrics"]
 assert "completion_tokens" not in agent_sync_span["metrics"]
 # Tokens live on the leaf chat span only.
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


def test_agent_to_cli_sync(memory_logger, monkeypatch):
Expand DownExpand Up@@ -544,9 +585,12 @@ async def test_agent_run_stream(memory_logger):
 print(f"span_parents: {chat_span['span_parents']}")
 print(f"metrics: {chat_span['metrics']}")

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper stream span must not log token metrics (would double-count at rollup).
 # time_to_first_token is asserted above; it's a non-summable timing metric and stays.
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -741,6 +785,15 @@ async def test_direct_model_request_stream(memory_logger, direct):
 assert ttft <= duration, f"time_to_first_token ({ttft}s) should be <= duration ({duration}s)"
 assert ttft < 3.0, f"time_to_first_token should be < 3s for API call, got {ttft}s"

 # Regression: model_request_stream wrapper span (TASK) must NOT log token metrics.
 # `_DirectStreamWrapper` here is constructed with span_type=TASK, so __aexit__
 # routes to `_wrapper_span_metrics`. The leaf chat span (when present) is the only
 # place tokens should be recorded.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in direct_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 print(f"✓ Direct stream time_to_first_token: {ttft}s (duration: {duration}s)")


Expand DownExpand Up@@ -842,9 +895,11 @@ def is_descendant(child_span, ancestor_id):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper agent_run span must not log token metrics (would double-count at rollup).
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -1143,9 +1198,9 @@ def is_descendant(child_span, ancestor_id):
 # Chat span may not have complete metrics since it's an intermediate span
 assert "start" in chat_span["metrics"]

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper agent_run_stream_sync span must not log token metrics.
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]


@pytest.mark.vcr
Expand DownExpand Up@@ -1183,6 +1238,14 @@ async def test_agent_run_stream_events(memory_logger):
 assert agent_span["metrics"]["event_count"] == event_count
 _assert_metrics_are_valid(agent_span["metrics"], start, end)

 # Regression: wrapper agent_run_stream_events span must NOT log token metrics.
 # Trace-tree rollup over non-scorer spans plus experiment-level token sums would
 # otherwise double-count against the leaf chat span.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in agent_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )


@pytest.mark.vcr
def test_direct_model_request_stream_sync(memory_logger, direct):
Expand DownExpand Up@@ -1223,6 +1286,15 @@ def test_direct_model_request_stream_sync(memory_logger, direct):
 assert ttft <= duration, f"time_to_first_token ({ttft}s) should be <= duration ({duration}s)"
 assert ttft < 3.0, f"time_to_first_token should be < 3s for API call, got {ttft}s"

 # Regression: model_request_stream_sync wrapper span must NOT log token metrics.
 # `_DirectStreamWrapperSync.__exit__` always uses `_wrapper_span_metrics`; sync
 # direct streaming has no leaf chat span, so this is the only span we emit and
 # token totals would otherwise be derived twice once a leaf is added.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 print(f"✓ Direct sync stream time_to_first_token: {ttft}s (duration: {duration}s)")


Expand DownExpand Up@@ -2030,53 +2102,40 @@ def helper_tool() -> str:
def test_reasoning_tokens_extraction(memory_logger):
 """Test that reasoning tokens are extracted from model responses.

 For reasoning models like o1/o3, usage.details.reasoning_tokens should be
 captured in the metrics field.
 For reasoning models like o1/o3, providers (e.g. OpenAI) populate
 `usage.details["reasoning_tokens"]` (a dict entry, not an attribute) and
 `_extract_response_metrics` must surface it as `completion_reasoning_tokens`.
 """
 assert not memory_logger.pop()

 # Mock a response that has reasoning tokens
 from unittest.mock import MagicMock
 from types import SimpleNamespace

 # Create a mock response with reasoning tokens
 mock_response = MagicMock()
 mock_response.parts = [
 MagicMock(
 part_kind="thinking",
 content="Let me think about this...",
 ),
 MagicMock(
 part_kind="text",
 content="The answer is 42",
 ),
 ]
 mock_response.usage = MagicMock()
 mock_response.usage.input_tokens = 10
 mock_response.usage.output_tokens = 20
 mock_response.usage.total_tokens = 30
 mock_response.usage.cache_read_tokens = 0
 mock_response.usage.cache_write_tokens = 0
 mock_response.usage.details = MagicMock()
 mock_response.usage.details.reasoning_tokens = 128

 # Test the metric extraction function directly
 from braintrust.integrations.pydantic_ai.tracing import _extract_response_metrics
 from pydantic_ai.usage import RequestUsage

 # Use a real RequestUsage so `details` is the actual dict[str, int] shape
 # pydantic_ai produces. A MagicMock would falsely satisfy attribute-style
 # lookups and hide real-world dict-only behavior.
 usage = RequestUsage(
 input_tokens=10,
 output_tokens=20,
 cache_read_tokens=0,
 cache_write_tokens=0,
 details={"reasoning_tokens": 128},
 )
 response = SimpleNamespace(parts=[], usage=usage)

 start_time = time.time()
 end_time = start_time + 5.0

 metrics = _extract_response_metrics(mock_response, start_time, end_time)
 metrics = _extract_response_metrics(response, start_time, end_time)

 # Verify all metrics are present
 assert metrics is not None, "Should extract metrics"
 # pylint: disable=unsupported-membership-test,unsubscriptable-object
 assert "prompt_tokens" in metrics, "Should have prompt_tokens"
 assert metrics["prompt_tokens"] == 10.0
 assert "completion_tokens" in metrics, "Should have completion_tokens"
 assert metrics["completion_tokens"] == 20.0
 assert "tokens" in metrics, "Should have total tokens"
 assert metrics["tokens"] == 30.0
 assert "completion_reasoning_tokens" in metrics, "Should have completion_reasoning_tokens"
 assert metrics["tokens"] == 30.0 # total_tokens is a property: input + output
 assert metrics["completion_reasoning_tokens"] == 128.0, (
 f"Expected 128.0, got {metrics['completion_reasoning_tokens']}"
 )
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Force GitHub README to respect dark mode
(function() {
var style = document.createElement('style');
style.textContent = '
.markdown-body {
color-scheme: dark light;
}
.markdown-body pre { background: #161b22 !important; }
.markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }
.markdown-body table th, .markdown-body table td { border-color: #30363d !important; }
.markdown-body img { background: #0d1117; }
.markdown-body blockquote { border-left-color: #8b949e; }
.markdown-body hr { border-color: #30363d; }
';
document.head.appendChild(style);
})();
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
fix(pydantic_ai): strip unnecessary wrapper tokens by colinbennettbrain · Pull Request #321 · braintrustdata/braintrust-sdk-python · GitHub
Skip to content
fix(pydantic_ai): strip unnecessary wrapper tokens#321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

Abhijeet Prasad (AbhiPrasad)
merged 5 commits into

from

Apr 20, 2026
 Merged


fix(pydantic_ai): strip unnecessary wrapper tokens
#321
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
6a6cb76

fix(pydantic_ai): strip token metrics from wrapper spans


colinbennettbrain Apr 16, 2026
61be582

fix(pydantic_ai): restore audio + dict-style usage details on leaf span


colinbennettbrain Apr 17, 2026
171897c

Merge branch 'main' into colin/pydantic-strip-wrapper-tokens


colinbennettbrain Apr 17, 2026
246fcc2

fix(pydantic-ai): narrow Optional metrics type for pylint

Apr 17, 2026
39c970a

fix(py): fix pylint in pydantic ai metrics test


colinbennettbrain Apr 17, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


173 changes: 116 additions & 57 deletions
py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -67,6 +67,46 @@ def _assert_metrics_are_valid(metrics, start, end):
 assert metrics["completion_tokens"] > 0


def test_extract_response_metrics_leaf_fields():
 """`_extract_response_metrics` is the only path that populates token metrics on
 the leaf chat span post-#312/strip-wrapper-tokens. Verify the audio-token and
 dict-style `details` fields that were dropped when `_extract_usage_metrics`
 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

 usage = RequestUsage(
 input_tokens=100,
 output_tokens=50,
 cache_read_tokens=10,
 cache_write_tokens=20,
 input_audio_tokens=7,
 output_audio_tokens=11,
 details={"reasoning_tokens": 13, "cached_tokens": 17},
 )
 response = SimpleNamespace(usage=usage)

 maybe_metrics = _extract_response_metrics(response, start_time=1.0, end_time=2.0)
 assert maybe_metrics is not None

 # pylint: disable=unsupported-membership-test,unsubscriptable-object
 assert maybe_metrics["prompt_tokens"] == 100
 assert maybe_metrics["completion_tokens"] == 50
 assert maybe_metrics["tokens"] == 150 # total_tokens is a property: input + output
 assert maybe_metrics["prompt_cache_creation_tokens"] == 20
 assert maybe_metrics["prompt_audio_tokens"] == 7
 assert maybe_metrics["completion_audio_tokens"] == 11
 assert maybe_metrics["completion_reasoning_tokens"] == 13
 # details["cached_tokens"] is read after cache_read_tokens, so it wins when both
 # are populated; this preserves the pre-strip behavior of `_extract_usage_metrics`.
 assert maybe_metrics["prompt_cached_tokens"] == 17
 # pylint: enable=unsupported-membership-test,unsubscriptable-object


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_direct_model_request_creates_nested_chat_span_without_class_scan(memory_logger, direct):
Expand DownExpand Up@@ -140,23 +180,21 @@ async def test_agent_run_async(memory_logger):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 assert agent_span["metrics"]["prompt_tokens"] > 0
 assert agent_span["metrics"]["completion_tokens"] > 0
 # Regression: wrapper agent_run span must NOT log token metrics. The leaf chat
 # span already logs them, and trace-tree rollup (self + descendants) plus any
 # unfiltered sum over metrics would otherwise double-count tokens/cost at the
 # parent regardless of span type.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in agent_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 # Regression: no double-counting of cost/tokens. Experiment-level aggregations
 # sum metrics across type='llm' spans, so a single agent turn must contribute
 # its tokens exactly once. The wrapper agent_run span logs the same usage as
 # the leaf chat span; only the leaf should be type=LLM.
 # Only the leaf chat span should be type=LLM, and it must own the token totals.
 llm_spans = [s for s in spans if s["span_attributes"]["type"] == SpanTypeAttribute.LLM]
 assert len(llm_spans) == 1, f"expected exactly one LLM-typed span, got {len(llm_spans)}"
 assert llm_spans[0]["span_id"] == chat_span["span_id"]
 llm_prompt_tokens_sum = sum(s["metrics"].get("prompt_tokens", 0) for s in llm_spans)
 llm_completion_tokens_sum = sum(s["metrics"].get("completion_tokens", 0) for s in llm_spans)
 assert llm_prompt_tokens_sum == chat_span["metrics"]["prompt_tokens"]
 assert llm_completion_tokens_sum == chat_span["metrics"]["completion_tokens"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -245,9 +283,12 @@ def is_descendant(child_span, ancestor_id):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_sync_span["metrics"]
 assert "completion_tokens" in agent_sync_span["metrics"]
 # Wrapper agent_run_sync span must not log token metrics (would double-count at rollup).
 assert "prompt_tokens" not in agent_sync_span["metrics"]
 assert "completion_tokens" not in agent_sync_span["metrics"]
 # Tokens live on the leaf chat span only.
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


def test_agent_to_cli_sync(memory_logger, monkeypatch):
Expand DownExpand Up@@ -544,9 +585,12 @@ async def test_agent_run_stream(memory_logger):
 print(f"span_parents: {chat_span['span_parents']}")
 print(f"metrics: {chat_span['metrics']}")

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper stream span must not log token metrics (would double-count at rollup).
 # time_to_first_token is asserted above; it's a non-summable timing metric and stays.
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -741,6 +785,15 @@ async def test_direct_model_request_stream(memory_logger, direct):
 assert ttft <= duration, f"time_to_first_token ({ttft}s) should be <= duration ({duration}s)"
 assert ttft < 3.0, f"time_to_first_token should be < 3s for API call, got {ttft}s"

 # Regression: model_request_stream wrapper span (TASK) must NOT log token metrics.
 # `_DirectStreamWrapper` here is constructed with span_type=TASK, so __aexit__
 # routes to `_wrapper_span_metrics`. The leaf chat span (when present) is the only
 # place tokens should be recorded.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in direct_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 print(f"✓ Direct stream time_to_first_token: {ttft}s (duration: {duration}s)")


Expand DownExpand Up@@ -842,9 +895,11 @@ def is_descendant(child_span, ancestor_id):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper agent_run span must not log token metrics (would double-count at rollup).
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -1143,9 +1198,9 @@ def is_descendant(child_span, ancestor_id):
 # Chat span may not have complete metrics since it's an intermediate span
 assert "start" in chat_span["metrics"]

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper agent_run_stream_sync span must not log token metrics.
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]


@pytest.mark.vcr
Expand DownExpand Up@@ -1183,6 +1238,14 @@ async def test_agent_run_stream_events(memory_logger):
 assert agent_span["metrics"]["event_count"] == event_count
 _assert_metrics_are_valid(agent_span["metrics"], start, end)

 # Regression: wrapper agent_run_stream_events span must NOT log token metrics.
 # Trace-tree rollup over non-scorer spans plus experiment-level token sums would
 # otherwise double-count against the leaf chat span.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in agent_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )


@pytest.mark.vcr
def test_direct_model_request_stream_sync(memory_logger, direct):
Expand DownExpand Up@@ -1223,6 +1286,15 @@ def test_direct_model_request_stream_sync(memory_logger, direct):
 assert ttft <= duration, f"time_to_first_token ({ttft}s) should be <= duration ({duration}s)"
 assert ttft < 3.0, f"time_to_first_token should be < 3s for API call, got {ttft}s"

 # Regression: model_request_stream_sync wrapper span must NOT log token metrics.
 # `_DirectStreamWrapperSync.__exit__` always uses `_wrapper_span_metrics`; sync
 # direct streaming has no leaf chat span, so this is the only span we emit and
 # token totals would otherwise be derived twice once a leaf is added.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 print(f"✓ Direct sync stream time_to_first_token: {ttft}s (duration: {duration}s)")


Expand DownExpand Up@@ -2030,53 +2102,40 @@ def helper_tool() -> str:
def test_reasoning_tokens_extraction(memory_logger):
 """Test that reasoning tokens are extracted from model responses.

 For reasoning models like o1/o3, usage.details.reasoning_tokens should be
 captured in the metrics field.
 For reasoning models like o1/o3, providers (e.g. OpenAI) populate
 `usage.details["reasoning_tokens"]` (a dict entry, not an attribute) and
 `_extract_response_metrics` must surface it as `completion_reasoning_tokens`.
 """
 assert not memory_logger.pop()

 # Mock a response that has reasoning tokens
 from unittest.mock import MagicMock
 from types import SimpleNamespace

 # Create a mock response with reasoning tokens
 mock_response = MagicMock()
 mock_response.parts = [
 MagicMock(
 part_kind="thinking",
 content="Let me think about this...",
 ),
 MagicMock(
 part_kind="text",
 content="The answer is 42",
 ),
 ]
 mock_response.usage = MagicMock()
 mock_response.usage.input_tokens = 10
 mock_response.usage.output_tokens = 20
 mock_response.usage.total_tokens = 30
 mock_response.usage.cache_read_tokens = 0
 mock_response.usage.cache_write_tokens = 0
 mock_response.usage.details = MagicMock()
 mock_response.usage.details.reasoning_tokens = 128

 # Test the metric extraction function directly
 from braintrust.integrations.pydantic_ai.tracing import _extract_response_metrics
 from pydantic_ai.usage import RequestUsage

 # Use a real RequestUsage so `details` is the actual dict[str, int] shape
 # pydantic_ai produces. A MagicMock would falsely satisfy attribute-style
 # lookups and hide real-world dict-only behavior.
 usage = RequestUsage(
 input_tokens=10,
 output_tokens=20,
 cache_read_tokens=0,
 cache_write_tokens=0,
 details={"reasoning_tokens": 128},
 )
 response = SimpleNamespace(parts=[], usage=usage)

 start_time = time.time()
 end_time = start_time + 5.0

 metrics = _extract_response_metrics(mock_response, start_time, end_time)
 metrics = _extract_response_metrics(response, start_time, end_time)

 # Verify all metrics are present
 assert metrics is not None, "Should extract metrics"
 # pylint: disable=unsupported-membership-test,unsubscriptable-object
 assert "prompt_tokens" in metrics, "Should have prompt_tokens"
 assert metrics["prompt_tokens"] == 10.0
 assert "completion_tokens" in metrics, "Should have completion_tokens"
 assert metrics["completion_tokens"] == 20.0
 assert "tokens" in metrics, "Should have total tokens"
 assert metrics["tokens"] == 30.0
 assert "completion_reasoning_tokens" in metrics, "Should have completion_reasoning_tokens"
 assert metrics["tokens"] == 30.0 # total_tokens is a property: input + output
 assert metrics["completion_reasoning_tokens"] == 128.0, (
 f"Expected 128.0, got {metrics['completion_reasoning_tokens']}"
 )
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Highlight search terms from Google/DuckDuckGo/Bing referrer
(function() {
var ref = document.referrer;
var terms = [];
if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {
var url = new URL(ref);
var q = url.searchParams.get('q') || url.searchParams.get('p');
if (q) {
terms = q.split(/\s+/).filter(function(t) { return t.length > 2; });
}
}
if (terms.length === 0) return;
var style = document.createElement('style');
style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';
document.head.appendChild(style);
function highlight(node) {
if (node.nodeType === 3) { // text node
var text = node.textContent;
var found = false;
terms.forEach(function(term) {
var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\]\\]/g, '\\') + ')', 'gi');
if (regex.test(text)) {
found = true;
var frag = document.createDocumentFragment();
var parts = text.split(regex);
parts.forEach(function(part, i) {
if (i % 2 === 0) {
frag.appendChild(document.createTextNode(part));
} else {
var span = document.createElement('span');
span.className = 'userscript-highlight';
span.textContent = part;
frag.appendChild(span);
}
});
node.parentNode.replaceChild(frag, node);
}
});
} else if (node.nodeType === 1 && node.childNodes) { // element
var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];
if (!skipTags.includes(node.tagName)) {
Array.from(node.childNodes).forEach(highlight);
}
}
}
highlight(document.body);
// Re-highlight on dynamic content
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1 || node.nodeType === 3) highlight(node);
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
fix(pydantic_ai): strip unnecessary wrapper tokens by colinbennettbrain · Pull Request #321 · braintrustdata/braintrust-sdk-python · GitHub
Skip to content
fix(pydantic_ai): strip unnecessary wrapper tokens#321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

Abhijeet Prasad (AbhiPrasad)
merged 5 commits into

from

Apr 20, 2026
 Merged


fix(pydantic_ai): strip unnecessary wrapper tokens
#321
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
6a6cb76

fix(pydantic_ai): strip token metrics from wrapper spans


colinbennettbrain Apr 16, 2026
61be582

fix(pydantic_ai): restore audio + dict-style usage details on leaf span


colinbennettbrain Apr 17, 2026
171897c

Merge branch 'main' into colin/pydantic-strip-wrapper-tokens


colinbennettbrain Apr 17, 2026
246fcc2

fix(pydantic-ai): narrow Optional metrics type for pylint

Apr 17, 2026
39c970a

fix(py): fix pylint in pydantic ai metrics test


colinbennettbrain Apr 17, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


173 changes: 116 additions & 57 deletions
py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -67,6 +67,46 @@ def _assert_metrics_are_valid(metrics, start, end):
 assert metrics["completion_tokens"] > 0


def test_extract_response_metrics_leaf_fields():
 """`_extract_response_metrics` is the only path that populates token metrics on
 the leaf chat span post-#312/strip-wrapper-tokens. Verify the audio-token and
 dict-style `details` fields that were dropped when `_extract_usage_metrics`
 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

 usage = RequestUsage(
 input_tokens=100,
 output_tokens=50,
 cache_read_tokens=10,
 cache_write_tokens=20,
 input_audio_tokens=7,
 output_audio_tokens=11,
 details={"reasoning_tokens": 13, "cached_tokens": 17},
 )
 response = SimpleNamespace(usage=usage)

 maybe_metrics = _extract_response_metrics(response, start_time=1.0, end_time=2.0)
 assert maybe_metrics is not None

 # pylint: disable=unsupported-membership-test,unsubscriptable-object
 assert maybe_metrics["prompt_tokens"] == 100
 assert maybe_metrics["completion_tokens"] == 50
 assert maybe_metrics["tokens"] == 150 # total_tokens is a property: input + output
 assert maybe_metrics["prompt_cache_creation_tokens"] == 20
 assert maybe_metrics["prompt_audio_tokens"] == 7
 assert maybe_metrics["completion_audio_tokens"] == 11
 assert maybe_metrics["completion_reasoning_tokens"] == 13
 # details["cached_tokens"] is read after cache_read_tokens, so it wins when both
 # are populated; this preserves the pre-strip behavior of `_extract_usage_metrics`.
 assert maybe_metrics["prompt_cached_tokens"] == 17
 # pylint: enable=unsupported-membership-test,unsubscriptable-object


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_direct_model_request_creates_nested_chat_span_without_class_scan(memory_logger, direct):
Expand DownExpand Up@@ -140,23 +180,21 @@ async def test_agent_run_async(memory_logger):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 assert agent_span["metrics"]["prompt_tokens"] > 0
 assert agent_span["metrics"]["completion_tokens"] > 0
 # Regression: wrapper agent_run span must NOT log token metrics. The leaf chat
 # span already logs them, and trace-tree rollup (self + descendants) plus any
 # unfiltered sum over metrics would otherwise double-count tokens/cost at the
 # parent regardless of span type.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in agent_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 # Regression: no double-counting of cost/tokens. Experiment-level aggregations
 # sum metrics across type='llm' spans, so a single agent turn must contribute
 # its tokens exactly once. The wrapper agent_run span logs the same usage as
 # the leaf chat span; only the leaf should be type=LLM.
 # Only the leaf chat span should be type=LLM, and it must own the token totals.
 llm_spans = [s for s in spans if s["span_attributes"]["type"] == SpanTypeAttribute.LLM]
 assert len(llm_spans) == 1, f"expected exactly one LLM-typed span, got {len(llm_spans)}"
 assert llm_spans[0]["span_id"] == chat_span["span_id"]
 llm_prompt_tokens_sum = sum(s["metrics"].get("prompt_tokens", 0) for s in llm_spans)
 llm_completion_tokens_sum = sum(s["metrics"].get("completion_tokens", 0) for s in llm_spans)
 assert llm_prompt_tokens_sum == chat_span["metrics"]["prompt_tokens"]
 assert llm_completion_tokens_sum == chat_span["metrics"]["completion_tokens"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -245,9 +283,12 @@ def is_descendant(child_span, ancestor_id):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_sync_span["metrics"]
 assert "completion_tokens" in agent_sync_span["metrics"]
 # Wrapper agent_run_sync span must not log token metrics (would double-count at rollup).
 assert "prompt_tokens" not in agent_sync_span["metrics"]
 assert "completion_tokens" not in agent_sync_span["metrics"]
 # Tokens live on the leaf chat span only.
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


def test_agent_to_cli_sync(memory_logger, monkeypatch):
Expand DownExpand Up@@ -544,9 +585,12 @@ async def test_agent_run_stream(memory_logger):
 print(f"span_parents: {chat_span['span_parents']}")
 print(f"metrics: {chat_span['metrics']}")

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper stream span must not log token metrics (would double-count at rollup).
 # time_to_first_token is asserted above; it's a non-summable timing metric and stays.
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -741,6 +785,15 @@ async def test_direct_model_request_stream(memory_logger, direct):
 assert ttft <= duration, f"time_to_first_token ({ttft}s) should be <= duration ({duration}s)"
 assert ttft < 3.0, f"time_to_first_token should be < 3s for API call, got {ttft}s"

 # Regression: model_request_stream wrapper span (TASK) must NOT log token metrics.
 # `_DirectStreamWrapper` here is constructed with span_type=TASK, so __aexit__
 # routes to `_wrapper_span_metrics`. The leaf chat span (when present) is the only
 # place tokens should be recorded.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in direct_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 print(f"✓ Direct stream time_to_first_token: {ttft}s (duration: {duration}s)")


Expand DownExpand Up@@ -842,9 +895,11 @@ def is_descendant(child_span, ancestor_id):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper agent_run span must not log token metrics (would double-count at rollup).
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -1143,9 +1198,9 @@ def is_descendant(child_span, ancestor_id):
 # Chat span may not have complete metrics since it's an intermediate span
 assert "start" in chat_span["metrics"]

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper agent_run_stream_sync span must not log token metrics.
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]


@pytest.mark.vcr
Expand DownExpand Up@@ -1183,6 +1238,14 @@ async def test_agent_run_stream_events(memory_logger):
 assert agent_span["metrics"]["event_count"] == event_count
 _assert_metrics_are_valid(agent_span["metrics"], start, end)

 # Regression: wrapper agent_run_stream_events span must NOT log token metrics.
 # Trace-tree rollup over non-scorer spans plus experiment-level token sums would
 # otherwise double-count against the leaf chat span.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in agent_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )


@pytest.mark.vcr
def test_direct_model_request_stream_sync(memory_logger, direct):
Expand DownExpand Up@@ -1223,6 +1286,15 @@ def test_direct_model_request_stream_sync(memory_logger, direct):
 assert ttft <= duration, f"time_to_first_token ({ttft}s) should be <= duration ({duration}s)"
 assert ttft < 3.0, f"time_to_first_token should be < 3s for API call, got {ttft}s"

 # Regression: model_request_stream_sync wrapper span must NOT log token metrics.
 # `_DirectStreamWrapperSync.__exit__` always uses `_wrapper_span_metrics`; sync
 # direct streaming has no leaf chat span, so this is the only span we emit and
 # token totals would otherwise be derived twice once a leaf is added.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 print(f"✓ Direct sync stream time_to_first_token: {ttft}s (duration: {duration}s)")


Expand DownExpand Up@@ -2030,53 +2102,40 @@ def helper_tool() -> str:
def test_reasoning_tokens_extraction(memory_logger):
 """Test that reasoning tokens are extracted from model responses.

 For reasoning models like o1/o3, usage.details.reasoning_tokens should be
 captured in the metrics field.
 For reasoning models like o1/o3, providers (e.g. OpenAI) populate
 `usage.details["reasoning_tokens"]` (a dict entry, not an attribute) and
 `_extract_response_metrics` must surface it as `completion_reasoning_tokens`.
 """
 assert not memory_logger.pop()

 # Mock a response that has reasoning tokens
 from unittest.mock import MagicMock
 from types import SimpleNamespace

 # Create a mock response with reasoning tokens
 mock_response = MagicMock()
 mock_response.parts = [
 MagicMock(
 part_kind="thinking",
 content="Let me think about this...",
 ),
 MagicMock(
 part_kind="text",
 content="The answer is 42",
 ),
 ]
 mock_response.usage = MagicMock()
 mock_response.usage.input_tokens = 10
 mock_response.usage.output_tokens = 20
 mock_response.usage.total_tokens = 30
 mock_response.usage.cache_read_tokens = 0
 mock_response.usage.cache_write_tokens = 0
 mock_response.usage.details = MagicMock()
 mock_response.usage.details.reasoning_tokens = 128

 # Test the metric extraction function directly
 from braintrust.integrations.pydantic_ai.tracing import _extract_response_metrics
 from pydantic_ai.usage import RequestUsage

 # Use a real RequestUsage so `details` is the actual dict[str, int] shape
 # pydantic_ai produces. A MagicMock would falsely satisfy attribute-style
 # lookups and hide real-world dict-only behavior.
 usage = RequestUsage(
 input_tokens=10,
 output_tokens=20,
 cache_read_tokens=0,
 cache_write_tokens=0,
 details={"reasoning_tokens": 128},
 )
 response = SimpleNamespace(parts=[], usage=usage)

 start_time = time.time()
 end_time = start_time + 5.0

 metrics = _extract_response_metrics(mock_response, start_time, end_time)
 metrics = _extract_response_metrics(response, start_time, end_time)

 # Verify all metrics are present
 assert metrics is not None, "Should extract metrics"
 # pylint: disable=unsupported-membership-test,unsubscriptable-object
 assert "prompt_tokens" in metrics, "Should have prompt_tokens"
 assert metrics["prompt_tokens"] == 10.0
 assert "completion_tokens" in metrics, "Should have completion_tokens"
 assert metrics["completion_tokens"] == 20.0
 assert "tokens" in metrics, "Should have total tokens"
 assert metrics["tokens"] == 30.0
 assert "completion_reasoning_tokens" in metrics, "Should have completion_reasoning_tokens"
 assert metrics["tokens"] == 30.0 # total_tokens is a property: input + output
 assert metrics["completion_reasoning_tokens"] == 128.0, (
 f"Expected 128.0, got {metrics['completion_reasoning_tokens']}"
 )
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Strip utm_, fbclid, gclid, etc. from all links on page
(function() {
var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',
'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',
'ref', 'ref_src', 'source', 'medium', 'campaign'];
function cleanUrl(url) {
try {
var u = new URL(url, window.location.origin);
var changed = false;
trackingParams.forEach(function(p) {
if (u.searchParams.has(p)) {
u.searchParams.delete(p);
changed = true;
}
});
return changed ? u.toString() : url;
} catch (e) {
return url;
}
}
function cleanLinks() {
document.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
cleanLinks();
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1) {
if (node.tagName === 'A') cleanLinks();
node.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '
fix(pydantic_ai): strip unnecessary wrapper tokens by colinbennettbrain · Pull Request #321 · braintrustdata/braintrust-sdk-python · GitHub
Skip to content
fix(pydantic_ai): strip unnecessary wrapper tokens#321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

Abhijeet Prasad (AbhiPrasad)
merged 5 commits into

from

Apr 20, 2026
 Merged


fix(pydantic_ai): strip unnecessary wrapper tokens
#321
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
6a6cb76

fix(pydantic_ai): strip token metrics from wrapper spans


colinbennettbrain Apr 16, 2026
61be582

fix(pydantic_ai): restore audio + dict-style usage details on leaf span


colinbennettbrain Apr 17, 2026
171897c

Merge branch 'main' into colin/pydantic-strip-wrapper-tokens


colinbennettbrain Apr 17, 2026
246fcc2

fix(pydantic-ai): narrow Optional metrics type for pylint

Apr 17, 2026
39c970a

fix(py): fix pylint in pydantic ai metrics test


colinbennettbrain Apr 17, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


173 changes: 116 additions & 57 deletions
py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -67,6 +67,46 @@ def _assert_metrics_are_valid(metrics, start, end):
 assert metrics["completion_tokens"] > 0


def test_extract_response_metrics_leaf_fields():
 """`_extract_response_metrics` is the only path that populates token metrics on
 the leaf chat span post-#312/strip-wrapper-tokens. Verify the audio-token and
 dict-style `details` fields that were dropped when `_extract_usage_metrics`
 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

 usage = RequestUsage(
 input_tokens=100,
 output_tokens=50,
 cache_read_tokens=10,
 cache_write_tokens=20,
 input_audio_tokens=7,
 output_audio_tokens=11,
 details={"reasoning_tokens": 13, "cached_tokens": 17},
 )
 response = SimpleNamespace(usage=usage)

 maybe_metrics = _extract_response_metrics(response, start_time=1.0, end_time=2.0)
 assert maybe_metrics is not None

 # pylint: disable=unsupported-membership-test,unsubscriptable-object
 assert maybe_metrics["prompt_tokens"] == 100
 assert maybe_metrics["completion_tokens"] == 50
 assert maybe_metrics["tokens"] == 150 # total_tokens is a property: input + output
 assert maybe_metrics["prompt_cache_creation_tokens"] == 20
 assert maybe_metrics["prompt_audio_tokens"] == 7
 assert maybe_metrics["completion_audio_tokens"] == 11
 assert maybe_metrics["completion_reasoning_tokens"] == 13
 # details["cached_tokens"] is read after cache_read_tokens, so it wins when both
 # are populated; this preserves the pre-strip behavior of `_extract_usage_metrics`.
 assert maybe_metrics["prompt_cached_tokens"] == 17
 # pylint: enable=unsupported-membership-test,unsubscriptable-object


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_direct_model_request_creates_nested_chat_span_without_class_scan(memory_logger, direct):
Expand DownExpand Up@@ -140,23 +180,21 @@ async def test_agent_run_async(memory_logger):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 assert agent_span["metrics"]["prompt_tokens"] > 0
 assert agent_span["metrics"]["completion_tokens"] > 0
 # Regression: wrapper agent_run span must NOT log token metrics. The leaf chat
 # span already logs them, and trace-tree rollup (self + descendants) plus any
 # unfiltered sum over metrics would otherwise double-count tokens/cost at the
 # parent regardless of span type.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in agent_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 # Regression: no double-counting of cost/tokens. Experiment-level aggregations
 # sum metrics across type='llm' spans, so a single agent turn must contribute
 # its tokens exactly once. The wrapper agent_run span logs the same usage as
 # the leaf chat span; only the leaf should be type=LLM.
 # Only the leaf chat span should be type=LLM, and it must own the token totals.
 llm_spans = [s for s in spans if s["span_attributes"]["type"] == SpanTypeAttribute.LLM]
 assert len(llm_spans) == 1, f"expected exactly one LLM-typed span, got {len(llm_spans)}"
 assert llm_spans[0]["span_id"] == chat_span["span_id"]
 llm_prompt_tokens_sum = sum(s["metrics"].get("prompt_tokens", 0) for s in llm_spans)
 llm_completion_tokens_sum = sum(s["metrics"].get("completion_tokens", 0) for s in llm_spans)
 assert llm_prompt_tokens_sum == chat_span["metrics"]["prompt_tokens"]
 assert llm_completion_tokens_sum == chat_span["metrics"]["completion_tokens"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -245,9 +283,12 @@ def is_descendant(child_span, ancestor_id):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_sync_span["metrics"]
 assert "completion_tokens" in agent_sync_span["metrics"]
 # Wrapper agent_run_sync span must not log token metrics (would double-count at rollup).
 assert "prompt_tokens" not in agent_sync_span["metrics"]
 assert "completion_tokens" not in agent_sync_span["metrics"]
 # Tokens live on the leaf chat span only.
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


def test_agent_to_cli_sync(memory_logger, monkeypatch):
Expand DownExpand Up@@ -544,9 +585,12 @@ async def test_agent_run_stream(memory_logger):
 print(f"span_parents: {chat_span['span_parents']}")
 print(f"metrics: {chat_span['metrics']}")

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper stream span must not log token metrics (would double-count at rollup).
 # time_to_first_token is asserted above; it's a non-summable timing metric and stays.
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -741,6 +785,15 @@ async def test_direct_model_request_stream(memory_logger, direct):
 assert ttft <= duration, f"time_to_first_token ({ttft}s) should be <= duration ({duration}s)"
 assert ttft < 3.0, f"time_to_first_token should be < 3s for API call, got {ttft}s"

 # Regression: model_request_stream wrapper span (TASK) must NOT log token metrics.
 # `_DirectStreamWrapper` here is constructed with span_type=TASK, so __aexit__
 # routes to `_wrapper_span_metrics`. The leaf chat span (when present) is the only
 # place tokens should be recorded.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in direct_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 print(f"✓ Direct stream time_to_first_token: {ttft}s (duration: {duration}s)")


Expand DownExpand Up@@ -842,9 +895,11 @@ def is_descendant(child_span, ancestor_id):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper agent_run span must not log token metrics (would double-count at rollup).
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -1143,9 +1198,9 @@ def is_descendant(child_span, ancestor_id):
 # Chat span may not have complete metrics since it's an intermediate span
 assert "start" in chat_span["metrics"]

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper agent_run_stream_sync span must not log token metrics.
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]


@pytest.mark.vcr
Expand DownExpand Up@@ -1183,6 +1238,14 @@ async def test_agent_run_stream_events(memory_logger):
 assert agent_span["metrics"]["event_count"] == event_count
 _assert_metrics_are_valid(agent_span["metrics"], start, end)

 # Regression: wrapper agent_run_stream_events span must NOT log token metrics.
 # Trace-tree rollup over non-scorer spans plus experiment-level token sums would
 # otherwise double-count against the leaf chat span.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in agent_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )


@pytest.mark.vcr
def test_direct_model_request_stream_sync(memory_logger, direct):
Expand DownExpand Up@@ -1223,6 +1286,15 @@ def test_direct_model_request_stream_sync(memory_logger, direct):
 assert ttft <= duration, f"time_to_first_token ({ttft}s) should be <= duration ({duration}s)"
 assert ttft < 3.0, f"time_to_first_token should be < 3s for API call, got {ttft}s"

 # Regression: model_request_stream_sync wrapper span must NOT log token metrics.
 # `_DirectStreamWrapperSync.__exit__` always uses `_wrapper_span_metrics`; sync
 # direct streaming has no leaf chat span, so this is the only span we emit and
 # token totals would otherwise be derived twice once a leaf is added.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 print(f"✓ Direct sync stream time_to_first_token: {ttft}s (duration: {duration}s)")


Expand DownExpand Up@@ -2030,53 +2102,40 @@ def helper_tool() -> str:
def test_reasoning_tokens_extraction(memory_logger):
 """Test that reasoning tokens are extracted from model responses.

 For reasoning models like o1/o3, usage.details.reasoning_tokens should be
 captured in the metrics field.
 For reasoning models like o1/o3, providers (e.g. OpenAI) populate
 `usage.details["reasoning_tokens"]` (a dict entry, not an attribute) and
 `_extract_response_metrics` must surface it as `completion_reasoning_tokens`.
 """
 assert not memory_logger.pop()

 # Mock a response that has reasoning tokens
 from unittest.mock import MagicMock
 from types import SimpleNamespace

 # Create a mock response with reasoning tokens
 mock_response = MagicMock()
 mock_response.parts = [
 MagicMock(
 part_kind="thinking",
 content="Let me think about this...",
 ),
 MagicMock(
 part_kind="text",
 content="The answer is 42",
 ),
 ]
 mock_response.usage = MagicMock()
 mock_response.usage.input_tokens = 10
 mock_response.usage.output_tokens = 20
 mock_response.usage.total_tokens = 30
 mock_response.usage.cache_read_tokens = 0
 mock_response.usage.cache_write_tokens = 0
 mock_response.usage.details = MagicMock()
 mock_response.usage.details.reasoning_tokens = 128

 # Test the metric extraction function directly
 from braintrust.integrations.pydantic_ai.tracing import _extract_response_metrics
 from pydantic_ai.usage import RequestUsage

 # Use a real RequestUsage so `details` is the actual dict[str, int] shape
 # pydantic_ai produces. A MagicMock would falsely satisfy attribute-style
 # lookups and hide real-world dict-only behavior.
 usage = RequestUsage(
 input_tokens=10,
 output_tokens=20,
 cache_read_tokens=0,
 cache_write_tokens=0,
 details={"reasoning_tokens": 128},
 )
 response = SimpleNamespace(parts=[], usage=usage)

 start_time = time.time()
 end_time = start_time + 5.0

 metrics = _extract_response_metrics(mock_response, start_time, end_time)
 metrics = _extract_response_metrics(response, start_time, end_time)

 # Verify all metrics are present
 assert metrics is not None, "Should extract metrics"
 # pylint: disable=unsupported-membership-test,unsubscriptable-object
 assert "prompt_tokens" in metrics, "Should have prompt_tokens"
 assert metrics["prompt_tokens"] == 10.0
 assert "completion_tokens" in metrics, "Should have completion_tokens"
 assert metrics["completion_tokens"] == 20.0
 assert "tokens" in metrics, "Should have total tokens"
 assert metrics["tokens"] == 30.0
 assert "completion_reasoning_tokens" in metrics, "Should have completion_reasoning_tokens"
 assert metrics["tokens"] == 30.0 # total_tokens is a property: input + output
 assert metrics["completion_reasoning_tokens"] == 128.0, (
 f"Expected 128.0, got {metrics['completion_reasoning_tokens']}"
 )
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Auto-enable theater mode on YouTube
(function() {
function tryTheater() {
var btn = document.querySelector('button[aria-label="Theater mode"], ytd-player #player button[title="Theater mode"]');
if (btn && !btn.classList.contains('activated')) {
btn.click();
}
}
// Try immediately
tryTheater();
// Try after navigation (SPA)
var lastUrl = location.href;
setInterval(function() {
if (location.href !== lastUrl) {
lastUrl = location.href;
setTimeout(tryTheater, 500);
}
}, 1000);
// Also try on player load
var observer = new MutationObserver(tryTheater);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
fix(pydantic_ai): strip unnecessary wrapper tokens by colinbennettbrain · Pull Request #321 · braintrustdata/braintrust-sdk-python · GitHub
Skip to content
fix(pydantic_ai): strip unnecessary wrapper tokens#321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

Abhijeet Prasad (AbhiPrasad)
merged 5 commits into

from

Apr 20, 2026
 Merged


fix(pydantic_ai): strip unnecessary wrapper tokens
#321
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
6a6cb76

fix(pydantic_ai): strip token metrics from wrapper spans


colinbennettbrain Apr 16, 2026
61be582

fix(pydantic_ai): restore audio + dict-style usage details on leaf span


colinbennettbrain Apr 17, 2026
171897c

Merge branch 'main' into colin/pydantic-strip-wrapper-tokens


colinbennettbrain Apr 17, 2026
246fcc2

fix(pydantic-ai): narrow Optional metrics type for pylint

Apr 17, 2026
39c970a

fix(py): fix pylint in pydantic ai metrics test


colinbennettbrain Apr 17, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


173 changes: 116 additions & 57 deletions
py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -67,6 +67,46 @@ def _assert_metrics_are_valid(metrics, start, end):
 assert metrics["completion_tokens"] > 0


def test_extract_response_metrics_leaf_fields():
 """`_extract_response_metrics` is the only path that populates token metrics on
 the leaf chat span post-#312/strip-wrapper-tokens. Verify the audio-token and
 dict-style `details` fields that were dropped when `_extract_usage_metrics`
 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

 usage = RequestUsage(
 input_tokens=100,
 output_tokens=50,
 cache_read_tokens=10,
 cache_write_tokens=20,
 input_audio_tokens=7,
 output_audio_tokens=11,
 details={"reasoning_tokens": 13, "cached_tokens": 17},
 )
 response = SimpleNamespace(usage=usage)

 maybe_metrics = _extract_response_metrics(response, start_time=1.0, end_time=2.0)
 assert maybe_metrics is not None

 # pylint: disable=unsupported-membership-test,unsubscriptable-object
 assert maybe_metrics["prompt_tokens"] == 100
 assert maybe_metrics["completion_tokens"] == 50
 assert maybe_metrics["tokens"] == 150 # total_tokens is a property: input + output
 assert maybe_metrics["prompt_cache_creation_tokens"] == 20
 assert maybe_metrics["prompt_audio_tokens"] == 7
 assert maybe_metrics["completion_audio_tokens"] == 11
 assert maybe_metrics["completion_reasoning_tokens"] == 13
 # details["cached_tokens"] is read after cache_read_tokens, so it wins when both
 # are populated; this preserves the pre-strip behavior of `_extract_usage_metrics`.
 assert maybe_metrics["prompt_cached_tokens"] == 17
 # pylint: enable=unsupported-membership-test,unsubscriptable-object


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_direct_model_request_creates_nested_chat_span_without_class_scan(memory_logger, direct):
Expand DownExpand Up@@ -140,23 +180,21 @@ async def test_agent_run_async(memory_logger):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 assert agent_span["metrics"]["prompt_tokens"] > 0
 assert agent_span["metrics"]["completion_tokens"] > 0
 # Regression: wrapper agent_run span must NOT log token metrics. The leaf chat
 # span already logs them, and trace-tree rollup (self + descendants) plus any
 # unfiltered sum over metrics would otherwise double-count tokens/cost at the
 # parent regardless of span type.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in agent_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 # Regression: no double-counting of cost/tokens. Experiment-level aggregations
 # sum metrics across type='llm' spans, so a single agent turn must contribute
 # its tokens exactly once. The wrapper agent_run span logs the same usage as
 # the leaf chat span; only the leaf should be type=LLM.
 # Only the leaf chat span should be type=LLM, and it must own the token totals.
 llm_spans = [s for s in spans if s["span_attributes"]["type"] == SpanTypeAttribute.LLM]
 assert len(llm_spans) == 1, f"expected exactly one LLM-typed span, got {len(llm_spans)}"
 assert llm_spans[0]["span_id"] == chat_span["span_id"]
 llm_prompt_tokens_sum = sum(s["metrics"].get("prompt_tokens", 0) for s in llm_spans)
 llm_completion_tokens_sum = sum(s["metrics"].get("completion_tokens", 0) for s in llm_spans)
 assert llm_prompt_tokens_sum == chat_span["metrics"]["prompt_tokens"]
 assert llm_completion_tokens_sum == chat_span["metrics"]["completion_tokens"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -245,9 +283,12 @@ def is_descendant(child_span, ancestor_id):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_sync_span["metrics"]
 assert "completion_tokens" in agent_sync_span["metrics"]
 # Wrapper agent_run_sync span must not log token metrics (would double-count at rollup).
 assert "prompt_tokens" not in agent_sync_span["metrics"]
 assert "completion_tokens" not in agent_sync_span["metrics"]
 # Tokens live on the leaf chat span only.
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


def test_agent_to_cli_sync(memory_logger, monkeypatch):
Expand DownExpand Up@@ -544,9 +585,12 @@ async def test_agent_run_stream(memory_logger):
 print(f"span_parents: {chat_span['span_parents']}")
 print(f"metrics: {chat_span['metrics']}")

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper stream span must not log token metrics (would double-count at rollup).
 # time_to_first_token is asserted above; it's a non-summable timing metric and stays.
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -741,6 +785,15 @@ async def test_direct_model_request_stream(memory_logger, direct):
 assert ttft <= duration, f"time_to_first_token ({ttft}s) should be <= duration ({duration}s)"
 assert ttft < 3.0, f"time_to_first_token should be < 3s for API call, got {ttft}s"

 # Regression: model_request_stream wrapper span (TASK) must NOT log token metrics.
 # `_DirectStreamWrapper` here is constructed with span_type=TASK, so __aexit__
 # routes to `_wrapper_span_metrics`. The leaf chat span (when present) is the only
 # place tokens should be recorded.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in direct_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 print(f"✓ Direct stream time_to_first_token: {ttft}s (duration: {duration}s)")


Expand DownExpand Up@@ -842,9 +895,11 @@ def is_descendant(child_span, ancestor_id):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper agent_run span must not log token metrics (would double-count at rollup).
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -1143,9 +1198,9 @@ def is_descendant(child_span, ancestor_id):
 # Chat span may not have complete metrics since it's an intermediate span
 assert "start" in chat_span["metrics"]

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper agent_run_stream_sync span must not log token metrics.
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]


@pytest.mark.vcr
Expand DownExpand Up@@ -1183,6 +1238,14 @@ async def test_agent_run_stream_events(memory_logger):
 assert agent_span["metrics"]["event_count"] == event_count
 _assert_metrics_are_valid(agent_span["metrics"], start, end)

 # Regression: wrapper agent_run_stream_events span must NOT log token metrics.
 # Trace-tree rollup over non-scorer spans plus experiment-level token sums would
 # otherwise double-count against the leaf chat span.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in agent_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )


@pytest.mark.vcr
def test_direct_model_request_stream_sync(memory_logger, direct):
Expand DownExpand Up@@ -1223,6 +1286,15 @@ def test_direct_model_request_stream_sync(memory_logger, direct):
 assert ttft <= duration, f"time_to_first_token ({ttft}s) should be <= duration ({duration}s)"
 assert ttft < 3.0, f"time_to_first_token should be < 3s for API call, got {ttft}s"

 # Regression: model_request_stream_sync wrapper span must NOT log token metrics.
 # `_DirectStreamWrapperSync.__exit__` always uses `_wrapper_span_metrics`; sync
 # direct streaming has no leaf chat span, so this is the only span we emit and
 # token totals would otherwise be derived twice once a leaf is added.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 print(f"✓ Direct sync stream time_to_first_token: {ttft}s (duration: {duration}s)")


Expand DownExpand Up@@ -2030,53 +2102,40 @@ def helper_tool() -> str:
def test_reasoning_tokens_extraction(memory_logger):
 """Test that reasoning tokens are extracted from model responses.

 For reasoning models like o1/o3, usage.details.reasoning_tokens should be
 captured in the metrics field.
 For reasoning models like o1/o3, providers (e.g. OpenAI) populate
 `usage.details["reasoning_tokens"]` (a dict entry, not an attribute) and
 `_extract_response_metrics` must surface it as `completion_reasoning_tokens`.
 """
 assert not memory_logger.pop()

 # Mock a response that has reasoning tokens
 from unittest.mock import MagicMock
 from types import SimpleNamespace

 # Create a mock response with reasoning tokens
 mock_response = MagicMock()
 mock_response.parts = [
 MagicMock(
 part_kind="thinking",
 content="Let me think about this...",
 ),
 MagicMock(
 part_kind="text",
 content="The answer is 42",
 ),
 ]
 mock_response.usage = MagicMock()
 mock_response.usage.input_tokens = 10
 mock_response.usage.output_tokens = 20
 mock_response.usage.total_tokens = 30
 mock_response.usage.cache_read_tokens = 0
 mock_response.usage.cache_write_tokens = 0
 mock_response.usage.details = MagicMock()
 mock_response.usage.details.reasoning_tokens = 128

 # Test the metric extraction function directly
 from braintrust.integrations.pydantic_ai.tracing import _extract_response_metrics
 from pydantic_ai.usage import RequestUsage

 # Use a real RequestUsage so `details` is the actual dict[str, int] shape
 # pydantic_ai produces. A MagicMock would falsely satisfy attribute-style
 # lookups and hide real-world dict-only behavior.
 usage = RequestUsage(
 input_tokens=10,
 output_tokens=20,
 cache_read_tokens=0,
 cache_write_tokens=0,
 details={"reasoning_tokens": 128},
 )
 response = SimpleNamespace(parts=[], usage=usage)

 start_time = time.time()
 end_time = start_time + 5.0

 metrics = _extract_response_metrics(mock_response, start_time, end_time)
 metrics = _extract_response_metrics(response, start_time, end_time)

 # Verify all metrics are present
 assert metrics is not None, "Should extract metrics"
 # pylint: disable=unsupported-membership-test,unsubscriptable-object
 assert "prompt_tokens" in metrics, "Should have prompt_tokens"
 assert metrics["prompt_tokens"] == 10.0
 assert "completion_tokens" in metrics, "Should have completion_tokens"
 assert metrics["completion_tokens"] == 20.0
 assert "tokens" in metrics, "Should have total tokens"
 assert metrics["tokens"] == 30.0
 assert "completion_reasoning_tokens" in metrics, "Should have completion_reasoning_tokens"
 assert metrics["tokens"] == 30.0 # total_tokens is a property: input + output
 assert metrics["completion_reasoning_tokens"] == 128.0, (
 f"Expected 128.0, got {metrics['completion_reasoning_tokens']}"
 )
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Remove or un-stick sticky/fixed headers that block content
(function() {
function unstick() {
document.querySelectorAll('header, nav, [role="banner"], .header, .navbar, .sticky, .fixed-top, [style*="position: fixed"], [style*="position:sticky"]').forEach(function(el) {
if (el.style.position === 'fixed' || el.style.position === 'sticky' || getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {
el.style.position = 'static';
el.style.top = 'auto';
el.style.zIndex = 'auto';
}
});
}
unstick();
var observer = new MutationObserver(unstick);
observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });
})();
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
fix(pydantic_ai): strip unnecessary wrapper tokens by colinbennettbrain · Pull Request #321 · braintrustdata/braintrust-sdk-python · GitHub
Skip to content
fix(pydantic_ai): strip unnecessary wrapper tokens#321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

Abhijeet Prasad (AbhiPrasad)
merged 5 commits into

from

Apr 20, 2026
 Merged


fix(pydantic_ai): strip unnecessary wrapper tokens
#321
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
6a6cb76

fix(pydantic_ai): strip token metrics from wrapper spans


colinbennettbrain Apr 16, 2026
61be582

fix(pydantic_ai): restore audio + dict-style usage details on leaf span


colinbennettbrain Apr 17, 2026
171897c

Merge branch 'main' into colin/pydantic-strip-wrapper-tokens


colinbennettbrain Apr 17, 2026
246fcc2

fix(pydantic-ai): narrow Optional metrics type for pylint

Apr 17, 2026
39c970a

fix(py): fix pylint in pydantic ai metrics test


colinbennettbrain Apr 17, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


173 changes: 116 additions & 57 deletions
py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -67,6 +67,46 @@ def _assert_metrics_are_valid(metrics, start, end):
 assert metrics["completion_tokens"] > 0


def test_extract_response_metrics_leaf_fields():
 """`_extract_response_metrics` is the only path that populates token metrics on
 the leaf chat span post-#312/strip-wrapper-tokens. Verify the audio-token and
 dict-style `details` fields that were dropped when `_extract_usage_metrics`
 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

 usage = RequestUsage(
 input_tokens=100,
 output_tokens=50,
 cache_read_tokens=10,
 cache_write_tokens=20,
 input_audio_tokens=7,
 output_audio_tokens=11,
 details={"reasoning_tokens": 13, "cached_tokens": 17},
 )
 response = SimpleNamespace(usage=usage)

 maybe_metrics = _extract_response_metrics(response, start_time=1.0, end_time=2.0)
 assert maybe_metrics is not None

 # pylint: disable=unsupported-membership-test,unsubscriptable-object
 assert maybe_metrics["prompt_tokens"] == 100
 assert maybe_metrics["completion_tokens"] == 50
 assert maybe_metrics["tokens"] == 150 # total_tokens is a property: input + output
 assert maybe_metrics["prompt_cache_creation_tokens"] == 20
 assert maybe_metrics["prompt_audio_tokens"] == 7
 assert maybe_metrics["completion_audio_tokens"] == 11
 assert maybe_metrics["completion_reasoning_tokens"] == 13
 # details["cached_tokens"] is read after cache_read_tokens, so it wins when both
 # are populated; this preserves the pre-strip behavior of `_extract_usage_metrics`.
 assert maybe_metrics["prompt_cached_tokens"] == 17
 # pylint: enable=unsupported-membership-test,unsubscriptable-object


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_direct_model_request_creates_nested_chat_span_without_class_scan(memory_logger, direct):
Expand DownExpand Up@@ -140,23 +180,21 @@ async def test_agent_run_async(memory_logger):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 assert agent_span["metrics"]["prompt_tokens"] > 0
 assert agent_span["metrics"]["completion_tokens"] > 0
 # Regression: wrapper agent_run span must NOT log token metrics. The leaf chat
 # span already logs them, and trace-tree rollup (self + descendants) plus any
 # unfiltered sum over metrics would otherwise double-count tokens/cost at the
 # parent regardless of span type.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in agent_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 # Regression: no double-counting of cost/tokens. Experiment-level aggregations
 # sum metrics across type='llm' spans, so a single agent turn must contribute
 # its tokens exactly once. The wrapper agent_run span logs the same usage as
 # the leaf chat span; only the leaf should be type=LLM.
 # Only the leaf chat span should be type=LLM, and it must own the token totals.
 llm_spans = [s for s in spans if s["span_attributes"]["type"] == SpanTypeAttribute.LLM]
 assert len(llm_spans) == 1, f"expected exactly one LLM-typed span, got {len(llm_spans)}"
 assert llm_spans[0]["span_id"] == chat_span["span_id"]
 llm_prompt_tokens_sum = sum(s["metrics"].get("prompt_tokens", 0) for s in llm_spans)
 llm_completion_tokens_sum = sum(s["metrics"].get("completion_tokens", 0) for s in llm_spans)
 assert llm_prompt_tokens_sum == chat_span["metrics"]["prompt_tokens"]
 assert llm_completion_tokens_sum == chat_span["metrics"]["completion_tokens"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -245,9 +283,12 @@ def is_descendant(child_span, ancestor_id):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_sync_span["metrics"]
 assert "completion_tokens" in agent_sync_span["metrics"]
 # Wrapper agent_run_sync span must not log token metrics (would double-count at rollup).
 assert "prompt_tokens" not in agent_sync_span["metrics"]
 assert "completion_tokens" not in agent_sync_span["metrics"]
 # Tokens live on the leaf chat span only.
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


def test_agent_to_cli_sync(memory_logger, monkeypatch):
Expand DownExpand Up@@ -544,9 +585,12 @@ async def test_agent_run_stream(memory_logger):
 print(f"span_parents: {chat_span['span_parents']}")
 print(f"metrics: {chat_span['metrics']}")

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper stream span must not log token metrics (would double-count at rollup).
 # time_to_first_token is asserted above; it's a non-summable timing metric and stays.
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -741,6 +785,15 @@ async def test_direct_model_request_stream(memory_logger, direct):
 assert ttft <= duration, f"time_to_first_token ({ttft}s) should be <= duration ({duration}s)"
 assert ttft < 3.0, f"time_to_first_token should be < 3s for API call, got {ttft}s"

 # Regression: model_request_stream wrapper span (TASK) must NOT log token metrics.
 # `_DirectStreamWrapper` here is constructed with span_type=TASK, so __aexit__
 # routes to `_wrapper_span_metrics`. The leaf chat span (when present) is the only
 # place tokens should be recorded.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in direct_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 print(f"✓ Direct stream time_to_first_token: {ttft}s (duration: {duration}s)")


Expand DownExpand Up@@ -842,9 +895,11 @@ def is_descendant(child_span, ancestor_id):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper agent_run span must not log token metrics (would double-count at rollup).
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -1143,9 +1198,9 @@ def is_descendant(child_span, ancestor_id):
 # Chat span may not have complete metrics since it's an intermediate span
 assert "start" in chat_span["metrics"]

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper agent_run_stream_sync span must not log token metrics.
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]


@pytest.mark.vcr
Expand DownExpand Up@@ -1183,6 +1238,14 @@ async def test_agent_run_stream_events(memory_logger):
 assert agent_span["metrics"]["event_count"] == event_count
 _assert_metrics_are_valid(agent_span["metrics"], start, end)

 # Regression: wrapper agent_run_stream_events span must NOT log token metrics.
 # Trace-tree rollup over non-scorer spans plus experiment-level token sums would
 # otherwise double-count against the leaf chat span.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in agent_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )


@pytest.mark.vcr
def test_direct_model_request_stream_sync(memory_logger, direct):
Expand DownExpand Up@@ -1223,6 +1286,15 @@ def test_direct_model_request_stream_sync(memory_logger, direct):
 assert ttft <= duration, f"time_to_first_token ({ttft}s) should be <= duration ({duration}s)"
 assert ttft < 3.0, f"time_to_first_token should be < 3s for API call, got {ttft}s"

 # Regression: model_request_stream_sync wrapper span must NOT log token metrics.
 # `_DirectStreamWrapperSync.__exit__` always uses `_wrapper_span_metrics`; sync
 # direct streaming has no leaf chat span, so this is the only span we emit and
 # token totals would otherwise be derived twice once a leaf is added.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 print(f"✓ Direct sync stream time_to_first_token: {ttft}s (duration: {duration}s)")


Expand DownExpand Up@@ -2030,53 +2102,40 @@ def helper_tool() -> str:
def test_reasoning_tokens_extraction(memory_logger):
 """Test that reasoning tokens are extracted from model responses.

 For reasoning models like o1/o3, usage.details.reasoning_tokens should be
 captured in the metrics field.
 For reasoning models like o1/o3, providers (e.g. OpenAI) populate
 `usage.details["reasoning_tokens"]` (a dict entry, not an attribute) and
 `_extract_response_metrics` must surface it as `completion_reasoning_tokens`.
 """
 assert not memory_logger.pop()

 # Mock a response that has reasoning tokens
 from unittest.mock import MagicMock
 from types import SimpleNamespace

 # Create a mock response with reasoning tokens
 mock_response = MagicMock()
 mock_response.parts = [
 MagicMock(
 part_kind="thinking",
 content="Let me think about this...",
 ),
 MagicMock(
 part_kind="text",
 content="The answer is 42",
 ),
 ]
 mock_response.usage = MagicMock()
 mock_response.usage.input_tokens = 10
 mock_response.usage.output_tokens = 20
 mock_response.usage.total_tokens = 30
 mock_response.usage.cache_read_tokens = 0
 mock_response.usage.cache_write_tokens = 0
 mock_response.usage.details = MagicMock()
 mock_response.usage.details.reasoning_tokens = 128

 # Test the metric extraction function directly
 from braintrust.integrations.pydantic_ai.tracing import _extract_response_metrics
 from pydantic_ai.usage import RequestUsage

 # Use a real RequestUsage so `details` is the actual dict[str, int] shape
 # pydantic_ai produces. A MagicMock would falsely satisfy attribute-style
 # lookups and hide real-world dict-only behavior.
 usage = RequestUsage(
 input_tokens=10,
 output_tokens=20,
 cache_read_tokens=0,
 cache_write_tokens=0,
 details={"reasoning_tokens": 128},
 )
 response = SimpleNamespace(parts=[], usage=usage)

 start_time = time.time()
 end_time = start_time + 5.0

 metrics = _extract_response_metrics(mock_response, start_time, end_time)
 metrics = _extract_response_metrics(response, start_time, end_time)

 # Verify all metrics are present
 assert metrics is not None, "Should extract metrics"
 # pylint: disable=unsupported-membership-test,unsubscriptable-object
 assert "prompt_tokens" in metrics, "Should have prompt_tokens"
 assert metrics["prompt_tokens"] == 10.0
 assert "completion_tokens" in metrics, "Should have completion_tokens"
 assert metrics["completion_tokens"] == 20.0
 assert "tokens" in metrics, "Should have total tokens"
 assert metrics["tokens"] == 30.0
 assert "completion_reasoning_tokens" in metrics, "Should have completion_reasoning_tokens"
 assert metrics["tokens"] == 30.0 # total_tokens is a property: input + output
 assert metrics["completion_reasoning_tokens"] == 128.0, (
 f"Expected 128.0, got {metrics['completion_reasoning_tokens']}"
 )
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Universal Dark Mode - works on any site
(function() {
var enabled = true;
function applyDarkMode() {
if (!enabled) return;
// Create style element if it doesn't exist
var style = document.getElementById('universal-dark-mode-style');
if (!style) {
style = document.createElement('style');
style.id = 'universal-dark-mode-style';
document.head.appendChild(style);
}
// Dark mode CSS - inverts colors but preserves images/video
style.textContent = '
/* Invert everything except media */
html {
filter: invert(1) hue-rotate(180deg) !important;
background: #1a1a2e !important;
}
/* Restore images, videos, iframes, canvas */
img, video, iframe, canvas, svg, picture, [style*="background-image"] {
filter: invert(1) hue-rotate(180deg) !important;
}
/* Preserve specific elements that should not be inverted */
.no-dark-mode, .no-dark-mode *,
[data-theme="light"], [data-theme="light"],
.ace_editor, .ace_editor *,
.CodeMirror, .CodeMirror *,
.monaco-editor, .monaco-editor *,
.markdown-body pre, .markdown-body pre *,
.highlight, .highlight *,
pre code, pre code * {
filter: none !important;
}
/* Fix common UI elements */
.modal, .popup, .dropdown-menu, .tooltip, .popover {
filter: invert(1) hue-rotate(180deg) !important;
background: #2d2d44 !important;
border-color: #444 !important;
}
/* Scrollbars */
::-webkit-scrollbar { background: #1a1a2e !important; }
::-webkit-scrollbar-thumb { background: #444 !important; }
::-webkit-scrollbar-thumb:hover { background: #555 !important; }
/* Selection */
::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
';
}
function removeDarkMode() {
var style = document.getElementById('universal-dark-mode-style');
if (style) style.remove();
}
// Toggle with Alt+Shift+D
document.addEventListener('keydown', function(e) {
if (e.altKey && e.shiftKey && e.key === 'D') {
e.preventDefault();
enabled = !enabled;
if (enabled) {
applyDarkMode();
console.log('[Universal Dark Mode] Enabled');
} else {
removeDarkMode();
console.log('[Universal Dark Mode] Disabled');
}
}
});
// Apply on load
applyDarkMode();
// Re-apply on dynamic content
var observer = new MutationObserver(function(mutations) {
if (enabled && !document.getElementById('universal-dark-mode-style')) {
applyDarkMode();
}
});
observer.observe(document.head, { childList: true });
console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');
})();
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();
fix(pydantic_ai): strip unnecessary wrapper tokens by colinbennettbrain · Pull Request #321 · braintrustdata/braintrust-sdk-python · GitHub
Skip to content
fix(pydantic_ai): strip unnecessary wrapper tokens#321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

Abhijeet Prasad (AbhiPrasad)
merged 5 commits into

from

Apr 20, 2026
 Merged


fix(pydantic_ai): strip unnecessary wrapper tokens
#321
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
6a6cb76

fix(pydantic_ai): strip token metrics from wrapper spans


colinbennettbrain Apr 16, 2026
61be582

fix(pydantic_ai): restore audio + dict-style usage details on leaf span


colinbennettbrain Apr 17, 2026
171897c

Merge branch 'main' into colin/pydantic-strip-wrapper-tokens


colinbennettbrain Apr 17, 2026
246fcc2

fix(pydantic-ai): narrow Optional metrics type for pylint

Apr 17, 2026
39c970a

fix(py): fix pylint in pydantic ai metrics test


colinbennettbrain Apr 17, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


173 changes: 116 additions & 57 deletions
py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -67,6 +67,46 @@ def _assert_metrics_are_valid(metrics, start, end):
 assert metrics["completion_tokens"] > 0


def test_extract_response_metrics_leaf_fields():
 """`_extract_response_metrics` is the only path that populates token metrics on
 the leaf chat span post-#312/strip-wrapper-tokens. Verify the audio-token and
 dict-style `details` fields that were dropped when `_extract_usage_metrics`
 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

 usage = RequestUsage(
 input_tokens=100,
 output_tokens=50,
 cache_read_tokens=10,
 cache_write_tokens=20,
 input_audio_tokens=7,
 output_audio_tokens=11,
 details={"reasoning_tokens": 13, "cached_tokens": 17},
 )
 response = SimpleNamespace(usage=usage)

 maybe_metrics = _extract_response_metrics(response, start_time=1.0, end_time=2.0)
 assert maybe_metrics is not None

 # pylint: disable=unsupported-membership-test,unsubscriptable-object
 assert maybe_metrics["prompt_tokens"] == 100
 assert maybe_metrics["completion_tokens"] == 50
 assert maybe_metrics["tokens"] == 150 # total_tokens is a property: input + output
 assert maybe_metrics["prompt_cache_creation_tokens"] == 20
 assert maybe_metrics["prompt_audio_tokens"] == 7
 assert maybe_metrics["completion_audio_tokens"] == 11
 assert maybe_metrics["completion_reasoning_tokens"] == 13
 # details["cached_tokens"] is read after cache_read_tokens, so it wins when both
 # are populated; this preserves the pre-strip behavior of `_extract_usage_metrics`.
 assert maybe_metrics["prompt_cached_tokens"] == 17
 # pylint: enable=unsupported-membership-test,unsubscriptable-object


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_direct_model_request_creates_nested_chat_span_without_class_scan(memory_logger, direct):
Expand DownExpand Up@@ -140,23 +180,21 @@ async def test_agent_run_async(memory_logger):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 assert agent_span["metrics"]["prompt_tokens"] > 0
 assert agent_span["metrics"]["completion_tokens"] > 0
 # Regression: wrapper agent_run span must NOT log token metrics. The leaf chat
 # span already logs them, and trace-tree rollup (self + descendants) plus any
 # unfiltered sum over metrics would otherwise double-count tokens/cost at the
 # parent regardless of span type.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in agent_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 # Regression: no double-counting of cost/tokens. Experiment-level aggregations
 # sum metrics across type='llm' spans, so a single agent turn must contribute
 # its tokens exactly once. The wrapper agent_run span logs the same usage as
 # the leaf chat span; only the leaf should be type=LLM.
 # Only the leaf chat span should be type=LLM, and it must own the token totals.
 llm_spans = [s for s in spans if s["span_attributes"]["type"] == SpanTypeAttribute.LLM]
 assert len(llm_spans) == 1, f"expected exactly one LLM-typed span, got {len(llm_spans)}"
 assert llm_spans[0]["span_id"] == chat_span["span_id"]
 llm_prompt_tokens_sum = sum(s["metrics"].get("prompt_tokens", 0) for s in llm_spans)
 llm_completion_tokens_sum = sum(s["metrics"].get("completion_tokens", 0) for s in llm_spans)
 assert llm_prompt_tokens_sum == chat_span["metrics"]["prompt_tokens"]
 assert llm_completion_tokens_sum == chat_span["metrics"]["completion_tokens"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -245,9 +283,12 @@ def is_descendant(child_span, ancestor_id):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_sync_span["metrics"]
 assert "completion_tokens" in agent_sync_span["metrics"]
 # Wrapper agent_run_sync span must not log token metrics (would double-count at rollup).
 assert "prompt_tokens" not in agent_sync_span["metrics"]
 assert "completion_tokens" not in agent_sync_span["metrics"]
 # Tokens live on the leaf chat span only.
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


def test_agent_to_cli_sync(memory_logger, monkeypatch):
Expand DownExpand Up@@ -544,9 +585,12 @@ async def test_agent_run_stream(memory_logger):
 print(f"span_parents: {chat_span['span_parents']}")
 print(f"metrics: {chat_span['metrics']}")

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper stream span must not log token metrics (would double-count at rollup).
 # time_to_first_token is asserted above; it's a non-summable timing metric and stays.
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -741,6 +785,15 @@ async def test_direct_model_request_stream(memory_logger, direct):
 assert ttft <= duration, f"time_to_first_token ({ttft}s) should be <= duration ({duration}s)"
 assert ttft < 3.0, f"time_to_first_token should be < 3s for API call, got {ttft}s"

 # Regression: model_request_stream wrapper span (TASK) must NOT log token metrics.
 # `_DirectStreamWrapper` here is constructed with span_type=TASK, so __aexit__
 # routes to `_wrapper_span_metrics`. The leaf chat span (when present) is the only
 # place tokens should be recorded.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in direct_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 print(f"✓ Direct stream time_to_first_token: {ttft}s (duration: {duration}s)")


Expand DownExpand Up@@ -842,9 +895,11 @@ def is_descendant(child_span, ancestor_id):
 assert chat_span["metadata"]["provider"] == "openai"
 _assert_metrics_are_valid(chat_span["metrics"], start, end)

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper agent_run span must not log token metrics (would double-count at rollup).
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]
 assert chat_span["metrics"]["prompt_tokens"] > 0
 assert chat_span["metrics"]["completion_tokens"] > 0


@pytest.mark.vcr
Expand DownExpand Up@@ -1143,9 +1198,9 @@ def is_descendant(child_span, ancestor_id):
 # Chat span may not have complete metrics since it's an intermediate span
 assert "start" in chat_span["metrics"]

 # Agent spans should have token metrics
 assert "prompt_tokens" in agent_span["metrics"]
 assert "completion_tokens" in agent_span["metrics"]
 # Wrapper agent_run_stream_sync span must not log token metrics.
 assert "prompt_tokens" not in agent_span["metrics"]
 assert "completion_tokens" not in agent_span["metrics"]


@pytest.mark.vcr
Expand DownExpand Up@@ -1183,6 +1238,14 @@ async def test_agent_run_stream_events(memory_logger):
 assert agent_span["metrics"]["event_count"] == event_count
 _assert_metrics_are_valid(agent_span["metrics"], start, end)

 # Regression: wrapper agent_run_stream_events span must NOT log token metrics.
 # Trace-tree rollup over non-scorer spans plus experiment-level token sums would
 # otherwise double-count against the leaf chat span.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in agent_span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )


@pytest.mark.vcr
def test_direct_model_request_stream_sync(memory_logger, direct):
Expand DownExpand Up@@ -1223,6 +1286,15 @@ def test_direct_model_request_stream_sync(memory_logger, direct):
 assert ttft <= duration, f"time_to_first_token ({ttft}s) should be <= duration ({duration}s)"
 assert ttft < 3.0, f"time_to_first_token should be < 3s for API call, got {ttft}s"

 # Regression: model_request_stream_sync wrapper span must NOT log token metrics.
 # `_DirectStreamWrapperSync.__exit__` always uses `_wrapper_span_metrics`; sync
 # direct streaming has no leaf chat span, so this is the only span we emit and
 # token totals would otherwise be derived twice once a leaf is added.
 for token_key in ("prompt_tokens", "completion_tokens", "tokens", "prompt_cached_tokens"):
 assert token_key not in span["metrics"], (
 f"wrapper span must not log {token_key}; it duplicates the leaf chat span"
 )

 print(f"✓ Direct sync stream time_to_first_token: {ttft}s (duration: {duration}s)")


Expand DownExpand Up@@ -2030,53 +2102,40 @@ def helper_tool() -> str:
def test_reasoning_tokens_extraction(memory_logger):
 """Test that reasoning tokens are extracted from model responses.

 For reasoning models like o1/o3, usage.details.reasoning_tokens should be
 captured in the metrics field.
 For reasoning models like o1/o3, providers (e.g. OpenAI) populate
 `usage.details["reasoning_tokens"]` (a dict entry, not an attribute) and
 `_extract_response_metrics` must surface it as `completion_reasoning_tokens`.
 """
 assert not memory_logger.pop()

 # Mock a response that has reasoning tokens
 from unittest.mock import MagicMock
 from types import SimpleNamespace

 # Create a mock response with reasoning tokens
 mock_response = MagicMock()
 mock_response.parts = [
 MagicMock(
 part_kind="thinking",
 content="Let me think about this...",
 ),
 MagicMock(
 part_kind="text",
 content="The answer is 42",
 ),
 ]
 mock_response.usage = MagicMock()
 mock_response.usage.input_tokens = 10
 mock_response.usage.output_tokens = 20
 mock_response.usage.total_tokens = 30
 mock_response.usage.cache_read_tokens = 0
 mock_response.usage.cache_write_tokens = 0
 mock_response.usage.details = MagicMock()
 mock_response.usage.details.reasoning_tokens = 128

 # Test the metric extraction function directly
 from braintrust.integrations.pydantic_ai.tracing import _extract_response_metrics
 from pydantic_ai.usage import RequestUsage

 # Use a real RequestUsage so `details` is the actual dict[str, int] shape
 # pydantic_ai produces. A MagicMock would falsely satisfy attribute-style
 # lookups and hide real-world dict-only behavior.
 usage = RequestUsage(
 input_tokens=10,
 output_tokens=20,
 cache_read_tokens=0,
 cache_write_tokens=0,
 details={"reasoning_tokens": 128},
 )
 response = SimpleNamespace(parts=[], usage=usage)

 start_time = time.time()
 end_time = start_time + 5.0

 metrics = _extract_response_metrics(mock_response, start_time, end_time)
 metrics = _extract_response_metrics(response, start_time, end_time)

 # Verify all metrics are present
 assert metrics is not None, "Should extract metrics"
 # pylint: disable=unsupported-membership-test,unsubscriptable-object
 assert "prompt_tokens" in metrics, "Should have prompt_tokens"
 assert metrics["prompt_tokens"] == 10.0
 assert "completion_tokens" in metrics, "Should have completion_tokens"
 assert metrics["completion_tokens"] == 20.0
 assert "tokens" in metrics, "Should have total tokens"
 assert metrics["tokens"] == 30.0
 assert "completion_reasoning_tokens" in metrics, "Should have completion_reasoning_tokens"
 assert metrics["tokens"] == 30.0 # total_tokens is a property: input + output
 assert metrics["completion_reasoning_tokens"] == 128.0, (
 f"Expected 128.0, got {metrics['completion_reasoning_tokens']}"
 )
Expand Down
Loading