diff --git a/sentry_sdk/integrations/pydantic_ai/__init__.py b/sentry_sdk/integrations/pydantic_ai/__init__.py index 11dd171944..2f1808d14f 100644 --- a/sentry_sdk/integrations/pydantic_ai/__init__.py +++ b/sentry_sdk/integrations/pydantic_ai/__init__.py @@ -19,15 +19,20 @@ class PydanticAIIntegration(Integration): identifier = "pydantic_ai" origin = f"auto.ai.{identifier}" - def __init__(self, include_prompts: bool = True) -> None: + def __init__( + self, include_prompts: bool = True, handled_tool_call_exceptions: bool = True + ) -> None: """ Initialize the Pydantic AI integration. Args: include_prompts: Whether to include prompts and messages in span data. Requires send_default_pii=True. Defaults to True. + handled_tool_exceptions: Capture tool call exceptions that Pydantic AI + internally prevents from bubbling up. """ self.include_prompts = include_prompts + self.handled_tool_call_exceptions = handled_tool_call_exceptions @staticmethod def setup_once() -> None: diff --git a/sentry_sdk/integrations/pydantic_ai/patches/tools.py b/sentry_sdk/integrations/pydantic_ai/patches/tools.py index e4251d671c..b826a543fc 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/tools.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/tools.py @@ -4,10 +4,7 @@ import sentry_sdk from ..spans import execute_tool_span, update_execute_tool_span -from ..utils import ( - _capture_exception, - get_current_agent, -) +from ..utils import _capture_exception, get_current_agent from typing import TYPE_CHECKING @@ -23,6 +20,7 @@ try: from pydantic_ai._tool_manager import ToolManager # type: ignore + from pydantic_ai.exceptions import ToolRetryError # type: ignore except ImportError: raise DidNotEnable("pydantic-ai not installed") @@ -82,8 +80,21 @@ async def wrapped_call_tool( ) update_execute_tool_span(span, result) return result - except Exception as exc: - _capture_exception(exc) + except ToolRetryError as exc: + # 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 None + or not integration.handled_tool_call_exceptions + ): + raise exc from None + _capture_exception(exc, handled=True) raise exc from None # No span context - just call original diff --git a/sentry_sdk/integrations/pydantic_ai/utils.py b/sentry_sdk/integrations/pydantic_ai/utils.py index 743f3078f2..62d36fb912 100644 --- a/sentry_sdk/integrations/pydantic_ai/utils.py +++ b/sentry_sdk/integrations/pydantic_ai/utils.py @@ -206,12 +206,12 @@ def _set_available_tools(span: "sentry_sdk.tracing.Span", agent: "Any") -> None: pass -def _capture_exception(exc: "Any") -> None: +def _capture_exception(exc: "Any", handled: bool = False) -> None: set_span_errored() event, hint = event_from_exception( exc, client_options=sentry_sdk.get_client().options, - mechanism={"type": "pydantic_ai", "handled": False}, + 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 7f81769407..049bcde39c 100644 --- a/tests/integrations/pydantic_ai/test_pydantic_ai.py +++ b/tests/integrations/pydantic_ai/test_pydantic_ai.py @@ -1,10 +1,14 @@ import asyncio import pytest +from typing import Annotated +from pydantic import Field + from sentry_sdk.integrations.pydantic_ai import PydanticAIIntegration from pydantic_ai import Agent from pydantic_ai.models.test import TestModel +from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior @pytest.fixture @@ -277,6 +281,157 @@ def add_numbers(a: int, b: int) -> int: assert "add_numbers" in available_tools_str +@pytest.mark.parametrize( + "handled_tool_call_exceptions", + [False, True], +) +@pytest.mark.asyncio +async def test_agent_with_tool_model_retry( + sentry_init, capture_events, test_agent, handled_tool_call_exceptions +): + """ + Test that a handled exception is captured when a tool raises ModelRetry. + """ + + retries = 0 + + @test_agent.tool_plain + def add_numbers(a: int, b: int) -> float: + """Add two numbers together, but raises an exception on the first attempt.""" + nonlocal retries + if retries == 0: + retries += 1 + raise ModelRetry(message="Try again with the same arguments.") + return a + b + + sentry_init( + integrations=[ + PydanticAIIntegration( + handled_tool_call_exceptions=handled_tool_call_exceptions + ) + ], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + events = capture_events() + + result = await test_agent.run("What is 5 + 3?") + + assert result is not None + + if handled_tool_call_exceptions: + (error, transaction) = events + else: + (transaction,) = events + spans = transaction["spans"] + + if handled_tool_call_exceptions: + assert error["level"] == "error" + assert error["exception"]["values"][0]["mechanism"]["handled"] + + # Find child span types (invoke_agent is the transaction, not a child span) + chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] + + # Should have tool spans + assert len(tool_spans) >= 1 + + # Check tool spans + model_retry_tool_span = tool_spans[0] + assert "execute_tool" in model_retry_tool_span["description"] + assert model_retry_tool_span["data"]["gen_ai.operation.name"] == "execute_tool" + assert model_retry_tool_span["data"]["gen_ai.tool.type"] == "function" + assert model_retry_tool_span["data"]["gen_ai.tool.name"] == "add_numbers" + assert "gen_ai.tool.input" in model_retry_tool_span["data"] + + tool_span = tool_spans[1] + assert "execute_tool" in tool_span["description"] + assert tool_span["data"]["gen_ai.operation.name"] == "execute_tool" + assert tool_span["data"]["gen_ai.tool.type"] == "function" + assert tool_span["data"]["gen_ai.tool.name"] == "add_numbers" + assert "gen_ai.tool.input" in tool_span["data"] + assert "gen_ai.tool.output" in tool_span["data"] + + # Check chat spans have available_tools + for chat_span in chat_spans: + assert "gen_ai.request.available_tools" in chat_span["data"] + available_tools_str = chat_span["data"]["gen_ai.request.available_tools"] + # Available tools is serialized as a string + assert "add_numbers" in available_tools_str + + +@pytest.mark.parametrize( + "handled_tool_call_exceptions", + [False, True], +) +@pytest.mark.asyncio +async def test_agent_with_tool_validation_error( + sentry_init, capture_events, test_agent, handled_tool_call_exceptions +): + """ + Test that a handled exception is captured when a tool has unsatisfiable constraints. + """ + + @test_agent.tool_plain + def add_numbers(a: Annotated[int, Field(gt=0, lt=0)], b: int) -> int: + """Add two numbers together.""" + return a + b + + sentry_init( + integrations=[ + PydanticAIIntegration( + handled_tool_call_exceptions=handled_tool_call_exceptions + ) + ], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + events = capture_events() + + result = None + with pytest.raises(UnexpectedModelBehavior): + result = await test_agent.run("What is 5 + 3?") + + assert result is None + + if handled_tool_call_exceptions: + (error, model_behaviour_error, transaction) = events + else: + ( + model_behaviour_error, + transaction, + ) = events + spans = transaction["spans"] + + if handled_tool_call_exceptions: + assert error["level"] == "error" + assert error["exception"]["values"][0]["mechanism"]["handled"] + + # Find child span types (invoke_agent is the transaction, not a child span) + chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] + + # Should have tool spans + assert len(tool_spans) >= 1 + + # Check tool spans + model_retry_tool_span = tool_spans[0] + assert "execute_tool" in model_retry_tool_span["description"] + assert model_retry_tool_span["data"]["gen_ai.operation.name"] == "execute_tool" + assert model_retry_tool_span["data"]["gen_ai.tool.type"] == "function" + assert model_retry_tool_span["data"]["gen_ai.tool.name"] == "add_numbers" + assert "gen_ai.tool.input" in model_retry_tool_span["data"] + + # Check chat spans have available_tools + for chat_span in chat_spans: + assert "gen_ai.request.available_tools" in chat_span["data"] + available_tools_str = chat_span["data"]["gen_ai.request.available_tools"] + # Available tools is serialized as a string + assert "add_numbers" in available_tools_str + + @pytest.mark.asyncio async def test_agent_with_tools_streaming(sentry_init, capture_events, test_agent): """