Skip to content

Commit 908f2cf

Browse files
committed
Merge branch 'master' into feat/span-first
2 parents 1447ac2 + a76280b commit 908f2cf

File tree

13 files changed

+221
-68
lines changed

13 files changed

+221
-68
lines changed

.github/release.yml

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,8 @@
1+
# This configuration is used by Craft to categorize changelog entries based on
2+
# PR labels. To avoid some manual work, there is a PR labeling GitHub action in
3+
# .github/workflows/pr-labeler.yml that adds a changelog label to PRs based on
4+
# the title.
5+
16
changelog:
27
exclude:
38
labels:
@@ -15,6 +20,9 @@ changelog:
1520
labels:
1621
- "Changelog: Bugfix"
1722
- Bug
23+
- title: Deprecations 🏗️
24+
labels:
25+
- "Changelog: Deprecation"
1826
- title: Documentation 📚
1927
labels:
2028
- "Changelog: Docs"

.github/workflows/pr-labeler.yml

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
# This action adds changelog labels to PRs that are then used by the release
2+
# notes generator in Craft. The configuration for which labels map to what
3+
# changelog categories can be found in .github/release.yml.
4+
5+
name: Label PR for Changelog
6+
7+
on:
8+
pull_request:
9+
types: [opened, edited]
10+
11+
permissions:
12+
pull-requests: write
13+
14+
jobs:
15+
label:
16+
runs-on: ubuntu-latest
17+
steps:
18+
- name: Add changelog label
19+
uses: actions/github-script@v7
20+
with:
21+
script: |
22+
const title = context.payload.pull_request.title.toLowerCase();
23+
const prNumber = context.payload.pull_request.number;
24+
25+
// Get current labels
26+
const { data: currentLabels } = await github.rest.issues.listLabelsOnIssue({
27+
owner: context.repo.owner,
28+
repo: context.repo.repo,
29+
issue_number: prNumber,
30+
});
31+
32+
// Check if a Changelog label already exists
33+
const hasChangelogLabel = currentLabels.some(label =>
34+
label.name.startsWith('Changelog:') || label.name === 'skip-changelog'
35+
);
36+
37+
if (hasChangelogLabel) {
38+
console.log('PR already has a Changelog label, skipping');
39+
return;
40+
}
41+
42+
// Determine which label to apply
43+
let newLabel = null;
44+
45+
if (title.includes('deprecate')) {
46+
newLabel = 'Changelog: Deprecation';
47+
} else if (title.startsWith('feat')) {
48+
newLabel = 'Changelog: Feature';
49+
} else if (title.startsWith('fix') || title.startsWith('bugfix')) {
50+
newLabel = 'Changelog: Bugfix';
51+
} else if (title.startsWith('docs')) {
52+
newLabel = 'Changelog: Docs';
53+
} else if (title.startsWith('ref') || title.startsWith('test')) {
54+
newLabel = 'Changelog: Internal';
55+
} else if (title.startsWith('ci') || title.startsWith('build')) {
56+
newLabel = 'skip-changelog';
57+
}
58+
59+
// Apply the new label if one was determined
60+
if (newLabel) {
61+
await github.rest.issues.addLabels({
62+
owner: context.repo.owner,
63+
repo: context.repo.repo,
64+
issue_number: prNumber,
65+
labels: [newLabel],
66+
});
67+
68+
console.log(`Applied label: ${newLabel}`);
69+
} else {
70+
console.log('No matching label pattern found in PR title, please add manually');
71+
}
72+

sentry_sdk/integrations/grpc/aio/server.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ async def wrapped(request, context):
4444
return await handler(request, context)
4545

4646
# What if the headers are empty?
47-
transaction = Transaction.continue_from_headers(
47+
transaction = sentry_sdk.continue_trace(
4848
dict(context.invocation_metadata()),
4949
op=OP.GRPC_SERVER,
5050
name=name,

sentry_sdk/integrations/grpc/server.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ def behavior(request, context):
3838
if name:
3939
metadata = dict(context.invocation_metadata())
4040

41-
transaction = Transaction.continue_from_headers(
41+
transaction = sentry_sdk.continue_trace(
4242
metadata,
4343
op=OP.GRPC_SERVER,
4444
name=name,

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -71,13 +71,9 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
7171
# Exit the original context manager first
7272
await self.original_ctx_manager.__aexit__(exc_type, exc_val, exc_tb)
7373

74-
# Update span with output if successful
75-
if exc_type is None and self._result and hasattr(self._result, "output"):
76-
output = (
77-
self._result.output if hasattr(self._result, "output") else None
78-
)
79-
if self._span is not None:
80-
update_invoke_agent_span(self._span, output)
74+
# Update span with result if successful
75+
if exc_type is None and self._result and self._span is not None:
76+
update_invoke_agent_span(self._span, self._result)
8177
finally:
8278
# Pop agent from contextvar stack
8379
pop_agent()
@@ -123,9 +119,8 @@ async def wrapper(self, *args, **kwargs):
123119
try:
124120
result = await original_func(self, *args, **kwargs)
125121

126-
# Update span with output
127-
output = result.output if hasattr(result, "output") else None
128-
update_invoke_agent_span(span, output)
122+
# Update span with result
123+
update_invoke_agent_span(span, result)
129124

130125
return result
131126
except Exception as exc:

sentry_sdk/integrations/pydantic_ai/spans/ai_client.py

Lines changed: 1 addition & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
get_current_agent,
1414
get_is_streaming,
1515
)
16+
from .utils import _set_usage_data
1617

1718
from typing import TYPE_CHECKING
1819

@@ -39,22 +40,6 @@
3940
ThinkingPart = None
4041

4142

42-
def _set_usage_data(span, usage):
43-
# type: (sentry_sdk.tracing.Span, RequestUsage) -> None
44-
"""Set token usage data on a span."""
45-
if usage is None:
46-
return
47-
48-
if hasattr(usage, "input_tokens") and usage.input_tokens is not None:
49-
span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens)
50-
51-
if hasattr(usage, "output_tokens") and usage.output_tokens is not None:
52-
span.set_data(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage.output_tokens)
53-
54-
if hasattr(usage, "total_tokens") and usage.total_tokens is not None:
55-
span.set_data(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens)
56-
57-
5843
def _set_input_messages(span, messages):
5944
# type: (sentry_sdk.tracing.Span, Any) -> None
6045
"""Set input messages data on a span."""

sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
_set_model_data,
1010
_should_send_prompts,
1111
)
12+
from .utils import _set_usage_data
1213

1314
from typing import TYPE_CHECKING
1415

@@ -103,10 +104,37 @@ def invoke_agent_span(user_prompt, agent, model, model_settings, is_streaming=Fa
103104
return span
104105

105106

106-
def update_invoke_agent_span(span, output):
107+
def update_invoke_agent_span(span, result):
107108
# type: (sentry_sdk.tracing.Span, Any) -> None
108109
"""Update and close the invoke agent span."""
109-
if span and _should_send_prompts() and output:
110+
if not span or not result:
111+
return
112+
113+
# Extract output from result
114+
output = getattr(result, "output", None)
115+
116+
# Set response text if prompts are enabled
117+
if _should_send_prompts() and output:
110118
set_data_normalized(
111119
span, SPANDATA.GEN_AI_RESPONSE_TEXT, str(output), unpack=False
112120
)
121+
122+
# Set token usage data if available
123+
if hasattr(result, "usage") and callable(result.usage):
124+
try:
125+
usage = result.usage()
126+
if usage:
127+
_set_usage_data(span, usage)
128+
except Exception:
129+
# If usage() call fails, continue without setting usage data
130+
pass
131+
132+
# Set model name from response if available
133+
if hasattr(result, "response"):
134+
try:
135+
response = result.response
136+
if hasattr(response, "model_name") and response.model_name:
137+
span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name)
138+
except Exception:
139+
# If response access fails, continue without setting model name
140+
pass
Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
"""Utility functions for PydanticAI span instrumentation."""
2+
3+
import sentry_sdk
4+
from sentry_sdk.consts import SPANDATA
5+
6+
from typing import TYPE_CHECKING
7+
8+
if TYPE_CHECKING:
9+
from typing import Union
10+
from pydantic_ai.usage import RequestUsage, RunUsage # type: ignore
11+
12+
13+
def _set_usage_data(span, usage):
14+
# type: (sentry_sdk.tracing.Span, Union[RequestUsage, RunUsage]) -> None
15+
"""Set token usage data on a span.
16+
17+
This function works with both RequestUsage (single request) and
18+
RunUsage (agent run) objects from pydantic_ai.
19+
20+
Args:
21+
span: The Sentry span to set data on.
22+
usage: RequestUsage or RunUsage object containing token usage information.
23+
"""
24+
if usage is None:
25+
return
26+
27+
if hasattr(usage, "input_tokens") and usage.input_tokens is not None:
28+
span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens)
29+
30+
if hasattr(usage, "output_tokens") and usage.output_tokens is not None:
31+
span.set_data(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage.output_tokens)
32+
33+
if hasattr(usage, "total_tokens") and usage.total_tokens is not None:
34+
span.set_data(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens)

sentry_sdk/tracing.py

Lines changed: 6 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -485,6 +485,8 @@ def continue_from_environ(
485485
):
486486
# type: (...) -> Transaction
487487
"""
488+
DEPRECATED: Use :py:meth:`sentry_sdk.continue_trace`.
489+
488490
Create a Transaction with the given params, then add in data pulled from
489491
the ``sentry-trace`` and ``baggage`` headers from the environ (if any)
490492
before returning the Transaction.
@@ -496,11 +498,6 @@ def continue_from_environ(
496498
497499
:param environ: The ASGI/WSGI environ to pull information from.
498500
"""
499-
if cls is Span:
500-
logger.warning(
501-
"Deprecated: use Transaction.continue_from_environ "
502-
"instead of Span.continue_from_environ."
503-
)
504501
return Transaction.continue_from_headers(EnvironHeaders(environ), **kwargs)
505502

506503
@classmethod
@@ -513,19 +510,16 @@ def continue_from_headers(
513510
):
514511
# type: (...) -> Transaction
515512
"""
513+
DEPRECATED: Use :py:meth:`sentry_sdk.continue_trace`.
514+
516515
Create a transaction with the given params (including any data pulled from
517516
the ``sentry-trace`` and ``baggage`` headers).
518517
519518
:param headers: The dictionary with the HTTP headers to pull information from.
520519
:param _sample_rand: If provided, we override the sample_rand value from the
521520
incoming headers with this value. (internal use only)
522521
"""
523-
# TODO move this to the Transaction class
524-
if cls is Span:
525-
logger.warning(
526-
"Deprecated: use Transaction.continue_from_headers "
527-
"instead of Span.continue_from_headers."
528-
)
522+
logger.warning("Deprecated: use sentry_sdk.continue_trace instead.")
529523

530524
# TODO-neel move away from this kwargs stuff, it's confusing and opaque
531525
# make more explicit
@@ -579,16 +573,11 @@ def from_traceparent(
579573
):
580574
# type: (...) -> Optional[Transaction]
581575
"""
582-
DEPRECATED: Use :py:meth:`sentry_sdk.tracing.Span.continue_from_headers`.
576+
DEPRECATED: Use :py:meth:`sentry_sdk.continue_trace`.
583577
584578
Create a ``Transaction`` with the given params, then add in data pulled from
585579
the given ``sentry-trace`` header value before returning the ``Transaction``.
586580
"""
587-
logger.warning(
588-
"Deprecated: Use Transaction.continue_from_headers(headers, **kwargs) "
589-
"instead of from_traceparent(traceparent, **kwargs)"
590-
)
591-
592581
if not traceparent:
593582
return None
594583

tests/integrations/pydantic_ai/test_pydantic_ai.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,51 @@ async def test_agent_run_async(sentry_init, capture_events, test_agent):
7676
assert "gen_ai.usage.output_tokens" in chat_span["data"]
7777

7878

79+
@pytest.mark.asyncio
80+
async def test_agent_run_async_usage_data(sentry_init, capture_events, test_agent):
81+
"""
82+
Test that the invoke_agent span includes token usage and model data.
83+
"""
84+
sentry_init(
85+
integrations=[PydanticAIIntegration()],
86+
traces_sample_rate=1.0,
87+
send_default_pii=True,
88+
)
89+
90+
events = capture_events()
91+
92+
result = await test_agent.run("Test input")
93+
94+
assert result is not None
95+
assert result.output is not None
96+
97+
(transaction,) = events
98+
99+
# Verify transaction (the transaction IS the invoke_agent span)
100+
assert transaction["transaction"] == "invoke_agent test_agent"
101+
102+
# The invoke_agent span should have token usage data
103+
trace_data = transaction["contexts"]["trace"].get("data", {})
104+
assert "gen_ai.usage.input_tokens" in trace_data, (
105+
"Missing input_tokens on invoke_agent span"
106+
)
107+
assert "gen_ai.usage.output_tokens" in trace_data, (
108+
"Missing output_tokens on invoke_agent span"
109+
)
110+
assert "gen_ai.usage.total_tokens" in trace_data, (
111+
"Missing total_tokens on invoke_agent span"
112+
)
113+
assert "gen_ai.response.model" in trace_data, (
114+
"Missing response.model on invoke_agent span"
115+
)
116+
117+
# Verify the values are reasonable
118+
assert trace_data["gen_ai.usage.input_tokens"] > 0
119+
assert trace_data["gen_ai.usage.output_tokens"] > 0
120+
assert trace_data["gen_ai.usage.total_tokens"] > 0
121+
assert trace_data["gen_ai.response.model"] == "test" # Test model name
122+
123+
79124
def test_agent_run_sync(sentry_init, capture_events, test_agent):
80125
"""
81126
Test that the integration creates spans for sync agent runs.

0 commit comments

Comments
 (0)