From cd97892fb6deea0748de8216ca57c786590334fa Mon Sep 17 00:00:00 2001 From: Erica Pisani Date: Wed, 12 Aug 2026 17:31:19 -0400 Subject: [PATCH] ref(pydantic-ai): Resolve version divergence in _compat and restructure modules Introduce _compat.py, which resolves the installed pydantic-ai version and every version-dependent decision once at import time: the hooks-vs-graph-nodes model backend, the ToolManager tool-call method name, and the message part classes (now resolved individually so one upstream rename degrades only the paths needing that class). The using_request_hooks class flag and the circular-import workarounds are gone; setup_once is three composed calls with deferred imports, and importing the package loads nothing but __init__. Restructure to the target layout: _spans.py absorbs utils.py and the spans/ package, _wrap_agent.py and _wrap_model.py (both model backends behind one install_model_backend()) and _wrap_tools.py (the two duplicated tool wrappers unified into one) replace patches/. Review-driven fixes folded in: unknown ToolManager method names and a missing _agent_graph module now degrade gracefully instead of crashing sentry_sdk.init(), agent_run_scope removes exactly its own run from the stack so non-LIFO streaming exits cannot corrupt it, the streaming wrapper propagates exception suppression from the wrapped context manager, the after_model_request span close is guarded, and chat spans extract model info once instead of twice. --- .../integrations/pydantic_ai/__init__.py | 170 +------ .../integrations/pydantic_ai/_compat.py | 70 +++ .../integrations/pydantic_ai/_extract.py | 32 +- .../integrations/pydantic_ai/_run_context.py | 22 +- sentry_sdk/integrations/pydantic_ai/_spans.py | 446 ++++++++++++++++++ .../{patches/agent_run.py => _wrap_agent.py} | 58 ++- .../integrations/pydantic_ai/_wrap_model.py | 220 +++++++++ .../integrations/pydantic_ai/_wrap_tools.py | 90 ++++ .../pydantic_ai/patches/__init__.py | 3 - .../pydantic_ai/patches/graph_nodes.py | 100 ---- .../integrations/pydantic_ai/patches/tools.py | 173 ------- .../pydantic_ai/spans/__init__.py | 3 - .../pydantic_ai/spans/ai_client.py | 172 ------- .../pydantic_ai/spans/execute_tool.py | 84 ---- .../pydantic_ai/spans/invoke_agent.py | 110 ----- .../integrations/pydantic_ai/spans/utils.py | 35 -- sentry_sdk/integrations/pydantic_ai/utils.py | 109 ----- .../pydantic_ai/test_pydantic_ai.py | 60 +-- 18 files changed, 928 insertions(+), 1029 deletions(-) create mode 100644 sentry_sdk/integrations/pydantic_ai/_compat.py create mode 100644 sentry_sdk/integrations/pydantic_ai/_spans.py rename sentry_sdk/integrations/pydantic_ai/{patches/agent_run.py => _wrap_agent.py} (81%) create mode 100644 sentry_sdk/integrations/pydantic_ai/_wrap_model.py create mode 100644 sentry_sdk/integrations/pydantic_ai/_wrap_tools.py delete mode 100644 sentry_sdk/integrations/pydantic_ai/patches/__init__.py delete mode 100644 sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py delete mode 100644 sentry_sdk/integrations/pydantic_ai/patches/tools.py delete mode 100644 sentry_sdk/integrations/pydantic_ai/spans/__init__.py delete mode 100644 sentry_sdk/integrations/pydantic_ai/spans/ai_client.py delete mode 100644 sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py delete mode 100644 sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py delete mode 100644 sentry_sdk/integrations/pydantic_ai/spans/utils.py delete mode 100644 sentry_sdk/integrations/pydantic_ai/utils.py diff --git a/sentry_sdk/integrations/pydantic_ai/__init__.py b/sentry_sdk/integrations/pydantic_ai/__init__.py index 58a2a7bfbc..307ae5e3ac 100644 --- a/sentry_sdk/integrations/pydantic_ai/__init__.py +++ b/sentry_sdk/integrations/pydantic_ai/__init__.py @@ -1,129 +1,11 @@ -import functools - from sentry_sdk.integrations import DidNotEnable, Integration -from sentry_sdk.utils import capture_internal_exceptions, parse_version try: import pydantic_ai # noqa: F401 - from pydantic_ai import Agent except ImportError: raise DidNotEnable("pydantic-ai not installed") -from importlib.metadata import PackageNotFoundError, version -from typing import TYPE_CHECKING - -from .patches import ( - _patch_agent_run, - _patch_graph_nodes, - _patch_tool_execution, -) -from .spans.ai_client import ai_client_span, update_ai_client_span - -if TYPE_CHECKING: - from typing import Any - - from pydantic_ai import ModelRequestContext, RunContext - from pydantic_ai.capabilities import Hooks - from pydantic_ai.messages import ModelResponse - - -def register_hooks(hooks: "Hooks") -> None: - """ - Creates hooks for chat model calls and register the hooks by adding the hooks to the `capabilities` argument passed to `Agent.__init__()`. - - The chat span opened in on_request is stored in the run's `RunContext.metadata` - dict, which pydantic-ai shares by reference between the hooks of one run. This - keeps span pairing correct per run (even for overlapping runs in one task) and - covers every entry point that fires request hooks (including `Agent.iter()`, - which the Agent.run/run_stream wrappers never see). It requires seeding a - metadata dict in `patched_init` below when the user did not provide one. - """ - - @hooks.on.before_model_request - async def on_request( - ctx: "RunContext[None]", request_context: "ModelRequestContext" - ) -> "ModelRequestContext": - run_context_metadata = ctx.metadata - if not isinstance(run_context_metadata, dict): - return request_context - - span = None - with capture_internal_exceptions(): - span = ai_client_span( - messages=request_context.messages, - agent=None, - model=request_context.model, - model_settings=request_context.model_settings, - ) - - if span is None: - return request_context - - run_context_metadata["_sentry_span"] = span - span.__enter__() - - return request_context - - @hooks.on.after_model_request - async def on_response( - ctx: "RunContext[None]", - *, - request_context: "ModelRequestContext", - response: "ModelResponse", - ) -> "ModelResponse": - run_context_metadata = ctx.metadata - if not isinstance(run_context_metadata, dict): - return response - - span = run_context_metadata.pop("_sentry_span", None) - if span is None: - return response - - with capture_internal_exceptions(): - update_ai_client_span(span, response) - span.__exit__(None, None, None) - - return response - - @hooks.on.model_request_error - async def on_error( - ctx: "RunContext[None]", - *, - request_context: "ModelRequestContext", - error: "Exception", - ) -> "ModelResponse": - run_context_metadata = ctx.metadata - - if not isinstance(run_context_metadata, dict): - raise error - - span = run_context_metadata.pop("_sentry_span", None) - if span is None: - raise error - - with capture_internal_exceptions(): - span.__exit__(type(error), error, error.__traceback__) - - raise error - - original_init = Agent.__init__ - - @functools.wraps(original_init) - def patched_init(self: "Agent[Any, Any]", *args: "Any", **kwargs: "Any") -> None: - caps = list(kwargs.get("capabilities") or []) - caps.append(hooks) - kwargs["capabilities"] = caps - - metadata = kwargs.get("metadata") - if metadata is None: - kwargs["metadata"] = {} # Used as shared reference between hooks - - return original_init(self, *args, **kwargs) - - Agent.__init__ = patched_init # type: ignore[method-assign] - - class PydanticAIIntegration(Integration): """ Typical interaction with the library: @@ -131,18 +13,18 @@ class PydanticAIIntegration(Integration): 2. The user calls `Agent.run()` or `Agent.run_stream()` to start an agent run. The latter can be used to incrementally receive progress. 3. In a loop, the agent repeatedly calls the model, maintaining a conversation history that includes previous messages and tool results, which is passed to each call. - Internally, Pydantic AI maintains an execution graph in which ModelRequestNode are responsible for model calls, including retries. - Hooks using the decorators provided by `pydantic_ai.capabilities` create and manage spans for model calls when these hooks are available (newer library versions); - older versions are instrumented by patching the graph nodes directly (see patches/graph_nodes.py). - - The wrappers around `Agent.run()` and `Agent.run_stream()` track each in-flight run on a contextvar stack (see _run_context.py); the tool patches and span - helpers read the current agent from there. The request hooks pair each chat span with its model request through the run's `RunContext.metadata` dict - (see register_hooks), which stays correct per run and also covers entry points the wrappers don't instrument, such as `Agent.iter()`. + How the integration is put together: + - _compat.py resolves the installed pydantic-ai version and every version-dependent decision, once, at import time. + - _extract.py is the only module that reads pydantic-ai object internals; it returns plain data structures. + - _spans.py creates spans and writes extracted data onto them. + - _run_context.py tracks each in-flight run on a contextvar stack; the tool wrapper and span helpers read the current agent from there. + - _wrap_agent.py instruments Agent.run / Agent.run_stream (invoke_agent spans, isolation scopes, run tracking). + - _wrap_model.py emits chat spans for model requests via one of two backends chosen in _compat: request hooks (>= 1.73), paired per run through RunContext.metadata, or graph-node patching (older versions). + - _wrap_tools.py instruments the single ToolManager method all tool calls flow through (execute_tool spans). """ identifier = "pydantic_ai" origin = f"auto.ai.{identifier}" - using_request_hooks = False def __init__( self, include_prompts: bool = True, handled_tool_call_exceptions: bool = True @@ -169,32 +51,14 @@ def setup_once() -> None: - Model requests (AI client calls) - Tool executions """ + # Deferred imports keep `import sentry_sdk.integrations.pydantic_ai` + # cheap when the integration is never enabled; they are the only + # intra-package imports in this module, keeping the import graph + # acyclic. + from ._wrap_agent import _patch_agent_run + from ._wrap_model import install_model_backend + from ._wrap_tools import _patch_tool_execution + _patch_agent_run() _patch_tool_execution() - - PydanticAIIntegration.using_request_hooks = False - try: - PYDANTIC_AI_VERSION = version("pydantic-ai-slim") - except PackageNotFoundError: - return - - PYDANTIC_AI_VERSION = parse_version(PYDANTIC_AI_VERSION) - if PYDANTIC_AI_VERSION is None: - return - - # ModelRequestContext.model added in https://github.com/pydantic/pydantic-ai/commit/f1260dfe09907f17688eee1646daf898fc428d4c - if PYDANTIC_AI_VERSION < ( - 1, - 73, - ): - _patch_graph_nodes() - return - - try: - from pydantic_ai.capabilities import Hooks - except ImportError: - return - - PydanticAIIntegration.using_request_hooks = True - hooks = Hooks() - register_hooks(hooks) + install_model_backend() diff --git a/sentry_sdk/integrations/pydantic_ai/_compat.py b/sentry_sdk/integrations/pydantic_ai/_compat.py new file mode 100644 index 0000000000..003458fb8a --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/_compat.py @@ -0,0 +1,70 @@ +"""Version detection and version-dependent imports for pydantic-ai. + +Everything here is resolved once at import time. The rest of the integration +consumes the resulting constants instead of probing versions or attributes at +call time, so "what runs on version X" is answered entirely by this module. +""" + +from typing import TYPE_CHECKING + +from sentry_sdk.integrations import DidNotEnable +from sentry_sdk.utils import package_version + +try: + from pydantic_ai import messages as _messages + from pydantic_ai.agent import Agent # noqa: F401 + from pydantic_ai.exceptions import ToolRetryError # noqa: F401 + + try: + from pydantic_ai.tool_manager import ToolManager + except ImportError: + # older versions + from pydantic_ai._tool_manager import ToolManager # type: ignore +except ImportError: + raise DidNotEnable("pydantic-ai not installed") + +if TYPE_CHECKING: + from typing import Optional + +# Message part classes are resolved individually so that a single upstream +# rename degrades only the extraction paths that need that class, instead of +# silently disabling all of them at once. +BaseToolCallPart = getattr(_messages, "BaseToolCallPart", None) +BaseToolReturnPart = getattr(_messages, "BaseToolReturnPart", None) +BinaryContent = getattr(_messages, "BinaryContent", None) +ImageUrl = getattr(_messages, "ImageUrl", None) +SystemPromptPart = getattr(_messages, "SystemPromptPart", None) +TextPart = getattr(_messages, "TextPart", None) +ThinkingPart = getattr(_messages, "ThinkingPart", None) + +PYDANTIC_AI_VERSION = package_version("pydantic-ai-slim") + +# The ToolManager method through which all tool calls flow; renamed from +# _call_tool to execute_tool_call in newer versions. None means the method +# could not be found and tool instrumentation is skipped. +TOOL_CALL_METHOD: "Optional[str]" = None +if hasattr(ToolManager, "execute_tool_call"): + TOOL_CALL_METHOD = "execute_tool_call" +elif hasattr(ToolManager, "_call_tool"): + TOOL_CALL_METHOD = "_call_tool" + +# Request hooks (pydantic_ai.capabilities) are usable from 1.73 on, when +# ModelRequestContext.model was added: +# https://github.com/pydantic/pydantic-ai/commit/f1260dfe09907f17688eee1646daf898fc428d4c +USES_REQUEST_HOOKS = False +if PYDANTIC_AI_VERSION is not None and PYDANTIC_AI_VERSION >= (1, 73): + try: + from pydantic_ai.capabilities import Hooks # noqa: F401 + + USES_REQUEST_HOOKS = True + except ImportError: + USES_REQUEST_HOOKS = False + +# Which mechanism emits chat spans for model requests: request hooks on new +# versions, graph-node patching on old ones. None (unknown version, or hooks +# unavailable on a new version) means only agent and tool spans are emitted. +MODEL_BACKEND: "Optional[str]" = None +if USES_REQUEST_HOOKS: + MODEL_BACKEND = "hooks" +elif PYDANTIC_AI_VERSION is not None and PYDANTIC_AI_VERSION < (1, 73): + MODEL_BACKEND = "graph_nodes" diff --git a/sentry_sdk/integrations/pydantic_ai/_extract.py b/sentry_sdk/integrations/pydantic_ai/_extract.py index ff0a85ae65..ff362b5bcb 100644 --- a/sentry_sdk/integrations/pydantic_ai/_extract.py +++ b/sentry_sdk/integrations/pydantic_ai/_extract.py @@ -4,8 +4,8 @@ private attributes and version-dependent shapes) so that upstream library changes are absorbed here rather than throughout the integration. The one exception is control-flow state read at the patch points themselves (e.g. -ModelRequestNode._did_stream in patches/graph_nodes.py and Tool.tool_def in -patches/tools.py); everything else consumes the plain data structures +ModelRequestNode._did_stream in _wrap_model.py and Tool.tool_def in +_wrap_tools.py); everything else consumes the plain data structures returned here. """ @@ -18,25 +18,15 @@ from sentry_sdk.consts import SPANDATA from sentry_sdk.utils import safe_serialize -try: - from pydantic_ai.messages import ( - BaseToolCallPart, - BaseToolReturnPart, - BinaryContent, - ImageUrl, - SystemPromptPart, - TextPart, - ThinkingPart, - ) -except ImportError: - # Fallback if these classes are not available - BaseToolCallPart = None # type: ignore[misc,assignment] - BaseToolReturnPart = None # type: ignore[misc,assignment] - BinaryContent = None # type: ignore[misc,assignment] - ImageUrl = None # type: ignore[misc,assignment] - SystemPromptPart = None # type: ignore[misc,assignment] - TextPart = None # type: ignore[misc,assignment] - ThinkingPart = None # type: ignore[misc,assignment] +from ._compat import ( + BaseToolCallPart, + BaseToolReturnPart, + BinaryContent, + ImageUrl, + SystemPromptPart, + TextPart, + ThinkingPart, +) if TYPE_CHECKING: from typing import Any, Dict, List, Optional diff --git a/sentry_sdk/integrations/pydantic_ai/_run_context.py b/sentry_sdk/integrations/pydantic_ai/_run_context.py index 4253ac50da..4c8c30850f 100644 --- a/sentry_sdk/integrations/pydantic_ai/_run_context.py +++ b/sentry_sdk/integrations/pydantic_ai/_run_context.py @@ -46,18 +46,18 @@ def get_is_streaming() -> bool: @contextmanager def agent_run_scope(agent: "Any", is_streaming: bool = False) -> "Iterator[AgentRun]": """Track an agent run on the contextvar stack for the duration of the - with block.""" + with block. + + On exit, exactly this run is removed from the stack (by identity, not a + token reset), so streaming runs that exit out of LIFO order or in a + different asyncio task never erase other still-active runs. + """ run = AgentRun(agent=agent, is_streaming=is_streaming) - token = _agent_run_stack.set(_agent_run_stack.get() + (run,)) + _agent_run_stack.set(_agent_run_stack.get() + (run,)) try: yield run finally: - try: - _agent_run_stack.reset(token) - except (LookupError, ValueError): - # A streaming run's context manager can be exited in a different - # asyncio task (and therefore a different Context) than it was - # entered in, in which case the token cannot be reset. The stack - # entry only lives in the entering task's context copy, so there - # is nothing to clean up. - pass + stack = _agent_run_stack.get() + new_stack = tuple(r for r in stack if r is not run) + if len(new_stack) != len(stack): + _agent_run_stack.set(new_stack) diff --git a/sentry_sdk/integrations/pydantic_ai/_spans.py b/sentry_sdk/integrations/pydantic_ai/_spans.py new file mode 100644 index 0000000000..fc72c0899d --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/_spans.py @@ -0,0 +1,446 @@ +"""Span creation and data population for the pydantic-ai integration. + +Functions here consume the plain data structures returned by _extract and +write them onto Sentry spans; they contain no pydantic-ai object probing of +their own. +""" + +import json +from typing import TYPE_CHECKING + +import sentry_sdk +from sentry_sdk.ai.monitoring import record_token_usage +from sentry_sdk.ai.utils import ( + _set_span_data_attribute, + get_start_span_function, + normalize_message_roles, + set_data_normalized, + truncate_and_annotate_messages, +) +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing_utils import ( + has_span_streaming_enabled, + should_truncate_gen_ai_input, +) +from sentry_sdk.utils import event_from_exception, safe_serialize + +from ._extract import ( + MODEL_SETTINGS_TO_SPANDATA, + ModelInfo, + extract_agent_name, + extract_agent_prompt_messages, + extract_available_tools, + extract_model_info, + extract_request_messages, + extract_response_model_name, + extract_response_parts, + extract_system_instructions, + extract_usage_kwargs, +) +from ._run_context import get_current_agent, get_is_streaming +from .consts import SPAN_ORIGIN + +if TYPE_CHECKING: + from typing import Any, Optional, Union + + from pydantic_ai._tool_manager import ToolDefinition # type: ignore + from pydantic_ai.messages import ModelResponse + from pydantic_ai.usage import RequestUsage, RunUsage + + from sentry_sdk.traces import StreamedSpan + + +def _should_send_prompts() -> bool: + """ + Check if prompts should be sent to Sentry. + + This checks both send_default_pii and the include_prompts integration setting. + """ + if not should_send_default_pii(): + return False + + from . import PydanticAIIntegration + + # Get the integration instance from the client + integration = sentry_sdk.get_client().get_integration(PydanticAIIntegration) + + if integration is None: + return False + + return getattr(integration, "include_prompts", False) + + +def _capture_exception(exc: "Any", handled: bool = False) -> None: + event, hint = event_from_exception( + exc, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "pydantic_ai", "handled": handled}, + ) + sentry_sdk.capture_event(event, hint=hint) + + +def _set_agent_data( + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", agent: "Any" +) -> 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) + """ + # Extract agent name from agent object or contextvar + agent_name = extract_agent_name(agent or get_current_agent()) + if agent_name: + _set_span_data_attribute(span, SPANDATA.GEN_AI_AGENT_NAME, agent_name) + + +def _set_model_data( + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", + model: "Any", + model_settings: "Any", + agent: "Any" = None, + model_info: "Optional[ModelInfo]" = None, +) -> None: + """Set model-related data on a span. + + Args: + span: The span to set data on + 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) + agent: Agent to fall back to for model and settings (defaults to the + agent of the current run) + model_info: Already-extracted model info; passing it avoids a second + extraction when the caller needed it anyway + """ + if model_info is None: + model_info = extract_model_info( + model, model_settings, agent or get_current_agent() + ) + + if model_info.system is not None: + _set_span_data_attribute(span, SPANDATA.GEN_AI_SYSTEM, model_info.system) + + if model_info.name: + _set_span_data_attribute(span, SPANDATA.GEN_AI_REQUEST_MODEL, model_info.name) + + for setting_name, value in model_info.settings.items(): + spandata_key = MODEL_SETTINGS_TO_SPANDATA.get(setting_name) + if spandata_key is not None: + _set_span_data_attribute(span, spandata_key, value) + + +def _set_available_tools( + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", agent: "Any" +) -> None: + """Set available tools data on a span from an agent's function toolset. + + Args: + span: The span to set data on + agent: Agent object with _function_toolset attribute + """ + tools = extract_available_tools(agent) + if tools: + _set_span_data_attribute( + span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools) + ) + + +def _set_request_messages_data( + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", + messages: "list[dict[str, Any]]", +) -> None: + """Normalize, truncate if configured, and set gen_ai.request.messages.""" + normalized_messages = normalize_message_roles(messages) + client = sentry_sdk.get_client() + scope = sentry_sdk.get_current_scope() + messages_data = ( + truncate_and_annotate_messages(normalized_messages, span, scope) + if should_truncate_gen_ai_input(client.options) + else normalized_messages + ) + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False + ) + + +def _set_usage_data( + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", + usage: "Union[RequestUsage, RunUsage]", +) -> None: + """Set token usage data on a span. + + This function works with both RequestUsage (single request) and + RunUsage (agent run) objects from pydantic_ai. + + Args: + span: The Sentry span to set data on. + usage: RequestUsage or RunUsage object containing token usage information. + """ + usage_kwargs = extract_usage_kwargs(usage) + if usage_kwargs is None: + return + + record_token_usage(span, **usage_kwargs) + + +def _set_input_messages( + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", messages: "Any" +) -> None: + """Set input messages data on a span.""" + if not _should_send_prompts(): + return + + if not messages: + return + + try: + system_instructions = extract_system_instructions(messages) + if system_instructions: + _set_span_data_attribute( + span, + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + json.dumps(system_instructions), + ) + + formatted_messages = extract_request_messages(messages) + + if formatted_messages: + _set_request_messages_data(span, formatted_messages) + except Exception: + # If we fail to format messages, just skip it + pass + + +def _set_output_data( + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", + response: "Optional[ModelResponse]", +) -> None: + """Set output data on a span.""" + if not _should_send_prompts(): + return + + if not response: + return + + if response.model_name: + _set_span_data_attribute( + span, SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name + ) + + try: + parts = extract_response_parts(response) + if parts: + _set_span_data_attribute( + span, + SPANDATA.GEN_AI_OUTPUT_MESSAGES, + json.dumps([{"role": "assistant", "parts": parts}]), + ) + except Exception: + # If we fail to format output, just skip it + pass + + +def ai_client_span( + messages: "Any", agent: "Any", model: "Any", model_settings: "Any" +) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]": + """Create a span for an AI client call (model request). + + Args: + messages: Full conversation history (list of messages) + agent: Agent object + model: Model object + model_settings: Model settings + """ + # Resolve the agent and model info once so the span name and every + # attribute derived below (gen_ai.request.model, agent data, available + # tools) agree + agent_obj = agent or get_current_agent() + model_info = extract_model_info(model, model_settings, agent_obj) + model_name = model_info.name or "unknown" + + span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) + if span_streaming: + span = sentry_sdk.traces.start_span( + name=f"chat {model_name}", + attributes={ + "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: + span = sentry_sdk.start_span( + op=OP.GEN_AI_CHAT, + name=f"chat {model_name}", + origin=SPAN_ORIGIN, + ) + + 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_obj) + _set_model_data(span, model, model_settings, model_info=model_info) + + # Add available tools if agent is available + _set_available_tools(span, agent_obj) + + # Set input messages (full conversation history) + if messages: + _set_input_messages(span, messages) + + return span + + +def update_ai_client_span( + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", + model_response: "Optional[ModelResponse]", +) -> None: + """Update the AI client span with response data.""" + if not span: + return + + # Set usage data if available + if model_response and hasattr(model_response, "usage"): + _set_usage_data(span, model_response.usage) + + # Set output data + _set_output_data(span, model_response) + + +def invoke_agent_span( + user_prompt: "Any", + agent: "Any", + model: "Any", + model_settings: "Any", + is_streaming: bool = False, +) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]": + """Create a span for invoking the agent.""" + # Determine agent name for span + name = extract_agent_name(agent) or "agent" + + span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) + if span_streaming: + span = sentry_sdk.traces.start_span( + name=f"invoke_agent {name}", + attributes={ + "sentry.op": OP.GEN_AI_INVOKE_AGENT, + "sentry.origin": SPAN_ORIGIN, + SPANDATA.GEN_AI_OPERATION_NAME: "invoke_agent", + }, + ) + else: + span = get_start_span_function()( + op=OP.GEN_AI_INVOKE_AGENT, + name=f"invoke_agent {name}", + origin=SPAN_ORIGIN, + ) + + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") + + _set_agent_data(span, agent) + _set_model_data(span, model, model_settings, agent=agent) + _set_available_tools(span, agent) + + # Add user prompt and system prompts if available and prompts are enabled + if _should_send_prompts(): + messages = extract_agent_prompt_messages(agent, user_prompt) + + if messages: + _set_request_messages_data(span, messages) + + return span + + +def update_invoke_agent_span( + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", + result: "Any", +) -> None: + """Update and close the invoke agent span.""" + if not span or not result: + return + + # Extract output from result + output = getattr(result, "output", None) + + # Set response text if prompts are enabled + if _should_send_prompts() and output: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TEXT, str(output), unpack=False + ) + + # Set model name from response if available + response_model_name = extract_response_model_name(result) + if response_model_name: + _set_span_data_attribute( + span, SPANDATA.GEN_AI_RESPONSE_MODEL, response_model_name + ) + + +def execute_tool_span( + tool_name: str, + tool_args: "Any", + agent: "Any", + tool_definition: "Optional[ToolDefinition]" = None, +) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]": + """Create a span for tool execution. + + Args: + tool_name: The name of the tool being executed + tool_args: The arguments passed to the tool + agent: The agent executing the tool + tool_definition: The definition of the tool, if available + """ + span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) + if span_streaming: + # Both keys must be present at span start so that attribute-based + # ignore_spans / traces_sampler rules can match this span. + span = sentry_sdk.traces.start_span( + name=f"execute_tool {tool_name}", + attributes={ + "sentry.op": OP.GEN_AI_EXECUTE_TOOL, + "sentry.origin": SPAN_ORIGIN, + SPANDATA.GEN_AI_OPERATION_NAME: "execute_tool", + SPANDATA.GEN_AI_TOOL_NAME: tool_name, + }, + ) + else: + span = sentry_sdk.start_span( + op=OP.GEN_AI_EXECUTE_TOOL, + name=f"execute_tool {tool_name}", + origin=SPAN_ORIGIN, + ) + + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "execute_tool") + span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool_name) + + if tool_definition is not None and hasattr(tool_definition, "description"): + _set_span_data_attribute( + span, + SPANDATA.GEN_AI_TOOL_DESCRIPTION, + tool_definition.description, + ) + + _set_agent_data(span, agent) + + if _should_send_prompts() and tool_args is not None: + _set_span_data_attribute( + span, SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_args) + ) + + return span + + +def update_execute_tool_span( + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", result: "Any" +) -> None: + """Update the execute tool span with the result.""" + if not span: + return + + if not _should_send_prompts() or result is None: + return + + _set_span_data_attribute(span, SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result)) diff --git a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py b/sentry_sdk/integrations/pydantic_ai/_wrap_agent.py similarity index 81% rename from sentry_sdk/integrations/pydantic_ai/patches/agent_run.py rename to sentry_sdk/integrations/pydantic_ai/_wrap_agent.py index cde48326a0..5c80a244f2 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py +++ b/sentry_sdk/integrations/pydantic_ai/_wrap_agent.py @@ -1,20 +1,20 @@ +"""Instrumentation of the Agent.run / Agent.run_stream entry points.""" + import sys from contextlib import ExitStack 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 .._run_context import agent_run_scope -from ..spans import invoke_agent_span, update_invoke_agent_span -from ..utils import _capture_exception - -try: - from pydantic_ai.agent import Agent -except ImportError: - raise DidNotEnable("pydantic-ai not installed") +from ._compat import USES_REQUEST_HOOKS, Agent +from ._run_context import agent_run_scope +from ._spans import ( + _capture_exception, + invoke_agent_span, + update_invoke_agent_span, +) if TYPE_CHECKING: from typing import Any, Callable, Optional, Union @@ -28,6 +28,17 @@ def _extract_run_params( return user_prompt, kwargs.get("model"), kwargs.get("model_settings") +def _seed_run_metadata(kwargs: "dict[str, Any]") -> None: + """Seed the run's metadata dict when the request hooks are in use. + + The hooks pair each chat span with its model request through the run's + RunContext.metadata dict (see _wrap_model.py), which requires the + metadata object to be a dict shared by reference between hooks. + """ + if USES_REQUEST_HOOKS and kwargs.get("metadata") is None: + kwargs["metadata"] = {"_sentry_span": None} + + class _StreamingContextManagerWrapper: """Wrapper for streaming methods that return async context managers.""" @@ -72,14 +83,21 @@ async def __aenter__(self) -> "Any": self._result = result return result - async def __aexit__(self, exc_type: "Any", exc_val: "Any", exc_tb: "Any") -> None: + async def __aexit__(self, exc_type: "Any", exc_val: "Any", exc_tb: "Any") -> "Any": try: - # Exit the original context manager first - await self.original_ctx_manager.__aexit__(exc_type, exc_val, exc_tb) + # Exit the original context manager first; propagate its exception + # suppression so the integration never changes control flow. + suppressed = await self.original_ctx_manager.__aexit__( + exc_type, exc_val, exc_tb + ) + if suppressed: + exc_type = exc_val = exc_tb = None # 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) + + return suppressed finally: if self._contexts is not None: self._contexts.__exit__(exc_type, exc_val, exc_tb) @@ -89,17 +107,11 @@ def _create_run_wrapper(original_func: "Callable[..., Any]") -> "Callable[..., A """ Wraps the Agent.run method to create an invoke_agent span. """ - from sentry_sdk.integrations.pydantic_ai import ( - PydanticAIIntegration, - ) # Required to avoid circular import @wraps(original_func) async def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": user_prompt, model, model_settings = _extract_run_params(args, kwargs) - - if PydanticAIIntegration.using_request_hooks: - if kwargs.get("metadata") is None: - kwargs["metadata"] = {"_sentry_span": None} + _seed_run_metadata(kwargs) # Isolate each workflow so that when agents are run in asyncio tasks # they don't touch each other's scopes @@ -129,17 +141,11 @@ def _create_streaming_wrapper( """ Wraps run_stream method that returns an async context manager. """ - from sentry_sdk.integrations.pydantic_ai import ( - PydanticAIIntegration, - ) # Required to avoid circular import @wraps(original_func) def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": user_prompt, model, model_settings = _extract_run_params(args, kwargs) - - if PydanticAIIntegration.using_request_hooks: - if kwargs.get("metadata") is None: - kwargs["metadata"] = {"_sentry_span": None} + _seed_run_metadata(kwargs) # Call original function to get the context manager original_ctx_manager = original_func(self, *args, **kwargs) diff --git a/sentry_sdk/integrations/pydantic_ai/_wrap_model.py b/sentry_sdk/integrations/pydantic_ai/_wrap_model.py new file mode 100644 index 0000000000..9550942fb1 --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/_wrap_model.py @@ -0,0 +1,220 @@ +"""Chat-span emission for model requests. + +Two backends, selected once in _compat, implement the same duty — emit a +gen_ai.chat span around each model request: + +- "hooks" (pydantic-ai >= 1.73): request hooks registered through + pydantic_ai.capabilities, paired per run through RunContext.metadata. +- "graph_nodes" (older versions): direct patching of ModelRequestNode. +""" + +import functools +from contextlib import asynccontextmanager +from functools import wraps +from typing import TYPE_CHECKING + +from sentry_sdk.utils import capture_internal_exceptions + +from ._compat import MODEL_BACKEND, Agent +from ._extract import extract_graph_request_data +from ._spans import ai_client_span, update_ai_client_span + +if TYPE_CHECKING: + from typing import Any, Callable, Optional + + from pydantic_ai import ModelRequestContext, RunContext + from pydantic_ai.capabilities import Hooks as HooksType + from pydantic_ai.messages import ModelResponse + + +def install_model_backend() -> None: + """Install the model-request instrumentation for the installed version.""" + if MODEL_BACKEND == "hooks": + _install_request_hooks() + elif MODEL_BACKEND == "graph_nodes": + _patch_graph_nodes() + + +def _install_request_hooks() -> None: + """ + Creates hooks for chat model calls and registers them by adding them to the + `capabilities` argument passed to `Agent.__init__()`. + + The chat span opened in on_request is stored in the run's + `RunContext.metadata` dict, which pydantic-ai shares by reference between + the hooks of one run. This keeps span pairing correct per run (even for + overlapping runs in one task). It requires seeding a metadata dict in + `patched_init` below (and in the run wrappers, see _wrap_agent.py) when + the user did not provide one. + """ + from pydantic_ai.capabilities import Hooks + + hooks: "HooksType" = Hooks() + + @hooks.on.before_model_request + async def on_request( + ctx: "RunContext[None]", request_context: "ModelRequestContext" + ) -> "ModelRequestContext": + run_context_metadata = ctx.metadata + if not isinstance(run_context_metadata, dict): + return request_context + + span = None + with capture_internal_exceptions(): + span = ai_client_span( + messages=request_context.messages, + agent=None, + model=request_context.model, + model_settings=request_context.model_settings, + ) + + if span is None: + return request_context + + run_context_metadata["_sentry_span"] = span + span.__enter__() + + return request_context + + @hooks.on.after_model_request + async def on_response( + ctx: "RunContext[None]", + *, + request_context: "ModelRequestContext", + response: "ModelResponse", + ) -> "ModelResponse": + run_context_metadata = ctx.metadata + if not isinstance(run_context_metadata, dict): + return response + + span = run_context_metadata.pop("_sentry_span", None) + if span is None: + return response + + with capture_internal_exceptions(): + update_ai_client_span(span, response) + with capture_internal_exceptions(): + span.__exit__(None, None, None) + + return response + + @hooks.on.model_request_error + async def on_error( + ctx: "RunContext[None]", + *, + request_context: "ModelRequestContext", + error: "Exception", + ) -> "ModelResponse": + run_context_metadata = ctx.metadata + + if not isinstance(run_context_metadata, dict): + raise error + + span = run_context_metadata.pop("_sentry_span", None) + if span is None: + raise error + + with capture_internal_exceptions(): + span.__exit__(type(error), error, error.__traceback__) + + raise error + + original_init = Agent.__init__ + + @functools.wraps(original_init) + def patched_init(self: "Agent[Any, Any]", *args: "Any", **kwargs: "Any") -> None: + caps = list(kwargs.get("capabilities") or []) + caps.append(hooks) + kwargs["capabilities"] = caps + + metadata = kwargs.get("metadata") + if metadata is None: + kwargs["metadata"] = {} # Used as shared reference between hooks + + return original_init(self, *args, **kwargs) + + Agent.__init__ = patched_init # type: ignore[method-assign] + + +def _patch_graph_nodes() -> None: + """ + Patches the graph node execution to create chat spans on pydantic-ai + versions that predate the request hooks. + + ModelRequestNode -> Creates ai_client span for model requests + """ + try: + from pydantic_ai._agent_graph import ModelRequestNode + except ImportError: + # Private module moved or renamed; degrade to agent + tool spans + # rather than crashing setup. + return + + # Patch ModelRequestNode to create ai_client spans + original_model_request_run = ModelRequestNode.run + + @wraps(original_model_request_run) + async def wrapped_model_request_run(self: "Any", ctx: "Any") -> "Any": + # Avoid creating a duplicate span if run() is invoked after stream(). + # This fails here: https://github.com/pydantic/pydantic-ai/blob/916fc83e8929470679db5ac1b3065bda5d5f4253/pydantic_ai_slim/pydantic_ai/_agent_graph.py#L1119 + did_stream = getattr(self, "_did_stream", False) + # Do not create a duplicate span when a cached result is served. + 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_graph_request_data(self, ctx) + + with ai_client_span(messages, None, model, model_settings) as span: + result = await original_model_request_run(self, ctx) + + # Extract response from result if available + model_response: "Optional[ModelResponse]" = None + if hasattr(result, "model_response"): + model_response = result.model_response + + update_ai_client_span(span, model_response) + return result + + ModelRequestNode.run = wrapped_model_request_run # type: ignore[method-assign] + + # Patch ModelRequestNode.stream for streaming requests + original_model_request_stream = ModelRequestNode.stream + + def create_wrapped_stream( + original_stream_method: "Callable[..., Any]", + ) -> "Callable[..., Any]": + """Create a wrapper for ModelRequestNode.stream that creates chat spans.""" + + @asynccontextmanager + @wraps(original_stream_method) + async def wrapped_model_request_stream(self: "Any", ctx: "Any") -> "Any": + # Avoid creating a duplicate span if the function is invoked twice. + # This fails here: https://github.com/pydantic/pydantic-ai/blob/916fc83e8929470679db5ac1b3065bda5d5f4253/pydantic_ai_slim/pydantic_ai/_agent_graph.py#L1128 + did_stream = getattr(self, "_did_stream", False) + if did_stream: + async with original_stream_method(self, ctx) as stream: + yield stream + return + + messages, model, model_settings = extract_graph_request_data(self, ctx) + + # Create chat span for streaming request + with ai_client_span(messages, None, model, model_settings) as span: + # Call the original stream method + async with original_stream_method(self, ctx) as stream: + yield stream + + # After streaming completes, update span with response data + # The ModelRequestNode stores the final response in _result + model_response: "Optional[ModelResponse]" = None + if hasattr(self, "_result") and self._result is not None: + # _result is a NextNode containing the model_response + if hasattr(self._result, "model_response"): + model_response = self._result.model_response + + update_ai_client_span(span, model_response) + + return wrapped_model_request_stream + + ModelRequestNode.stream = create_wrapped_stream(original_model_request_stream) # type: ignore[method-assign] diff --git a/sentry_sdk/integrations/pydantic_ai/_wrap_tools.py b/sentry_sdk/integrations/pydantic_ai/_wrap_tools.py new file mode 100644 index 0000000000..6ceb30810c --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/_wrap_tools.py @@ -0,0 +1,90 @@ +"""Instrumentation of tool execution. + +All tool calls in pydantic-ai flow through one ToolManager method (named +execute_tool_call on newer versions, _call_tool on older ones — resolved in +_compat), regardless of toolset type (function, MCP, combined, wrapper, ...). +Patching there avoids patching multiple toolset classes and dealing with +signature mismatches from instrumented MCP servers. +""" + +import sys +from functools import wraps +from typing import TYPE_CHECKING + +import sentry_sdk +from sentry_sdk.utils import capture_internal_exceptions, reraise + +from ._compat import TOOL_CALL_METHOD, ToolManager, ToolRetryError +from ._extract import extract_tool_call_args +from ._run_context import get_current_agent +from ._spans import _capture_exception, execute_tool_span, update_execute_tool_span + +if TYPE_CHECKING: + from typing import Any + + +def _patch_tool_execution() -> None: + if TOOL_CALL_METHOD is None: + # No known tool-call method on this version; skip tool instrumentation + # rather than crashing setup. + return + + original_method = getattr(ToolManager, TOOL_CALL_METHOD) + + @wraps(original_method) + async def wrapped_tool_call( + self: "Any", first_arg: "Any", *args: "Any", **kwargs: "Any" + ) -> "Any": + # execute_tool_call receives a validated wrapper holding the call; + # the older _call_tool receives the call directly. + if TOOL_CALL_METHOD == "execute_tool_call": + if not first_arg or not hasattr(first_arg, "call"): + return await original_method(self, first_arg, *args, **kwargs) + call = first_arg.call + else: + call = first_arg + + name = call.tool_name + tool = self.tools.get(name) if self.tools else None + selected_tool_definition = getattr(tool, "tool_def", None) + + # Get agent from contextvar + agent = get_current_agent() + + if not (agent and tool): + # No span context - just call original + return await original_method(self, first_arg, *args, **kwargs) + + args_dict = extract_tool_call_args(call) + + # Create execute_tool span + # Nesting is handled by isolation_scope() to ensure proper parent-child relationships + with sentry_sdk.isolation_scope(): + with execute_tool_span( + name, + args_dict, + agent, + tool_definition=selected_tool_definition, + ) as span: + try: + result = await original_method(self, first_arg, *args, **kwargs) + update_execute_tool_span(span, result) + return result + except ToolRetryError as exc: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + from sentry_sdk.integrations.pydantic_ai import ( + PydanticAIIntegration, + ) + + integration = sentry_sdk.get_client().get_integration( + PydanticAIIntegration + ) + if ( + integration is not None + and integration.handled_tool_call_exceptions + ): + _capture_exception(exc, handled=True) + reraise(*exc_info) + + setattr(ToolManager, TOOL_CALL_METHOD, wrapped_tool_call) 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 d0ea6242b4..0000000000 --- a/sentry_sdk/integrations/pydantic_ai/patches/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .agent_run import _patch_agent_run # noqa: F401 -from .graph_nodes import _patch_graph_nodes # noqa: F401 -from .tools import _patch_tool_execution # noqa: F401 diff --git a/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py b/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py deleted file mode 100644 index bb4ace697d..0000000000 --- a/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py +++ /dev/null @@ -1,100 +0,0 @@ -from contextlib import asynccontextmanager -from functools import wraps - -from sentry_sdk.integrations import DidNotEnable - -from .._extract import extract_graph_request_data -from ..spans import ( - ai_client_span, - update_ai_client_span, -) - -try: - from pydantic_ai._agent_graph import ModelRequestNode -except ImportError: - raise DidNotEnable("pydantic-ai not installed") - -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from typing import Any, Callable, Optional - - from pydantic_ai.messages import ModelResponse - - -def _patch_graph_nodes() -> None: - """ - Patches the graph node execution to create appropriate spans. - - ModelRequestNode -> Creates ai_client span for model requests - CallToolsNode -> Handles tool calls (spans created in tool patching) - """ - - # Patch ModelRequestNode to create ai_client spans - original_model_request_run = ModelRequestNode.run - - @wraps(original_model_request_run) - async def wrapped_model_request_run(self: "Any", ctx: "Any") -> "Any": - # Avoid creating a duplicate span if run() is invoked after stream(). - # This fails here: https://github.com/pydantic/pydantic-ai/blob/916fc83e8929470679db5ac1b3065bda5d5f4253/pydantic_ai_slim/pydantic_ai/_agent_graph.py#L1119 - did_stream = getattr(self, "_did_stream", False) - # Do not create a duplicate span when a cached result is served. - 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_graph_request_data(self, ctx) - - with ai_client_span(messages, None, model, model_settings) as span: - result = await original_model_request_run(self, ctx) - - # Extract response from result if available - model_response: "Optional[ModelResponse]" = None - if hasattr(result, "model_response"): - model_response = result.model_response - - update_ai_client_span(span, model_response) - return result - - ModelRequestNode.run = wrapped_model_request_run # type: ignore[method-assign] - - # Patch ModelRequestNode.stream for streaming requests - original_model_request_stream = ModelRequestNode.stream - - def create_wrapped_stream( - original_stream_method: "Callable[..., Any]", - ) -> "Callable[..., Any]": - """Create a wrapper for ModelRequestNode.stream that creates chat spans.""" - - @asynccontextmanager - @wraps(original_stream_method) - async def wrapped_model_request_stream(self: "Any", ctx: "Any") -> "Any": - # Avoid creating a duplicate span if the function is invoked twice. - # This fails here: https://github.com/pydantic/pydantic-ai/blob/916fc83e8929470679db5ac1b3065bda5d5f4253/pydantic_ai_slim/pydantic_ai/_agent_graph.py#L1128 - did_stream = getattr(self, "_did_stream", False) - if did_stream: - async with original_stream_method(self, ctx) as stream: - yield stream - return - - messages, model, model_settings = extract_graph_request_data(self, ctx) - - # Create chat span for streaming request - with ai_client_span(messages, None, model, model_settings) as span: - # Call the original stream method - async with original_stream_method(self, ctx) as stream: - yield stream - - # After streaming completes, update span with response data - # The ModelRequestNode stores the final response in _result - model_response: "Optional[ModelResponse]" = None - if hasattr(self, "_result") and self._result is not None: - # _result is a NextNode containing the model_response - if hasattr(self._result, "model_response"): - model_response = self._result.model_response - - update_ai_client_span(span, model_response) - - return wrapped_model_request_stream - - ModelRequestNode.stream = create_wrapped_stream(original_model_request_stream) # type: ignore[method-assign] diff --git a/sentry_sdk/integrations/pydantic_ai/patches/tools.py b/sentry_sdk/integrations/pydantic_ai/patches/tools.py deleted file mode 100644 index 81c7cf80e7..0000000000 --- a/sentry_sdk/integrations/pydantic_ai/patches/tools.py +++ /dev/null @@ -1,173 +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 .._extract import extract_tool_call_args -from .._run_context import get_current_agent -from ..spans import execute_tool_span, update_execute_tool_span -from ..utils import _capture_exception - -if TYPE_CHECKING: - from typing import Any - -try: - try: - from pydantic_ai.tool_manager import ToolManager - except ImportError: - from pydantic_ai._tool_manager import ToolManager # type: ignore - - from pydantic_ai.exceptions import ToolRetryError -except ImportError: - raise DidNotEnable("pydantic-ai not installed") - - -def _patch_tool_execution() -> None: - if hasattr(ToolManager, "execute_tool_call"): - _patch_execute_tool_call() - - elif hasattr(ToolManager, "_call_tool"): - # older versions - _patch_call_tool() - - -def _patch_execute_tool_call() -> None: - original_execute_tool_call = ToolManager.execute_tool_call - - @wraps(original_execute_tool_call) - async def wrapped_execute_tool_call( - self: "Any", validated: "Any", *args: "Any", **kwargs: "Any" - ) -> "Any": - if not validated or not hasattr(validated, "call"): - return await original_execute_tool_call(self, validated, *args, **kwargs) - - # Extract tool info before calling original - call = validated.call - name = call.tool_name - tool = self.tools.get(name) if self.tools else None - selected_tool_definition = getattr(tool, "tool_def", None) - - # Get agent from contextvar - agent = get_current_agent() - - if agent and tool: - args_dict = extract_tool_call_args(call) - - # Create execute_tool span - # Nesting is handled by isolation_scope() to ensure proper parent-child relationships - with sentry_sdk.isolation_scope(): - with execute_tool_span( - name, - args_dict, - agent, - tool_definition=selected_tool_definition, - ) as span: - try: - result = await original_execute_tool_call( - self, - validated, - *args, - **kwargs, - ) - update_execute_tool_span(span, result) - return result - except ToolRetryError as exc: - exc_info = sys.exc_info() - with capture_internal_exceptions(): - # Avoid circular import due to multi-file integration structure - from sentry_sdk.integrations.pydantic_ai import ( - PydanticAIIntegration, - ) - - integration = sentry_sdk.get_client().get_integration( - PydanticAIIntegration - ) - if ( - integration is not None - and integration.handled_tool_call_exceptions - ): - _capture_exception(exc, handled=True) - reraise(*exc_info) - - return await original_execute_tool_call(self, validated, *args, **kwargs) - - ToolManager.execute_tool_call = wrapped_execute_tool_call # type: ignore[method-assign] - - -def _patch_call_tool() -> None: - """ - Patch ToolManager._call_tool to create execute_tool spans. - - This is the single point where ALL tool calls flow through in pydantic_ai, - regardless of toolset type (function, MCP, combined, wrapper, etc.). - - By patching here, we avoid: - - Patching multiple toolset classes - - Dealing with signature mismatches from instrumented MCP servers - - Complex nested toolset handling - """ - original_call_tool = ToolManager._call_tool # type: ignore[attr-defined] - - @wraps(original_call_tool) - async def wrapped_call_tool( - self: "Any", call: "Any", *args: "Any", **kwargs: "Any" - ) -> "Any": - # Extract tool info before calling original - name = call.tool_name - tool = self.tools.get(name) if self.tools else None - selected_tool_definition = getattr(tool, "tool_def", None) - - # Get agent from contextvar - agent = get_current_agent() - - if agent and tool: - args_dict = extract_tool_call_args(call) - - # Create execute_tool span - # Nesting is handled by isolation_scope() to ensure proper parent-child relationships - with sentry_sdk.isolation_scope(): - with execute_tool_span( - name, - args_dict, - agent, - tool_definition=selected_tool_definition, - ) as span: - try: - result = await original_call_tool( - self, - call, - *args, - **kwargs, - ) - update_execute_tool_span(span, result) - return result - except ToolRetryError as exc: - exc_info = sys.exc_info() - with capture_internal_exceptions(): - # Avoid circular import due to multi-file integration structure - from sentry_sdk.integrations.pydantic_ai import ( - PydanticAIIntegration, - ) - - integration = sentry_sdk.get_client().get_integration( - PydanticAIIntegration - ) - if ( - integration is not None - and integration.handled_tool_call_exceptions - ): - _capture_exception(exc, handled=True) - reraise(*exc_info) - - # No span context - just call original - return await original_call_tool( - self, - call, - *args, - **kwargs, - ) - - ToolManager._call_tool = wrapped_call_tool # type: ignore[attr-defined] diff --git a/sentry_sdk/integrations/pydantic_ai/spans/__init__.py b/sentry_sdk/integrations/pydantic_ai/spans/__init__.py deleted file mode 100644 index 574046d645..0000000000 --- a/sentry_sdk/integrations/pydantic_ai/spans/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .ai_client import ai_client_span, update_ai_client_span # noqa: F401 -from .execute_tool import execute_tool_span, update_execute_tool_span # noqa: F401 -from .invoke_agent import invoke_agent_span, update_invoke_agent_span # noqa: F401 diff --git a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py deleted file mode 100644 index 4720a33726..0000000000 --- a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py +++ /dev/null @@ -1,172 +0,0 @@ -import json -from typing import TYPE_CHECKING - -import sentry_sdk -from sentry_sdk.ai.utils import ( - _set_span_data_attribute, - normalize_message_roles, - set_data_normalized, - truncate_and_annotate_messages, -) -from sentry_sdk.consts import OP, SPANDATA -from sentry_sdk.tracing_utils import ( - has_span_streaming_enabled, - should_truncate_gen_ai_input, -) - -from .._extract import ( - extract_model_info, - extract_request_messages, - extract_response_parts, - extract_system_instructions, -) -from .._run_context import get_current_agent, get_is_streaming -from ..consts import SPAN_ORIGIN -from ..utils import ( - _set_agent_data, - _set_available_tools, - _set_model_data, - _should_send_prompts, -) -from .utils import _set_usage_data - -if TYPE_CHECKING: - from typing import Any, Optional, Union - - from pydantic_ai.messages import ModelResponse - - from sentry_sdk.traces import StreamedSpan - - -def _set_input_messages( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", messages: "Any" -) -> None: - """Set input messages data on a span.""" - if not _should_send_prompts(): - return - - if not messages: - return - - try: - system_instructions = extract_system_instructions(messages) - if system_instructions: - _set_span_data_attribute( - span, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - json.dumps(system_instructions), - ) - - formatted_messages = extract_request_messages(messages) - - if formatted_messages: - normalized_messages = normalize_message_roles(formatted_messages) - client = sentry_sdk.get_client() - scope = sentry_sdk.get_current_scope() - messages_data = ( - truncate_and_annotate_messages(normalized_messages, span, scope) - if should_truncate_gen_ai_input(client.options) - else normalized_messages - ) - set_data_normalized( - span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False - ) - except Exception: - # If we fail to format messages, just skip it - pass - - -def _set_output_data( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", - response: "Optional[ModelResponse]", -) -> None: - """Set output data on a span.""" - if not _should_send_prompts(): - return - - if not response: - return - - if response.model_name: - _set_span_data_attribute( - span, SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name - ) - - try: - parts = extract_response_parts(response) - if parts: - _set_span_data_attribute( - span, - SPANDATA.GEN_AI_OUTPUT_MESSAGES, - json.dumps([{"role": "assistant", "parts": parts}]), - ) - except Exception: - # If we fail to format output, just skip it - pass - - -def ai_client_span( - messages: "Any", agent: "Any", model: "Any", model_settings: "Any" -) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]": - """Create a span for an AI client call (model request). - - Args: - messages: Full conversation history (list of messages) - agent: Agent object - model: Model object - model_settings: Model settings - """ - # Resolve the agent once so the span name and every attribute derived - # below (gen_ai.request.model, agent data, available tools) agree - agent_obj = agent or get_current_agent() - model_name = extract_model_info(model, model_settings, agent_obj).name or "unknown" - - span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) - if span_streaming: - span = sentry_sdk.traces.start_span( - name=f"chat {model_name}", - attributes={ - "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: - span = sentry_sdk.start_span( - op=OP.GEN_AI_CHAT, - name=f"chat {model_name}", - origin=SPAN_ORIGIN, - ) - - 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_obj) - _set_model_data(span, model, model_settings, agent=agent_obj) - - # Add available tools if agent is available - _set_available_tools(span, agent_obj) - - # Set input messages (full conversation history) - if messages: - _set_input_messages(span, messages) - - return span - - -def update_ai_client_span( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", - model_response: "Optional[ModelResponse]", -) -> None: - """Update the AI client span with response data.""" - if not span: - return - - # Set usage data if available - if model_response and hasattr(model_response, "usage"): - _set_usage_data(span, model_response.usage) - - # Set output data - _set_output_data(span, model_response) diff --git a/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py b/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py deleted file mode 100644 index bb4cfd8cc1..0000000000 --- a/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py +++ /dev/null @@ -1,84 +0,0 @@ -from typing import TYPE_CHECKING - -import sentry_sdk -from sentry_sdk.ai.utils import _set_span_data_attribute -from sentry_sdk.consts import OP, SPANDATA -from sentry_sdk.tracing_utils import has_span_streaming_enabled -from sentry_sdk.utils import safe_serialize - -from ..consts import SPAN_ORIGIN -from ..utils import _set_agent_data, _should_send_prompts - -if TYPE_CHECKING: - from typing import Any, Optional, Union - - from pydantic_ai._tool_manager import ToolDefinition # type: ignore - - from sentry_sdk.traces import StreamedSpan - - -def execute_tool_span( - tool_name: str, - tool_args: "Any", - agent: "Any", - tool_definition: "Optional[ToolDefinition]" = None, -) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]": - """Create a span for tool execution. - - Args: - tool_name: The name of the tool being executed - tool_args: The arguments passed to the tool - agent: The agent executing the tool - tool_definition: The definition of the tool, if available - """ - span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) - if span_streaming: - # Both keys must be present at span start so that attribute-based - # ignore_spans / traces_sampler rules can match this span. - span = sentry_sdk.traces.start_span( - name=f"execute_tool {tool_name}", - attributes={ - "sentry.op": OP.GEN_AI_EXECUTE_TOOL, - "sentry.origin": SPAN_ORIGIN, - SPANDATA.GEN_AI_OPERATION_NAME: "execute_tool", - SPANDATA.GEN_AI_TOOL_NAME: tool_name, - }, - ) - else: - span = sentry_sdk.start_span( - op=OP.GEN_AI_EXECUTE_TOOL, - name=f"execute_tool {tool_name}", - origin=SPAN_ORIGIN, - ) - - span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "execute_tool") - span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool_name) - - if tool_definition is not None and hasattr(tool_definition, "description"): - _set_span_data_attribute( - span, - SPANDATA.GEN_AI_TOOL_DESCRIPTION, - tool_definition.description, - ) - - _set_agent_data(span, agent) - - if _should_send_prompts() and tool_args is not None: - _set_span_data_attribute( - span, SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_args) - ) - - return span - - -def update_execute_tool_span( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", result: "Any" -) -> None: - """Update the execute tool span with the result.""" - if not span: - return - - if not _should_send_prompts() or result is None: - return - - _set_span_data_attribute(span, SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result)) diff --git a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py deleted file mode 100644 index fb795339b0..0000000000 --- a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py +++ /dev/null @@ -1,110 +0,0 @@ -from typing import TYPE_CHECKING - -import sentry_sdk -from sentry_sdk.ai.utils import ( - _set_span_data_attribute, - get_start_span_function, - normalize_message_roles, - set_data_normalized, - truncate_and_annotate_messages, -) -from sentry_sdk.consts import OP, SPANDATA -from sentry_sdk.tracing_utils import ( - has_span_streaming_enabled, - should_truncate_gen_ai_input, -) - -from .._extract import extract_agent_prompt_messages, extract_response_model_name -from ..consts import SPAN_ORIGIN -from ..utils import ( - _set_agent_data, - _set_available_tools, - _set_model_data, - _should_send_prompts, -) - -if TYPE_CHECKING: - from typing import Any, Union - - from sentry_sdk.traces import StreamedSpan - - -def invoke_agent_span( - user_prompt: "Any", - agent: "Any", - model: "Any", - model_settings: "Any", - is_streaming: bool = False, -) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]": - """Create a span for invoking the agent.""" - # Determine agent name for span - name = "agent" - if agent and getattr(agent, "name", None): - name = agent.name - - span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) - if span_streaming: - span = sentry_sdk.traces.start_span( - name=f"invoke_agent {name}", - attributes={ - "sentry.op": OP.GEN_AI_INVOKE_AGENT, - "sentry.origin": SPAN_ORIGIN, - SPANDATA.GEN_AI_OPERATION_NAME: "invoke_agent", - }, - ) - else: - span = get_start_span_function()( - op=OP.GEN_AI_INVOKE_AGENT, - name=f"invoke_agent {name}", - origin=SPAN_ORIGIN, - ) - - span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") - - _set_agent_data(span, agent) - _set_model_data(span, model, model_settings, agent=agent) - _set_available_tools(span, agent) - - # Add user prompt and system prompts if available and prompts are enabled - if _should_send_prompts(): - messages = extract_agent_prompt_messages(agent, user_prompt) - - if messages: - normalized_messages = normalize_message_roles(messages) - client = sentry_sdk.get_client() - scope = sentry_sdk.get_current_scope() - messages_data = ( - truncate_and_annotate_messages(normalized_messages, span, scope) - if should_truncate_gen_ai_input(client.options) - else normalized_messages - ) - set_data_normalized( - span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False - ) - - return span - - -def update_invoke_agent_span( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", - result: "Any", -) -> None: - """Update and close the invoke agent span.""" - if not span or not result: - return - - # Extract output from result - output = getattr(result, "output", None) - - # Set response text if prompts are enabled - if _should_send_prompts() and output: - set_data_normalized( - span, SPANDATA.GEN_AI_RESPONSE_TEXT, str(output), unpack=False - ) - - # Set model name from response if available - response_model_name = extract_response_model_name(result) - if response_model_name: - _set_span_data_attribute( - span, SPANDATA.GEN_AI_RESPONSE_MODEL, response_model_name - ) diff --git a/sentry_sdk/integrations/pydantic_ai/spans/utils.py b/sentry_sdk/integrations/pydantic_ai/spans/utils.py deleted file mode 100644 index 2b23a459c6..0000000000 --- a/sentry_sdk/integrations/pydantic_ai/spans/utils.py +++ /dev/null @@ -1,35 +0,0 @@ -"""Utility functions for PydanticAI span instrumentation.""" - -from typing import TYPE_CHECKING - -import sentry_sdk -from sentry_sdk.ai.monitoring import record_token_usage - -from .._extract import extract_usage_kwargs - -if TYPE_CHECKING: - from typing import Union - - from pydantic_ai.usage import RequestUsage, RunUsage - - from sentry_sdk.traces import StreamedSpan - - -def _set_usage_data( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", - usage: "Union[RequestUsage, RunUsage]", -) -> None: - """Set token usage data on a span. - - This function works with both RequestUsage (single request) and - RunUsage (agent run) objects from pydantic_ai. - - Args: - span: The Sentry span to set data on. - usage: RequestUsage or RunUsage object containing token usage information. - """ - usage_kwargs = extract_usage_kwargs(usage) - if usage_kwargs is None: - return - - record_token_usage(span, **usage_kwargs) diff --git a/sentry_sdk/integrations/pydantic_ai/utils.py b/sentry_sdk/integrations/pydantic_ai/utils.py deleted file mode 100644 index b4840603b9..0000000000 --- a/sentry_sdk/integrations/pydantic_ai/utils.py +++ /dev/null @@ -1,109 +0,0 @@ -from typing import TYPE_CHECKING - -import sentry_sdk -from sentry_sdk.ai.utils import _set_span_data_attribute -from sentry_sdk.consts import SPANDATA -from sentry_sdk.scope import should_send_default_pii -from sentry_sdk.utils import event_from_exception, safe_serialize - -from ._extract import ( - MODEL_SETTINGS_TO_SPANDATA, - extract_agent_name, - extract_available_tools, - extract_model_info, -) -from ._run_context import get_current_agent - -if TYPE_CHECKING: - from typing import Any, Union - - from sentry_sdk.traces import StreamedSpan - - -def _should_send_prompts() -> bool: - """ - Check if prompts should be sent to Sentry. - - This checks both send_default_pii and the include_prompts integration setting. - """ - if not should_send_default_pii(): - return False - - from . import PydanticAIIntegration - - # Get the integration instance from the client - integration = sentry_sdk.get_client().get_integration(PydanticAIIntegration) - - if integration is None: - return False - - return getattr(integration, "include_prompts", False) - - -def _set_agent_data( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", agent: "Any" -) -> 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) - """ - # Extract agent name from agent object or contextvar - agent_name = extract_agent_name(agent or get_current_agent()) - if agent_name: - _set_span_data_attribute(span, SPANDATA.GEN_AI_AGENT_NAME, agent_name) - - -def _set_model_data( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", - model: "Any", - model_settings: "Any", - agent: "Any" = None, -) -> None: - """Set model-related data on a span. - - Args: - span: The span to set data on - 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) - agent: Agent to fall back to for model and settings (defaults to the - agent of the current run) - """ - model_info = extract_model_info(model, model_settings, agent or get_current_agent()) - - if model_info.system is not None: - _set_span_data_attribute(span, SPANDATA.GEN_AI_SYSTEM, model_info.system) - - if model_info.name: - _set_span_data_attribute(span, SPANDATA.GEN_AI_REQUEST_MODEL, model_info.name) - - for setting_name, value in model_info.settings.items(): - spandata_key = MODEL_SETTINGS_TO_SPANDATA.get(setting_name) - if spandata_key is not None: - _set_span_data_attribute(span, spandata_key, value) - - -def _set_available_tools( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", agent: "Any" -) -> None: - """Set available tools data on a span from an agent's function toolset. - - Args: - span: The span to set data on - agent: Agent object with _function_toolset attribute - """ - tools = extract_available_tools(agent) - if tools: - _set_span_data_attribute( - span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools) - ) - - -def _capture_exception(exc: "Any", handled: bool = False) -> None: - event, hint = event_from_exception( - exc, - client_options=sentry_sdk.get_client().options, - mechanism={"type": "pydantic_ai", "handled": handled}, - ) - sentry_sdk.capture_event(event, hint=hint) diff --git a/tests/integrations/pydantic_ai/test_pydantic_ai.py b/tests/integrations/pydantic_ai/test_pydantic_ai.py index bcb0d36ab2..53d6e1fa34 100644 --- a/tests/integrations/pydantic_ai/test_pydantic_ai.py +++ b/tests/integrations/pydantic_ai/test_pydantic_ai.py @@ -24,8 +24,10 @@ from sentry_sdk._types import BLOB_DATA_SUBSTITUTE from sentry_sdk.consts import SPANDATA from sentry_sdk.integrations.pydantic_ai import PydanticAIIntegration -from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_input_messages -from sentry_sdk.integrations.pydantic_ai.spans.utils import _set_usage_data +from sentry_sdk.integrations.pydantic_ai._spans import ( + _set_input_messages, + _set_usage_data, +) from sentry_sdk.utils import package_version PYDANTIC_AI_VERSION = package_version("pydantic-ai") @@ -2360,7 +2362,7 @@ async def test_model_settings_object_style(sentry_init, capture_items): from unittest.mock import MagicMock import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.utils import _set_model_data + from sentry_sdk.integrations.pydantic_ai._spans import _set_model_data sentry_init( integrations=[PydanticAIIntegration()], @@ -2714,7 +2716,7 @@ async def test_update_invoke_agent_span_with_none_output(sentry_init, capture_it Test that update_invoke_agent_span handles None output gracefully. """ import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.spans.invoke_agent import ( + from sentry_sdk.integrations.pydantic_ai._spans import ( update_invoke_agent_span, ) @@ -2742,7 +2744,7 @@ async def test_update_ai_client_span_with_none_response(sentry_init, capture_ite Test that update_ai_client_span handles None response gracefully. """ import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.spans.ai_client import ( + from sentry_sdk.integrations.pydantic_ai._spans import ( update_ai_client_span, ) @@ -2855,7 +2857,7 @@ async def test_available_tools_error_handling(sentry_init, capture_items): from unittest.mock import MagicMock import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.utils import _set_available_tools + from sentry_sdk.integrations.pydantic_ai._spans import _set_available_tools sentry_init( integrations=[PydanticAIIntegration()], @@ -2884,7 +2886,7 @@ async def test_set_usage_data_with_none_usage(sentry_init, capture_items): Test that _set_usage_data handles None usage gracefully. """ import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_usage_data + from sentry_sdk.integrations.pydantic_ai._spans import _set_usage_data sentry_init( integrations=[PydanticAIIntegration()], @@ -2911,7 +2913,7 @@ async def test_set_usage_data_with_partial_fields(sentry_init, capture_items): from unittest.mock import MagicMock import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_usage_data + from sentry_sdk.integrations.pydantic_ai._spans import _set_usage_data sentry_init( integrations=[PydanticAIIntegration()], @@ -3203,7 +3205,7 @@ async def test_output_data_error_handling(sentry_init, capture_items): from unittest.mock import MagicMock import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_output_data + from sentry_sdk.integrations.pydantic_ai._spans import _set_output_data sentry_init( integrations=[PydanticAIIntegration()], @@ -3399,7 +3401,7 @@ async def test_set_model_data_with_system(sentry_init, capture_items): from unittest.mock import MagicMock import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.utils import _set_model_data + from sentry_sdk.integrations.pydantic_ai._spans import _set_model_data sentry_init( integrations=[PydanticAIIntegration()], @@ -3431,7 +3433,7 @@ async def test_set_model_data_from_agent_scope(sentry_init, capture_items): from unittest.mock import MagicMock import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.utils import _set_model_data + from sentry_sdk.integrations.pydantic_ai._spans import _set_model_data sentry_init( integrations=[PydanticAIIntegration()], @@ -3464,7 +3466,7 @@ async def test_set_model_data_with_none_settings_values(sentry_init, capture_ite Test that _set_model_data skips None values in settings. """ import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.utils import _set_model_data + from sentry_sdk.integrations.pydantic_ai._spans import _set_model_data sentry_init( integrations=[PydanticAIIntegration()], @@ -3495,7 +3497,7 @@ async def test_should_send_prompts_without_pii(sentry_init, capture_items): """ Test that _should_send_prompts returns False when PII disabled. """ - from sentry_sdk.integrations.pydantic_ai.utils import _should_send_prompts + from sentry_sdk.integrations.pydantic_ai._spans import _should_send_prompts sentry_init( integrations=[PydanticAIIntegration(include_prompts=True)], @@ -3514,7 +3516,7 @@ async def test_set_agent_data_without_agent(sentry_init, capture_items): Test that _set_agent_data handles None agent gracefully. """ import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.utils import _set_agent_data + from sentry_sdk.integrations.pydantic_ai._spans import _set_agent_data sentry_init( integrations=[PydanticAIIntegration()], @@ -3541,7 +3543,7 @@ async def test_set_agent_data_from_scope(sentry_init, capture_items): from unittest.mock import MagicMock import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.utils import _set_agent_data + from sentry_sdk.integrations.pydantic_ai._spans import _set_agent_data sentry_init( integrations=[PydanticAIIntegration()], @@ -3574,7 +3576,7 @@ async def test_set_agent_data_without_name(sentry_init, capture_items): from unittest.mock import MagicMock import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.utils import _set_agent_data + from sentry_sdk.integrations.pydantic_ai._spans import _set_agent_data sentry_init( integrations=[PydanticAIIntegration()], @@ -3605,7 +3607,7 @@ async def test_set_available_tools_without_toolset(sentry_init, capture_items): from unittest.mock import MagicMock import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.utils import _set_available_tools + from sentry_sdk.integrations.pydantic_ai._spans import _set_available_tools sentry_init( integrations=[PydanticAIIntegration()], @@ -3636,7 +3638,7 @@ async def test_set_available_tools_with_schema(sentry_init, capture_items): from unittest.mock import MagicMock import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.utils import _set_available_tools + from sentry_sdk.integrations.pydantic_ai._spans import _set_available_tools sentry_init( integrations=[PydanticAIIntegration()], @@ -3671,7 +3673,7 @@ async def test_execute_tool_span_creation(sentry_init, capture_items): Test direct creation of execute_tool span. """ import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.spans.execute_tool import ( + from sentry_sdk.integrations.pydantic_ai._spans import ( execute_tool_span, update_execute_tool_span, ) @@ -3698,7 +3700,7 @@ async def test_execute_tool_span_with_mcp_type(sentry_init, capture_items): Test execute_tool span with MCP tool type. """ import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.spans.execute_tool import ( + from sentry_sdk.integrations.pydantic_ai._spans import ( execute_tool_span, ) @@ -3724,7 +3726,7 @@ async def test_execute_tool_span_without_prompts(sentry_init, capture_items): Test that execute_tool span respects _should_send_prompts(). """ import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.spans.execute_tool import ( + from sentry_sdk.integrations.pydantic_ai._spans import ( execute_tool_span, update_execute_tool_span, ) @@ -3751,7 +3753,7 @@ async def test_execute_tool_span_with_none_args(sentry_init, capture_items): Test execute_tool span with None args. """ import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.spans.execute_tool import execute_tool_span + from sentry_sdk.integrations.pydantic_ai._spans import execute_tool_span sentry_init( integrations=[PydanticAIIntegration()], @@ -3773,7 +3775,7 @@ async def test_update_execute_tool_span_with_none_span(sentry_init, capture_item """ Test that update_execute_tool_span handles None span gracefully. """ - from sentry_sdk.integrations.pydantic_ai.spans.execute_tool import ( + from sentry_sdk.integrations.pydantic_ai._spans import ( update_execute_tool_span, ) @@ -3795,7 +3797,7 @@ async def test_update_execute_tool_span_with_none_result(sentry_init, capture_it Test that update_execute_tool_span handles None result gracefully. """ import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.spans.execute_tool import ( + from sentry_sdk.integrations.pydantic_ai._spans import ( execute_tool_span, update_execute_tool_span, ) @@ -3853,7 +3855,7 @@ async def test_invoke_agent_span_with_callable_instruction(sentry_init, capture_ from unittest.mock import MagicMock import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.spans.invoke_agent import invoke_agent_span + from sentry_sdk.integrations.pydantic_ai._spans import invoke_agent_span sentry_init( integrations=[PydanticAIIntegration()], @@ -3887,7 +3889,7 @@ async def test_invoke_agent_span_with_string_instructions(sentry_init, capture_i from unittest.mock import MagicMock import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.spans.invoke_agent import invoke_agent_span + from sentry_sdk.integrations.pydantic_ai._spans import invoke_agent_span sentry_init( integrations=[PydanticAIIntegration()], @@ -3916,7 +3918,7 @@ async def test_ai_client_span_with_streaming_flag(sentry_init, capture_items): Test that ai_client_span reads streaming flag from scope. """ import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.spans.ai_client import ai_client_span + from sentry_sdk.integrations.pydantic_ai._spans import ai_client_span sentry_init( integrations=[PydanticAIIntegration()], @@ -3944,7 +3946,7 @@ async def test_ai_client_span_gets_agent_from_scope(sentry_init, capture_items): from unittest.mock import MagicMock import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.spans.ai_client import ai_client_span + from sentry_sdk.integrations.pydantic_ai._spans import ai_client_span sentry_init( integrations=[PydanticAIIntegration()], @@ -4392,7 +4394,7 @@ def test_image_url_base64_content_in_span( stream_gen_ai_spans, span_streaming, ): - from sentry_sdk.integrations.pydantic_ai.spans.ai_client import ai_client_span + from sentry_sdk.integrations.pydantic_ai._spans import ai_client_span sentry_init( integrations=[PydanticAIIntegration()],