Skip to content

Commit c686455

Browse files
alexander-alderman-webbalexander-alderman-webb
authored
feat(pydantic-ai): Support span streaming (getsentry#6389)
1 parent 4e30693 commit c686455

7 files changed

Lines changed: 737 additions & 149 deletions

File tree

‎sentry_sdk/integrations/pydantic_ai/patches/agent_run.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
raise DidNotEnable("pydantic-ai not installed")
1616

1717
if TYPE_CHECKING:
18-
from typing import Any, Callable, Optional
18+
from typing import Any, Callable, Optional, Union
1919

2020

2121
class _StreamingContextManagerWrapper:
@@ -37,7 +37,7 @@ def __init__(
3737
self.model_settings = model_settings
3838
self.is_streaming = is_streaming
3939
self._isolation_scope: "Any" = None
40-
self._span: "Optional[sentry_sdk.tracing.Span]" = None
40+
self._span: "Optional[Union[sentry_sdk.tracing.Span, sentry_sdk.traces.StreamedSpan]]" = None
4141
self._result: "Any" = None
4242

4343
async def __aenter__(self) -> "Any":

‎sentry_sdk/integrations/pydantic_ai/spans/ai_client.py‎

Lines changed: 55 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
truncate_and_annotate_messages,
99
)
1010
from sentry_sdk.consts import OP, SPANDATA
11+
from sentry_sdk.traces import StreamedSpan
12+
from sentry_sdk.tracing_utils import has_span_streaming_enabled
1113
from sentry_sdk.utils import safe_serialize
1214

1315
from ..consts import SPAN_ORIGIN
@@ -27,7 +29,7 @@
2729
)
2830

2931
if TYPE_CHECKING:
30-
from typing import Any, Dict, List
32+
from typing import Any, Dict, List, Union
3133

3234
from pydantic_ai.messages import ModelMessage, SystemPromptPart # type: ignore
3335

@@ -97,7 +99,9 @@ def _get_system_instructions(
9799
return permanent_instructions, current_instructions
98100

99101

100-
def _set_input_messages(span: "sentry_sdk.tracing.Span", messages: "Any") -> None:
102+
def _set_input_messages(
103+
span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", messages: "Any"
104+
) -> None:
101105
"""Set input messages data on a span."""
102106
if not _should_send_prompts():
103107
return
@@ -107,14 +111,24 @@ def _set_input_messages(span: "sentry_sdk.tracing.Span", messages: "Any") -> Non
107111

108112
permanent_instructions, current_instructions = _get_system_instructions(messages)
109113
if len(permanent_instructions) > 0 or len(current_instructions) > 0:
110-
span.set_data(
111-
SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
112-
json.dumps(
113-
_transform_system_instructions(
114-
permanent_instructions, current_instructions
115-
)
116-
),
117-
)
114+
if isinstance(span, StreamedSpan):
115+
span.set_attribute(
116+
SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
117+
json.dumps(
118+
_transform_system_instructions(
119+
permanent_instructions, current_instructions
120+
)
121+
),
122+
)
123+
else:
124+
span.set_data(
125+
SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
126+
json.dumps(
127+
_transform_system_instructions(
128+
permanent_instructions, current_instructions
129+
)
130+
),
131+
)
118132

119133
try:
120134
formatted_messages = []
@@ -198,15 +212,21 @@ def _set_input_messages(span: "sentry_sdk.tracing.Span", messages: "Any") -> Non
198212
pass
199213

200214

201-
def _set_output_data(span: "sentry_sdk.tracing.Span", response: "Any") -> None:
215+
def _set_output_data(
216+
span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", response: "Any"
217+
) -> None:
202218
"""Set output data on a span."""
203219
if not _should_send_prompts():
204220
return
205221

206222
if not response:
207223
return
208224

209-
span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name)
225+
set_on_span = (
226+
span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
227+
)
228+
set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name)
229+
210230
try:
211231
# Extract text from ModelResponse
212232
if hasattr(response, "parts"):
@@ -230,7 +250,7 @@ def _set_output_data(span: "sentry_sdk.tracing.Span", response: "Any") -> None:
230250
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, texts)
231251

232252
if tool_calls:
233-
span.set_data(
253+
set_on_span(
234254
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls)
235255
)
236256

@@ -241,7 +261,7 @@ def _set_output_data(span: "sentry_sdk.tracing.Span", response: "Any") -> None:
241261

242262
def ai_client_span(
243263
messages: "Any", agent: "Any", model: "Any", model_settings: "Any"
244-
) -> "sentry_sdk.tracing.Span":
264+
) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]":
245265
"""Create a span for an AI client call (model request).
246266
247267
Args:
@@ -257,20 +277,31 @@ def ai_client_span(
257277

258278
model_name = _get_model_name(model_obj) or "unknown"
259279

260-
span = sentry_sdk.start_span(
261-
op=OP.GEN_AI_CHAT,
262-
name=f"chat {model_name}",
263-
origin=SPAN_ORIGIN,
264-
)
280+
span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
281+
if span_streaming:
282+
span = sentry_sdk.traces.start_span(
283+
name=f"chat {model_name}",
284+
attributes={
285+
"sentry.op": OP.GEN_AI_CHAT,
286+
"sentry.origin": SPAN_ORIGIN,
287+
SPANDATA.GEN_AI_OPERATION_NAME: "chat",
288+
SPANDATA.GEN_AI_RESPONSE_STREAMING: get_is_streaming(),
289+
},
290+
)
291+
else:
292+
span = sentry_sdk.start_span(
293+
op=OP.GEN_AI_CHAT,
294+
name=f"chat {model_name}",
295+
origin=SPAN_ORIGIN,
296+
)
265297

266-
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
298+
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
299+
# Set streaming flag from contextvar
300+
span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, get_is_streaming())
267301

268302
_set_agent_data(span, agent)
269303
_set_model_data(span, model, model_settings)
270304

271-
# Set streaming flag from contextvar
272-
span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, get_is_streaming())
273-
274305
# Add available tools if agent is available
275306
agent_obj = agent or get_current_agent()
276307
_set_available_tools(span, agent_obj)
@@ -283,7 +314,7 @@ def ai_client_span(
283314

284315

285316
def update_ai_client_span(
286-
span: "sentry_sdk.tracing.Span", model_response: "Any"
317+
span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", model_response: "Any"
287318
) -> None:
288319
"""Update the AI client span with response data."""
289320
if not span:

‎sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py‎

Lines changed: 38 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,15 @@
22

33
import sentry_sdk
44
from sentry_sdk.consts import OP, SPANDATA
5+
from sentry_sdk.traces import StreamedSpan
6+
from sentry_sdk.tracing_utils import has_span_streaming_enabled
57
from sentry_sdk.utils import safe_serialize
68

79
from ..consts import SPAN_ORIGIN
810
from ..utils import _set_agent_data, _should_send_prompts
911

1012
if TYPE_CHECKING:
11-
from typing import Any, Optional
13+
from typing import Any, Optional, Union
1214

1315
from pydantic_ai._tool_manager import ToolDefinition # type: ignore
1416

@@ -18,7 +20,7 @@ def execute_tool_span(
1820
tool_args: "Any",
1921
agent: "Any",
2022
tool_definition: "Optional[ToolDefinition]" = None,
21-
) -> "sentry_sdk.tracing.Span":
23+
) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]":
2224
"""Create a span for tool execution.
2325
2426
Args:
@@ -27,33 +29,56 @@ def execute_tool_span(
2729
agent: The agent executing the tool
2830
tool_definition: The definition of the tool, if available
2931
"""
30-
span = sentry_sdk.start_span(
31-
op=OP.GEN_AI_EXECUTE_TOOL,
32-
name=f"execute_tool {tool_name}",
33-
origin=SPAN_ORIGIN,
34-
)
32+
span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
33+
if span_streaming:
34+
span = sentry_sdk.traces.start_span(
35+
name=f"execute_tool {tool_name}",
36+
attributes={
37+
"sentry.op": OP.GEN_AI_EXECUTE_TOOL,
38+
"sentry.origin": SPAN_ORIGIN,
39+
SPANDATA.GEN_AI_OPERATION_NAME: "execute_tool",
40+
SPANDATA.GEN_AI_TOOL_NAME: tool_name,
41+
},
42+
)
43+
44+
set_on_span = span.set_attribute
45+
else:
46+
span = sentry_sdk.start_span(
47+
op=OP.GEN_AI_EXECUTE_TOOL,
48+
name=f"execute_tool {tool_name}",
49+
origin=SPAN_ORIGIN,
50+
)
3551

36-
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "execute_tool")
37-
span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool_name)
52+
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "execute_tool")
53+
span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool_name)
54+
55+
set_on_span = span.set_data
3856

3957
if tool_definition is not None and hasattr(tool_definition, "description"):
40-
span.set_data(
58+
set_on_span(
4159
SPANDATA.GEN_AI_TOOL_DESCRIPTION,
4260
tool_definition.description,
4361
)
4462

4563
_set_agent_data(span, agent)
4664

4765
if _should_send_prompts() and tool_args is not None:
48-
span.set_data(SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_args))
66+
set_on_span(SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_args))
4967

5068
return span
5169

5270

53-
def update_execute_tool_span(span: "sentry_sdk.tracing.Span", result: "Any") -> None:
71+
def update_execute_tool_span(
72+
span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", result: "Any"
73+
) -> None:
5474
"""Update the execute tool span with the result."""
5575
if not span:
5676
return
5777

58-
if _should_send_prompts() and result is not None:
78+
if not _should_send_prompts() or result is None:
79+
return
80+
81+
if isinstance(span, StreamedSpan):
82+
span.set_attribute(SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result))
83+
else:
5984
span.set_data(SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result))

‎sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py‎

Lines changed: 31 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
truncate_and_annotate_messages,
99
)
1010
from sentry_sdk.consts import OP, SPANDATA
11+
from sentry_sdk.traces import StreamedSpan
12+
from sentry_sdk.tracing_utils import has_span_streaming_enabled
1113

1214
from ..consts import SPAN_ORIGIN
1315
from ..utils import (
@@ -22,7 +24,7 @@
2224
)
2325

2426
if TYPE_CHECKING:
25-
from typing import Any
27+
from typing import Any, Union
2628

2729
try:
2830
from pydantic_ai.messages import BinaryContent, ImageUrl # type: ignore
@@ -37,20 +39,31 @@ def invoke_agent_span(
3739
model: "Any",
3840
model_settings: "Any",
3941
is_streaming: bool = False,
40-
) -> "sentry_sdk.tracing.Span":
42+
) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]":
4143
"""Create a span for invoking the agent."""
4244
# Determine agent name for span
4345
name = "agent"
4446
if agent and getattr(agent, "name", None):
4547
name = agent.name
4648

47-
span = get_start_span_function()(
48-
op=OP.GEN_AI_INVOKE_AGENT,
49-
name=f"invoke_agent {name}",
50-
origin=SPAN_ORIGIN,
51-
)
49+
span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
50+
if span_streaming:
51+
span = sentry_sdk.traces.start_span(
52+
name=f"invoke_agent {name}",
53+
attributes={
54+
"sentry.op": OP.GEN_AI_INVOKE_AGENT,
55+
"sentry.origin": SPAN_ORIGIN,
56+
SPANDATA.GEN_AI_OPERATION_NAME: "invoke_agent",
57+
},
58+
)
59+
else:
60+
span = get_start_span_function()(
61+
op=OP.GEN_AI_INVOKE_AGENT,
62+
name=f"invoke_agent {name}",
63+
origin=SPAN_ORIGIN,
64+
)
5265

53-
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
66+
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
5467

5568
_set_agent_data(span, agent)
5669
_set_model_data(span, model, model_settings)
@@ -135,7 +148,10 @@ def invoke_agent_span(
135148
return span
136149

137150

138-
def update_invoke_agent_span(span: "sentry_sdk.tracing.Span", result: "Any") -> None:
151+
def update_invoke_agent_span(
152+
span: "Union[sentry_sdk.tracing.Span, StreamedSpan]",
153+
result: "Any",
154+
) -> None:
139155
"""Update and close the invoke agent span."""
140156
if not span or not result:
141157
return
@@ -154,7 +170,12 @@ def update_invoke_agent_span(span: "sentry_sdk.tracing.Span", result: "Any") ->
154170
try:
155171
response = result.response
156172
if hasattr(response, "model_name") and response.model_name:
157-
span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name)
173+
if isinstance(span, StreamedSpan):
174+
span.set_attribute(
175+
SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name
176+
)
177+
else:
178+
span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name)
158179
except Exception:
159180
# If response access fails, continue without setting model name
160181
pass

0 commit comments

Comments
 (0)