diff --git a/sentry_sdk/integrations/pydantic_ai/__init__.py b/sentry_sdk/integrations/pydantic_ai/__init__.py index db21861e71..a3215aa284 100644 --- a/sentry_sdk/integrations/pydantic_ai/__init__.py +++ b/sentry_sdk/integrations/pydantic_ai/__init__.py @@ -1,7 +1,7 @@ import functools from sentry_sdk.integrations import DidNotEnable, Integration -from sentry_sdk.utils import capture_internal_exceptions, parse_version +from sentry_sdk.utils import parse_version try: import pydantic_ai # noqa: F401 @@ -24,7 +24,7 @@ from typing import Any from pydantic_ai import ModelRequestContext, RunContext - from pydantic_ai.capabilities import Hooks + from pydantic_ai.capabilities import Hooks, WrapModelRequestHandler from pydantic_ai.messages import ModelResponse @@ -33,67 +33,24 @@ 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__()`. """ - @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 = ai_client_span( - messages=request_context.messages, - agent=None, - model=request_context.model, - model_settings=request_context.model_settings, - ) - - run_context_metadata["_sentry_span"] = span - span.__enter__() - - return request_context - - @hooks.on.after_model_request - async def on_response( + @hooks.on.model_request + async def on_model_request( ctx: "RunContext[None]", *, request_context: "ModelRequestContext", - response: "ModelResponse", + handler: "WrapModelRequestHandler", ) -> "ModelResponse": - run_context_metadata = ctx.metadata - if not isinstance(run_context_metadata, dict): - return response + with ai_client_span( + messages=request_context.messages, + agent=None, + model=request_context.model, + model_settings=request_context.model_settings, + ) as span: + response = await handler(request_context) - span = run_context_metadata.pop("_sentry_span", None) - if span is None: + update_ai_client_span(span, response) return response - 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) @@ -102,10 +59,6 @@ def patched_init(self: "Agent[Any, Any]", *args: "Any", **kwargs: "Any") -> None 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] @@ -129,7 +82,6 @@ class PydanticAIIntegration(Integration): 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 @@ -159,7 +111,6 @@ def setup_once() -> None: _patch_agent_run() _patch_tool_execution() - PydanticAIIntegration.using_request_hooks = False try: PYDANTIC_AI_VERSION = version("pydantic-ai-slim") except PackageNotFoundError: @@ -182,6 +133,5 @@ def setup_once() -> None: except ImportError: return - PydanticAIIntegration.using_request_hooks = True hooks = Hooks() register_hooks(hooks) diff --git a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py index 864c83a506..b15b27f802 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py @@ -95,9 +95,6 @@ def _create_run_wrapper( original_func: The original run method is_streaming: Whether this is a streaming method (for future use) """ - 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": @@ -109,11 +106,6 @@ async def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": model = kwargs.get("model") model_settings = kwargs.get("model_settings") - if PydanticAIIntegration.using_request_hooks: - metadata = kwargs.get("metadata") - if metadata is None: - kwargs["metadata"] = {"_sentry_span": None} - # Create invoke_agent span with invoke_agent_span( user_prompt, self, model, model_settings, is_streaming @@ -147,9 +139,6 @@ 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": @@ -158,11 +147,6 @@ def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": model = kwargs.get("model") model_settings = kwargs.get("model_settings") - if PydanticAIIntegration.using_request_hooks: - metadata = kwargs.get("metadata") - if metadata is None: - kwargs["metadata"] = {"_sentry_span": None} - # Call original function to get the context manager original_ctx_manager = original_func(self, *args, **kwargs)