diff --git a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py index 864c83a506..d9fffadd9c 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py @@ -57,7 +57,7 @@ async def __aenter__(self) -> "Any": # Push agent to contextvar stack after span is successfully created and entered # This ensures proper pairing with pop_agent() in __aexit__ even if exceptions occur - push_agent(self.agent, self.is_streaming) + push_agent(self.agent) # Enter the original context manager result = await self.original_ctx_manager.__aenter__() @@ -120,7 +120,7 @@ async def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": ) as span: # Push agent to contextvar stack after span is successfully created and entered # This ensures proper pairing with pop_agent() in finally even if exceptions occur - push_agent(self, is_streaming) + push_agent(self) try: result = await original_func(self, *args, **kwargs) diff --git a/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py b/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py index 3a7d2fa635..d043acd03a 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py @@ -1,7 +1,9 @@ from contextlib import asynccontextmanager from functools import wraps +from sentry_sdk.consts import SPANDATA from sentry_sdk.integrations import DidNotEnable +from sentry_sdk.traces import StreamedSpan from ..spans import ( ai_client_span, @@ -59,9 +61,19 @@ def _patch_graph_nodes() -> None: @wraps(original_model_request_run) async def wrapped_model_request_run(self: "Any", ctx: "Any") -> "Any": + did_stream = getattr(self, "_did_stream", None) + cached_result = getattr(self, "_result", None) + if did_stream or cached_result is not None: + return await original_model_request_run(self, ctx) + messages, model, model_settings = _extract_span_data(self, ctx) with ai_client_span(messages, None, model, model_settings) as span: + if isinstance(span, StreamedSpan): + span.set_attribute(SPANDATA.GEN_AI_RESPONSE_STREAMING, False) + else: + span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, False) + result = await original_model_request_run(self, ctx) # Extract response from result if available @@ -85,10 +97,20 @@ def create_wrapped_stream( @asynccontextmanager @wraps(original_stream_method) async def wrapped_model_request_stream(self: "Any", ctx: "Any") -> "Any": + did_stream = getattr(self, "_did_stream", None) + if did_stream: + async with original_stream_method(self, ctx) as stream: + yield stream + messages, model, model_settings = _extract_span_data(self, ctx) # Create chat span for streaming request with ai_client_span(messages, None, model, model_settings) as span: + if isinstance(span, StreamedSpan): + span.set_attribute(SPANDATA.GEN_AI_RESPONSE_STREAMING, False) + else: + span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, False) + # Call the original stream method async with original_stream_method(self, ctx) as stream: yield stream diff --git a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py index c89b8609d8..429f7dae9d 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py @@ -23,7 +23,6 @@ _set_model_data, _should_send_prompts, get_current_agent, - get_is_streaming, ) from .utils import ( _serialize_binary_content_item, @@ -314,7 +313,6 @@ def ai_client_span( "sentry.op": OP.GEN_AI_CHAT, "sentry.origin": SPAN_ORIGIN, SPANDATA.GEN_AI_OPERATION_NAME: "chat", - SPANDATA.GEN_AI_RESPONSE_STREAMING: get_is_streaming(), }, ) else: @@ -325,8 +323,6 @@ def ai_client_span( ) span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat") - # Set streaming flag from contextvar - span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, get_is_streaming()) _set_agent_data(span, agent) _set_model_data(span, model, model_settings) diff --git a/sentry_sdk/integrations/pydantic_ai/utils.py b/sentry_sdk/integrations/pydantic_ai/utils.py index 340dcf8953..0da3a3133f 100644 --- a/sentry_sdk/integrations/pydantic_ai/utils.py +++ b/sentry_sdk/integrations/pydantic_ai/utils.py @@ -18,10 +18,10 @@ ) -def push_agent(agent: "Any", is_streaming: bool = False) -> None: - """Push an agent context onto the stack along with its streaming flag.""" +def push_agent(agent: "Any") -> None: + """Push an agent context onto the stack.""" stack = _agent_context_stack.get().copy() - stack.append({"agent": agent, "is_streaming": is_streaming}) + stack.append(agent) _agent_context_stack.set(stack) @@ -37,18 +37,10 @@ def get_current_agent() -> "Any": """Get the current agent from the contextvar stack.""" stack = _agent_context_stack.get() if stack: - return stack[-1]["agent"] + return stack[-1] return None -def get_is_streaming() -> bool: - """Get the streaming flag from the contextvar stack.""" - stack = _agent_context_stack.get() - if stack: - return stack[-1].get("is_streaming", False) - return False - - def _should_send_prompts() -> bool: """ Check if prompts should be sent to Sentry.