diff --git a/sentry_sdk/integrations/pydantic_ai/__init__.py b/sentry_sdk/integrations/pydantic_ai/__init__.py index 2988301791..2a19c1c5a4 100644 --- a/sentry_sdk/integrations/pydantic_ai/__init__.py +++ b/sentry_sdk/integrations/pydantic_ai/__init__.py @@ -19,10 +19,8 @@ from importlib.metadata import PackageNotFoundError, version from typing import TYPE_CHECKING -from .patches import ( - _patch_agent_run, -) from .spans.ai_client import ai_client_span, update_ai_client_span +from .spans.invoke_agent import invoke_agent_span, update_invoke_agent_span if TYPE_CHECKING: from typing import Any @@ -35,11 +33,13 @@ ToolCallPart, ToolDefinition, ) + from pydantic_ai.agent import AgentRunResult from pydantic_ai.capabilities import ( Hooks, RawToolArgs, ValidatedToolArgs, WrapModelRequestHandler, + WrapRunHandler, WrapToolExecuteHandler, ) from pydantic_ai.messages import ModelResponse @@ -59,7 +59,7 @@ async def on_model_request( ) -> "ModelResponse": with ai_client_span( messages=request_context.messages, - agent=None, + agent=ctx.agent, model=request_context.model, model_settings=request_context.model_settings, ) as span: @@ -118,6 +118,28 @@ async def sentry_wrap_tool_execute( _capture_exception(exc, handled=True) reraise(*exc_info) + @hooks.on.run + async def sentry_wrap_run( + ctx: "RunContext[Any]", + *, + handler: "WrapRunHandler", + ) -> "AgentRunResult[Any]": + with sentry_sdk.isolation_scope(), invoke_agent_span( + user_prompt=ctx.prompt, + agent=ctx.agent, + model=ctx.model, + model_settings=ctx.model_settings, + ) as span: + try: + result = await handler() + update_invoke_agent_span(span, result) + return result + except Exception as exc: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + _capture_exception(exc, handled=False) + reraise(*exc_info) + original_init = Agent.__init__ @functools.wraps(original_init) @@ -185,8 +207,6 @@ def setup_once() -> None: if PYDANTIC_AI_VERSION is None: return - _patch_agent_run() - try: from pydantic_ai.capabilities import Hooks except ImportError: diff --git a/sentry_sdk/integrations/pydantic_ai/patches/__init__.py b/sentry_sdk/integrations/pydantic_ai/patches/__init__.py deleted file mode 100644 index 0130dc1c4c..0000000000 --- a/sentry_sdk/integrations/pydantic_ai/patches/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .agent_run import _patch_agent_run # noqa: F401 diff --git a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py deleted file mode 100644 index 489f323679..0000000000 --- a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py +++ /dev/null @@ -1,181 +0,0 @@ -import sys -from functools import wraps -from typing import TYPE_CHECKING - -import sentry_sdk -from sentry_sdk.integrations import DidNotEnable -from sentry_sdk.utils import capture_internal_exceptions, reraise - -from ..spans import invoke_agent_span, update_invoke_agent_span -from ..utils import _capture_exception, pop_agent, push_agent - -try: - from pydantic_ai.agent import Agent -except ImportError: - raise DidNotEnable("pydantic-ai not installed") - -if TYPE_CHECKING: - from typing import Any, Callable, Optional - - -class _StreamingContextManagerWrapper: - """Wrapper for streaming methods that return async context managers.""" - - def __init__( - self, - agent: "Any", - original_ctx_manager: "Any", - user_prompt: "Any", - model: "Any", - model_settings: "Any", - is_streaming: bool = True, - ) -> None: - self.agent = agent - self.original_ctx_manager = original_ctx_manager - self.user_prompt = user_prompt - self.model = model - self.model_settings = model_settings - self.is_streaming = is_streaming - self._isolation_scope: "Any" = None - self._span: "Optional[sentry_sdk.traces.StreamedSpan]" = None - self._result: "Any" = None - - async def __aenter__(self) -> "Any": - # Set up isolation scope and invoke_agent span - self._isolation_scope = sentry_sdk.isolation_scope() - self._isolation_scope.__enter__() - - # Create invoke_agent span (will be closed in __aexit__) - self._span = invoke_agent_span( - self.user_prompt, - self.agent, - self.model, - self.model_settings, - self.is_streaming, - ) - - # 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) - - # Enter the original context manager - result = await self.original_ctx_manager.__aenter__() - self._result = result - return result - - async def __aexit__(self, exc_type: "Any", exc_val: "Any", exc_tb: "Any") -> None: - try: - # Exit the original context manager first - await self.original_ctx_manager.__aexit__(exc_type, exc_val, exc_tb) - - # Update span with result if successful - if exc_type is None and self._result and self._span is not None: - update_invoke_agent_span(self._span, self._result) - finally: - # Pop agent from contextvar stack - pop_agent() - - # Clean up invoke span - if self._span: - self._span.__exit__(exc_type, exc_val, exc_tb) - - # Clean up isolation scope - if self._isolation_scope: - self._isolation_scope.__exit__(exc_type, exc_val, exc_tb) - - -def _create_run_wrapper( - original_func: "Callable[..., Any]", is_streaming: bool = False -) -> "Callable[..., Any]": - """ - Wraps the Agent.run method to create an invoke_agent span. - - Args: - original_func: The original run method - is_streaming: Whether this is a streaming method (for future use) - """ - - @wraps(original_func) - async def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": - # Isolate each workflow so that when agents are run in asyncio tasks they - # don't touch each other's scopes - with sentry_sdk.isolation_scope(): - # Extract parameters for the span - user_prompt = kwargs.get("user_prompt") or (args[0] if args else None) - model = kwargs.get("model") - model_settings = kwargs.get("model_settings") - - # Create invoke_agent span - with invoke_agent_span( - user_prompt, self, model, model_settings, is_streaming - ) 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) - - try: - result = await original_func(self, *args, **kwargs) - - # Update span with result - update_invoke_agent_span(span, result) - - return result - except Exception as exc: - exc_info = sys.exc_info() - with capture_internal_exceptions(): - _capture_exception(exc) - reraise(*exc_info) - finally: - # Pop agent from contextvar stack - pop_agent() - - return wrapper - - -def _create_streaming_wrapper( - original_func: "Callable[..., Any]", -) -> "Callable[..., Any]": - """ - Wraps run_stream method that returns an async context manager. - """ - - @wraps(original_func) - def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": - # Extract parameters for the span - user_prompt = kwargs.get("user_prompt") or (args[0] if args else None) - model = kwargs.get("model") - model_settings = kwargs.get("model_settings") - - # Call original function to get the context manager - original_ctx_manager = original_func(self, *args, **kwargs) - - # Wrap it with our instrumentation - return _StreamingContextManagerWrapper( - agent=self, - original_ctx_manager=original_ctx_manager, - user_prompt=user_prompt, - model=model, - model_settings=model_settings, - is_streaming=True, - ) - - return wrapper - - -def _patch_agent_run() -> None: - """ - Patches the Agent run methods to create spans for agent execution. - - This patches both non-streaming (run, run_sync) and streaming - (run_stream, run_stream_events) methods. - """ - - # Store original methods - original_run = Agent.run - original_run_stream = Agent.run_stream - - # Wrap and apply patches for non-streaming methods - Agent.run = _create_run_wrapper(original_run, is_streaming=False) # type: ignore[method-assign] - - # Wrap and apply patches for streaming methods - Agent.run_stream = _create_streaming_wrapper(original_run_stream) # type: ignore[method-assign] diff --git a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py index 5576c7f5a3..9e7c5c31d2 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py @@ -18,8 +18,6 @@ _set_model_data, _should_send_inputs, _should_send_outputs, - get_current_agent, - get_is_streaming, ) from .utils import ( _serialize_binary_content_item, @@ -30,6 +28,7 @@ if TYPE_CHECKING: from typing import Any, Dict, List, Optional, Union + from pydantic_ai import Agent from pydantic_ai.messages import ModelMessage, ModelResponse, SystemPromptPart from pydantic_ai.models import Model from pydantic_ai.settings import ModelSettings @@ -271,7 +270,7 @@ def _set_output_data( def ai_client_span( messages: "list[ModelMessage]", - agent: "Any", + agent: "Optional[Agent[Any, Any]]", model: "Model", model_settings: "Optional[ModelSettings]", ) -> "StreamedSpan": @@ -283,12 +282,7 @@ def ai_client_span( model: Model object model_settings: Model settings """ - # Determine model name for span name - model_obj = model - if agent and hasattr(agent, "model"): - model_obj = agent.model - - model_name = _get_model_name(model_obj) or "unknown" + model_name = _get_model_name(model) or "unknown" span = sentry_sdk.traces.start_span( name=f"chat {model_name}", @@ -296,16 +290,12 @@ 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(), }, ) _set_agent_data(span, agent) - _set_model_data(span, model, model_settings) - - # Add available tools if agent is available - agent_obj = agent or get_current_agent() - _set_available_tools(span, agent_obj) + _set_model_data(span, agent, model, model_settings) + _set_available_tools(span, agent) # Set input messages (full conversation history) if messages: diff --git a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py index 8839bb8dc3..daef200cb9 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py @@ -22,7 +22,12 @@ ) if TYPE_CHECKING: - from typing import Any + from typing import Any, Optional, Sequence, Union + + from pydantic_ai import Agent, UserContent + from pydantic_ai.models import AbstractModel + from pydantic_ai.realtime.settings import RealtimeModelSettings + from pydantic_ai.settings import ModelSettings try: from pydantic_ai.messages import BinaryContent, ImageUrl @@ -32,11 +37,10 @@ def invoke_agent_span( - user_prompt: "Any", - agent: "Any", - model: "Any", - model_settings: "Any", - is_streaming: bool = False, + user_prompt: "Optional[Union[str, Sequence[UserContent]]]", + agent: "Optional[Agent]", + model: "AbstractModel", + model_settings: "Optional[Union[ModelSettings, RealtimeModelSettings]]", ) -> "StreamedSpan": """Create a span for invoking the agent.""" # Determine agent name for span @@ -54,7 +58,7 @@ def invoke_agent_span( ) _set_agent_data(span, agent) - _set_model_data(span, model, model_settings) + _set_model_data(span, agent, model, model_settings) _set_available_tools(span, agent) # Add user prompt and system prompts if available and prompts are enabled diff --git a/sentry_sdk/integrations/pydantic_ai/utils.py b/sentry_sdk/integrations/pydantic_ai/utils.py index 8628aa2e5a..29f7da3558 100644 --- a/sentry_sdk/integrations/pydantic_ai/utils.py +++ b/sentry_sdk/integrations/pydantic_ai/utils.py @@ -1,4 +1,3 @@ -from contextvars import ContextVar from typing import TYPE_CHECKING import sentry_sdk @@ -12,47 +11,12 @@ ) if TYPE_CHECKING: - from typing import Any, Optional + from typing import Any, Optional, Union - from pydantic_ai.models import Model - - -# Store the current agent context in a contextvar for re-entrant safety -# Using a list as a stack to support nested agent calls -_agent_context_stack: "ContextVar[list[dict[str, Any]]]" = ContextVar( - "pydantic_ai_agent_context_stack", default=[] -) - - -def push_agent(agent: "Any", is_streaming: bool = False) -> None: - """Push an agent context onto the stack along with its streaming flag.""" - stack = _agent_context_stack.get().copy() - stack.append({"agent": agent, "is_streaming": is_streaming}) - _agent_context_stack.set(stack) - - -def pop_agent() -> None: - """Pop an agent context from the stack.""" - stack = _agent_context_stack.get().copy() - if stack: - stack.pop() - _agent_context_stack.set(stack) - - -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 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 + from pydantic_ai import Agent + from pydantic_ai.models import AbstractModel, Model + from pydantic_ai.realtime.settings import RealtimeModelSettings + from pydantic_ai.settings import ModelSettings def _should_send_prompts_legacy() -> bool: @@ -92,24 +56,20 @@ def _should_send_outputs() -> bool: return _should_send_prompts_legacy() -def _set_agent_data(span: "StreamedSpan", agent: "Any") -> None: +def _set_agent_data(span: "StreamedSpan", agent: "Optional[Agent]") -> None: """Set agent-related data on a span. Args: span: The span to set data on - agent: Agent object (can be None, will try to get from contextvar if not provided) + agent: Agent object """ - # Extract agent name from agent object or contextvar - agent_obj = agent - if not agent_obj: - # Try to get from contextvar - agent_obj = get_current_agent() + if agent and hasattr(agent, "name") and agent.name: + span.set_attribute(SPANDATA.GEN_AI_AGENT_NAME, agent.name) - if agent_obj and hasattr(agent_obj, "name") and agent_obj.name: - span.set_attribute(SPANDATA.GEN_AI_AGENT_NAME, agent_obj.name) - -def _get_model_name(model_obj: "Model") -> "Optional[str]": +def _get_model_name( + model_obj: "Optional[Union[AbstractModel, Model, str]]", +) -> "Optional[str]": """Extract model name from a model object. Args: @@ -136,8 +96,9 @@ def _get_model_name(model_obj: "Model") -> "Optional[str]": def _set_model_data( span: "StreamedSpan", - model: "Any", - model_settings: "Any", + agent: "Optional[Agent]", + model: "Union[Model, AbstractModel]", + model_settings: "Optional[Union[ModelSettings, RealtimeModelSettings]]", ) -> None: """Set model-related data on a span. @@ -146,13 +107,10 @@ def _set_model_data( model: Model object (can be None, will try to get from agent if not provided) model_settings: Model settings (can be None, will try to get from agent if not provided) """ - # Try to get agent from contextvar if we need it - agent_obj = get_current_agent() - # Extract model information model_obj = model - if not model_obj and agent_obj and hasattr(agent_obj, "model"): - model_obj = agent_obj.model + if not model_obj and agent and hasattr(agent, "model"): + model_obj = agent.model if model_obj: # Set system from model @@ -166,8 +124,8 @@ def _set_model_data( # Extract model settings settings = model_settings - if not settings and agent_obj and hasattr(agent_obj, "model_settings"): - settings = agent_obj.model_settings + if not settings and agent and hasattr(agent, "model_settings"): + settings = agent.model_settings if settings: settings_map = { @@ -183,7 +141,7 @@ def _set_model_data( for setting_name, spandata_key in settings_map.items(): value = settings.get(setting_name) if value is not None: - span.set_attribute(spandata_key, value) + span.set_attribute(spandata_key, value) # type: ignore[arg-type] else: # Fallback for object-style settings for setting_name, spandata_key in settings_map.items(): @@ -193,7 +151,9 @@ def _set_model_data( span.set_attribute(spandata_key, value) -def _set_available_tools(span: "StreamedSpan", agent: "Any") -> None: +def _set_available_tools( + span: "StreamedSpan", agent: "Optional[Agent[Any, Any]]" +) -> None: """Set available tools data on a span from an agent's function toolset. Args: @@ -213,7 +173,7 @@ def _set_available_tools(span: "StreamedSpan", agent: "Any") -> None: # Get tools from the function toolset if hasattr(agent._function_toolset, "tools"): for tool_name, tool in agent._function_toolset.tools.items(): - tool_info = {"name": tool_name} + tool_info: "dict[str, Any]" = {"name": tool_name} # Add description from function_schema if available if hasattr(tool, "function_schema"): diff --git a/tests/integrations/pydantic_ai/test_pydantic_ai.py b/tests/integrations/pydantic_ai/test_pydantic_ai.py index 15280cf4d9..8255ab3386 100644 --- a/tests/integrations/pydantic_ai/test_pydantic_ai.py +++ b/tests/integrations/pydantic_ai/test_pydantic_ai.py @@ -121,7 +121,6 @@ async def test_agent_run_async( chat_span = chat_spans[0] assert "chat" in chat_span["name"] assert chat_span["attributes"]["gen_ai.operation.name"] == "chat" - assert chat_span["attributes"]["gen_ai.response.streaming"] is False assert json.loads(chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) == [ { "role": "user", @@ -211,9 +210,6 @@ def test_agent_run_sync( ] assert len(chat_spans) == 1 - # Verify streaming flag is False for sync - assert chat_spans[0]["attributes"]["gen_ai.response.streaming"] is False - def test_agent_run_sync_model_error( sentry_init, @@ -285,8 +281,6 @@ async def test_agent_run_stream( ] assert len(chat_spans) == 1 - # Verify streaming flag is True for streaming - assert chat_spans[0]["attributes"]["gen_ai.response.streaming"] is True assert json.loads( chat_spans[0]["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] ) == [ @@ -354,9 +348,6 @@ async def test_agent_run_stream_events( ] assert len(chat_spans) == 1 - # run_stream_events uses run() internally, so streaming should be False - assert chat_spans[0]["attributes"]["gen_ai.response.streaming"] is False - @pytest.mark.asyncio async def test_agent_with_tools( @@ -601,10 +592,6 @@ def multiply(a: int, b: int) -> int: sentry_sdk.flush() spans = [item.payload for item in items] - # Find span types - chat_spans = [ - s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" - ] tool_spans = [ s for s in spans @@ -614,9 +601,6 @@ def multiply(a: int, b: int) -> int: # Should have tool spans assert len(tool_spans) >= 1 - # Verify streaming flag is True - assert chat_spans[0]["attributes"]["gen_ai.response.streaming"] is True - # Check tool span tool_span = tool_spans[0] assert tool_span["attributes"]["gen_ai.tool.name"] == "multiply"