diff --git a/.github/workflows/test-integrations-ai.yml b/.github/workflows/test-integrations-ai.yml index 1b9a341f17..c8b5e70151 100644 --- a/.github/workflows/test-integrations-ai.yml +++ b/.github/workflows/test-integrations-ai.yml @@ -94,6 +94,10 @@ jobs: run: | set -x # print commands that are executed ./scripts/runtox.sh "py${{ matrix.python-version }}-openai_agents" + - name: Test pydantic_ai + run: | + set -x # print commands that are executed + ./scripts/runtox.sh "py${{ matrix.python-version }}-pydantic_ai" - name: Generate coverage XML (Python 3.6) if: ${{ !cancelled() && matrix.python-version == '3.6' }} run: | diff --git a/scripts/populate_tox/config.py b/scripts/populate_tox/config.py index 1f23b3fb08..45ee32348f 100644 --- a/scripts/populate_tox/config.py +++ b/scripts/populate_tox/config.py @@ -272,6 +272,12 @@ "package": "pure_eval", "num_versions": 2, }, + "pydantic_ai": { + "package": "pydantic-ai", + "deps": { + "*": ["pytest-asyncio"], + }, + }, "pymongo": { "package": "pymongo", "deps": { diff --git a/scripts/populate_tox/releases.jsonl b/scripts/populate_tox/releases.jsonl index ac5fe1de14..0c4bc3f974 100644 --- a/scripts/populate_tox/releases.jsonl +++ b/scripts/populate_tox/releases.jsonl @@ -129,6 +129,7 @@ {"info": {"classifiers": ["License :: OSI Approved :: Apache Software License", "Programming Language :: Python", "Programming Language :: Python :: 3"], "name": "openfeature-sdk", "requires_python": ">=3.9", "version": "0.8.3", "yanked": false}} {"info": {"classifiers": ["Intended Audience :: Developers", "License :: OSI Approved :: MIT License", "Operating System :: OS Independent", "Programming Language :: Python :: 3.5", "Programming Language :: Python :: 3.6", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8"], "name": "pure-eval", "requires_python": "", "version": "0.0.3", "yanked": false}} {"info": {"classifiers": ["Intended Audience :: Developers", "License :: OSI Approved :: MIT License", "Operating System :: OS Independent", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9"], "name": "pure-eval", "requires_python": null, "version": "0.2.3", "yanked": false}} +{"info": {"classifiers": ["Development Status :: 5 - Production/Stable", "Framework :: Pydantic", "Framework :: Pydantic :: 2", "Intended Audience :: Developers", "Intended Audience :: Information Technology", "Operating System :: OS Independent", "Programming Language :: Python", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3 :: Only", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Topic :: Internet", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development :: Libraries :: Python Modules"], "name": "pydantic-ai", "requires_python": ">=3.10", "version": "1.0.17", "yanked": false}} {"info": {"classifiers": ["Development Status :: 3 - Alpha", "Intended Audience :: Developers", "License :: OSI Approved :: Apache Software License", "Operating System :: MacOS :: MacOS X", "Operating System :: Microsoft :: Windows", "Operating System :: POSIX", "Programming Language :: Python", "Topic :: Database"], "name": "pymongo", "requires_python": null, "version": "0.6", "yanked": false}} {"info": {"classifiers": ["Development Status :: 5 - Production/Stable", "Intended Audience :: Developers", "License :: OSI Approved :: Apache Software License", "Operating System :: MacOS :: MacOS X", "Operating System :: Microsoft :: Windows", "Operating System :: POSIX", "Programming Language :: Python :: 2", "Programming Language :: Python :: 2.4", "Programming Language :: Python :: 2.5", "Programming Language :: Python :: 2.6", "Programming Language :: Python :: 2.7", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.1", "Programming Language :: Python :: 3.2", "Programming Language :: Python :: 3.3", "Programming Language :: Python :: 3.4", "Programming Language :: Python :: Implementation :: CPython", "Programming Language :: Python :: Implementation :: Jython", "Programming Language :: Python :: Implementation :: PyPy", "Topic :: Database"], "name": "pymongo", "requires_python": null, "version": "2.8.1", "yanked": false}} {"info": {"classifiers": ["Development Status :: 5 - Production/Stable", "Intended Audience :: Developers", "License :: OSI Approved :: Apache Software License", "Operating System :: MacOS :: MacOS X", "Operating System :: Microsoft :: Windows", "Operating System :: POSIX", "Programming Language :: Python :: 2", "Programming Language :: Python :: 2.7", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.5", "Programming Language :: Python :: 3.6", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: Implementation :: CPython", "Programming Language :: Python :: Implementation :: PyPy", "Topic :: Database"], "name": "pymongo", "requires_python": "", "version": "3.13.0", "yanked": false}} diff --git a/scripts/split_tox_gh_actions/split_tox_gh_actions.py b/scripts/split_tox_gh_actions/split_tox_gh_actions.py index 9dea95842b..241da35ce8 100755 --- a/scripts/split_tox_gh_actions/split_tox_gh_actions.py +++ b/scripts/split_tox_gh_actions/split_tox_gh_actions.py @@ -81,6 +81,7 @@ "openai-base", "openai-notiktoken", "openai_agents", + "pydantic_ai", ], "Cloud": [ "aws_lambda", diff --git a/sentry_sdk/integrations/__init__.py b/sentry_sdk/integrations/__init__.py index 9e279b8345..45a0f53eaf 100644 --- a/sentry_sdk/integrations/__init__.py +++ b/sentry_sdk/integrations/__init__.py @@ -152,6 +152,7 @@ def iter_default_integrations(with_auto_enabling_integrations): "openai": (1, 0, 0), "openai_agents": (0, 0, 19), "openfeature": (0, 7, 1), + "pydantic_ai": (1, 0, 0), "quart": (0, 16, 0), "ray": (2, 7, 0), "requests": (2, 0, 0), diff --git a/sentry_sdk/integrations/pydantic_ai/__init__.py b/sentry_sdk/integrations/pydantic_ai/__init__.py new file mode 100644 index 0000000000..733148dd08 --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/__init__.py @@ -0,0 +1,47 @@ +from sentry_sdk.integrations import DidNotEnable, Integration + +from .patches import ( + _patch_agent_run, + _patch_graph_nodes, + _patch_model_request, + _patch_tool_execution, +) + +try: + import pydantic_ai + +except ImportError: + raise DidNotEnable("pydantic-ai not installed") + + +class PydanticAIIntegration(Integration): + identifier = "pydantic_ai" + origin = f"auto.ai.{identifier}" + + def __init__(self, include_prompts=True): + # type: (bool) -> 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. + """ + self.include_prompts = include_prompts + + @staticmethod + def setup_once(): + # type: () -> None + """ + Set up the pydantic-ai integration. + + This patches the key methods in pydantic-ai to create Sentry spans for: + - Agent workflow execution (root span) + - Individual agent invocations + - Model requests (AI client calls) + - Tool executions + """ + _patch_agent_run() + _patch_graph_nodes() + _patch_model_request() + _patch_tool_execution() diff --git a/sentry_sdk/integrations/pydantic_ai/consts.py b/sentry_sdk/integrations/pydantic_ai/consts.py new file mode 100644 index 0000000000..afa66dc47d --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/consts.py @@ -0,0 +1 @@ +SPAN_ORIGIN = "auto.ai.pydantic_ai" diff --git a/sentry_sdk/integrations/pydantic_ai/patches/__init__.py b/sentry_sdk/integrations/pydantic_ai/patches/__init__.py new file mode 100644 index 0000000000..de28780728 --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/patches/__init__.py @@ -0,0 +1,4 @@ +from .agent_run import _patch_agent_run # noqa: F401 +from .graph_nodes import _patch_graph_nodes # noqa: F401 +from .model_request import _patch_model_request # noqa: F401 +from .tools import _patch_tool_execution # noqa: F401 diff --git a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py new file mode 100644 index 0000000000..84d1ac2bfd --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py @@ -0,0 +1,222 @@ +from functools import wraps + +import sentry_sdk +from sentry_sdk.integrations import DidNotEnable + +from ..spans import agent_workflow_span, invoke_agent_span, update_invoke_agent_span +from ..utils import _capture_exception + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Callable, Optional + +try: + import pydantic_ai +except ImportError: + raise DidNotEnable("pydantic-ai not installed") + + +class _StreamingContextManagerWrapper: + """Wrapper for streaming methods that return async context managers.""" + + def __init__(self, agent, original_ctx_manager, is_streaming=True): + # type: (Any, Any, bool) -> None + self.agent = agent + self.original_ctx_manager = original_ctx_manager + self.is_streaming = is_streaming + self._isolation_scope = None # type: Any + self._workflow_span = None # type: Optional[sentry_sdk.tracing.Span] + + async def __aenter__(self): + # type: () -> Any + # Set up isolation scope and workflow span + self._isolation_scope = sentry_sdk.isolation_scope() + self._isolation_scope.__enter__() + + # Store agent reference and streaming flag + sentry_sdk.get_current_scope().set_context( + "pydantic_ai_agent", {"_agent": self.agent, "_streaming": self.is_streaming} + ) + + # Create workflow span + self._workflow_span = agent_workflow_span(self.agent) + self._workflow_span.__enter__() + + # Enter the original context manager + result = await self.original_ctx_manager.__aenter__() + return result + + async def __aexit__(self, exc_type, exc_val, exc_tb): + # type: (Any, Any, Any) -> None + try: + # Exit the original context manager first + await self.original_ctx_manager.__aexit__(exc_type, exc_val, exc_tb) + finally: + # Clean up workflow span + if self._workflow_span: + self._workflow_span.__exit__(exc_type, exc_val, exc_tb) + + # Clean up isolation scope + if self._isolation_scope: + self._isolation_scope.__exit__(exc_type, exc_val, exc_tb) + + +def _create_run_wrapper(original_func, is_streaming=False): + # type: (Callable[..., Any], bool) -> Callable[..., Any] + """ + Wraps the Agent.run method to create a root span for the agent workflow. + + Args: + original_func: The original run method + is_streaming: Whether this is a streaming method (for future use) + """ + + @wraps(original_func) + async def wrapper(self, *args, **kwargs): + # type: (Any, *Any, **Any) -> Any + # Isolate each workflow so that when agents are run in asyncio tasks they + # don't touch each other's scopes + with sentry_sdk.isolation_scope(): + # Store agent reference and streaming flag in Sentry scope for access in nested spans + # We store the full agent to allow access to tools and system prompts + sentry_sdk.get_current_scope().set_context( + "pydantic_ai_agent", {"_agent": self, "_streaming": is_streaming} + ) + + with agent_workflow_span(self): + result = None + try: + result = await original_func(self, *args, **kwargs) + return result + except Exception as exc: + _capture_exception(exc) + + # It could be that there is an "invoke agent" span still open + current_span = sentry_sdk.get_current_span() + if current_span is not None and current_span.timestamp is None: + current_span.__exit__(None, None, None) + + raise exc from None + + return wrapper + + +def _create_run_sync_wrapper(original_func): + # type: (Callable[..., Any]) -> Callable[..., Any] + """ + Wraps the Agent.run_sync method to create a root span for the agent workflow. + Note: run_sync is always non-streaming. + """ + + @wraps(original_func) + def wrapper(self, *args, **kwargs): + # type: (Any, *Any, **Any) -> Any + # Isolate each workflow so that when agents are run they + # don't touch each other's scopes + with sentry_sdk.isolation_scope(): + # Store agent reference and streaming flag in Sentry scope for access in nested spans + # We store the full agent to allow access to tools and system prompts + sentry_sdk.get_current_scope().set_context( + "pydantic_ai_agent", {"_agent": self, "_streaming": False} + ) + + with agent_workflow_span(self): + result = None + try: + result = original_func(self, *args, **kwargs) + return result + except Exception as exc: + _capture_exception(exc) + + # It could be that there is an "invoke agent" span still open + current_span = sentry_sdk.get_current_span() + if current_span is not None and current_span.timestamp is None: + current_span.__exit__(None, None, None) + + raise exc from None + + return wrapper + + +def _create_streaming_wrapper(original_func): + # type: (Callable[..., Any]) -> Callable[..., Any] + """ + Wraps run_stream method that returns an async context manager. + """ + + @wraps(original_func) + def wrapper(self, *args, **kwargs): + # type: (Any, *Any, **Any) -> Any + # Call original function to get the context manager + original_ctx_manager = original_func(self, *args, **kwargs) + + # Wrap it with our instrumentation + return _StreamingContextManagerWrapper( + agent=self, original_ctx_manager=original_ctx_manager, is_streaming=True + ) + + return wrapper + + +def _create_streaming_events_wrapper(original_func): + # type: (Callable[..., Any]) -> Callable[..., Any] + """ + Wraps run_stream_events method that returns an async generator/iterator. + """ + + @wraps(original_func) + async def wrapper(self, *args, **kwargs): + # type: (Any, *Any, **Any) -> Any + # Isolate each workflow so that when agents are run in asyncio tasks they + # don't touch each other's scopes + with sentry_sdk.isolation_scope(): + # Store agent reference and streaming flag in Sentry scope for access in nested spans + sentry_sdk.get_current_scope().set_context( + "pydantic_ai_agent", {"_agent": self, "_streaming": True} + ) + + with agent_workflow_span(self): + try: + # Call the original generator and yield all events + async for event in original_func(self, *args, **kwargs): + yield event + except Exception as exc: + _capture_exception(exc) + + # It could be that there is an "invoke agent" span still open + current_span = sentry_sdk.get_current_span() + if current_span is not None and current_span.timestamp is None: + current_span.__exit__(None, None, None) + + raise exc from None + + return wrapper + + +def _patch_agent_run(): + # type: () -> None + """ + Patches the Agent run methods to create spans for agent execution. + + This patches both non-streaming (run, run_sync) and streaming + (run_stream, run_stream_events) methods. + """ + # Import here to avoid circular imports + from pydantic_ai.agent import Agent + + # Store original methods + original_run = Agent.run + original_run_sync = Agent.run_sync + original_run_stream = Agent.run_stream + original_run_stream_events = Agent.run_stream_events + + # Wrap and apply patches for non-streaming methods + Agent.run = _create_run_wrapper(original_run, is_streaming=False) # type: ignore + Agent.run_sync = _create_run_sync_wrapper(original_run_sync) # type: ignore + + # Wrap and apply patches for streaming methods + Agent.run_stream = _create_streaming_wrapper(original_run_stream) # type: ignore + Agent.run_stream_events = _create_streaming_events_wrapper( # type: ignore[method-assign] + original_run_stream_events + ) diff --git a/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py b/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py new file mode 100644 index 0000000000..21030013ef --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py @@ -0,0 +1,179 @@ +from functools import wraps + +from sentry_sdk.integrations import DidNotEnable + +from ..spans import ( + invoke_agent_span, + update_invoke_agent_span, + ai_client_span, + update_ai_client_span, +) + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Callable + +try: + import pydantic_ai + from pydantic_ai._agent_graph import UserPromptNode, ModelRequestNode, CallToolsNode +except ImportError: + raise DidNotEnable("pydantic-ai not installed") + + +def _patch_graph_nodes(): + # type: () -> None + """ + Patches the graph node execution to create appropriate spans. + + UserPromptNode -> Creates invoke_agent span + ModelRequestNode -> Creates ai_client span for model requests + CallToolsNode -> Handles tool calls (spans created in tool patching) + """ + + def _extract_span_data(node, ctx): + # type: (Any, Any) -> tuple[list[Any], Any, Any] + """Extract common data needed for creating chat spans. + + Returns: + Tuple of (messages, model, model_settings) + """ + # Extract model and settings from context + model = None + model_settings = None + if hasattr(ctx, "deps"): + model = getattr(ctx.deps, "model", None) + model_settings = getattr(ctx.deps, "model_settings", None) + + # Build full message list: history + current request + messages = [] + if hasattr(ctx, "state") and hasattr(ctx.state, "message_history"): + messages.extend(ctx.state.message_history) + + current_request = getattr(node, "request", None) + if current_request: + messages.append(current_request) + + return messages, model, model_settings + + # Patch UserPromptNode to create invoke_agent spans + original_user_prompt_run = UserPromptNode.run + + @wraps(original_user_prompt_run) + async def wrapped_user_prompt_run(self, ctx): + # type: (Any, Any) -> Any + # Extract data from context + user_prompt = getattr(self, "user_prompt", None) + model = None + model_settings = None + + if hasattr(ctx, "deps"): + model = getattr(ctx.deps, "model", None) + model_settings = getattr(ctx.deps, "model_settings", None) + + # Try to get agent from Sentry scope + import sentry_sdk + + agent_data = ( + sentry_sdk.get_current_scope()._contexts.get("pydantic_ai_agent") or {} + ) + agent = agent_data.get("_agent") + + # Create and store the invoke_agent span + span = invoke_agent_span(user_prompt, agent, model, model_settings) + # Store span in context for later use + if hasattr(ctx, "state"): + ctx.state._sentry_invoke_span = span + + result = await original_user_prompt_run(self, ctx) + return result + + UserPromptNode.run = wrapped_user_prompt_run # type: ignore[method-assign] + + # 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, ctx): + # type: (Any, Any) -> Any + messages, model, model_settings = _extract_span_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 = 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 + + # Patch ModelRequestNode.stream for streaming requests + original_model_request_stream = ModelRequestNode.stream + + def create_wrapped_stream(original_stream_method): + # type: (Callable[..., Any]) -> Callable[..., Any] + """Create a wrapper for ModelRequestNode.stream that creates chat spans.""" + from contextlib import asynccontextmanager + + @asynccontextmanager + @wraps(original_stream_method) + async def wrapped_model_request_stream(self, ctx): + # type: (Any, Any) -> Any + messages, model, model_settings = _extract_span_data(self, ctx) + + # Create chat span for streaming request + import sentry_sdk + + span = ai_client_span(messages, None, model, model_settings) + span.__enter__() + + try: + # 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 = 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) + finally: + # Close the span after streaming completes + span.__exit__(None, None, None) + + return wrapped_model_request_stream + + ModelRequestNode.stream = create_wrapped_stream(original_model_request_stream) # type: ignore + + # Patch CallToolsNode to close invoke_agent span when done + original_call_tools_run = CallToolsNode.run + + @wraps(original_call_tools_run) + async def wrapped_call_tools_run(self, ctx): + # type: (Any, Any) -> Any + result = await original_call_tools_run(self, ctx) + + # Check if this is an End node (final result) + from pydantic_graph import End + + if isinstance(result, End): + # Close the invoke_agent span if it exists + if hasattr(ctx, "state") and hasattr(ctx.state, "_sentry_invoke_span"): + span = ctx.state._sentry_invoke_span + output = None + if hasattr(result, "data") and hasattr(result.data, "output"): + output = result.data.output + update_invoke_agent_span(span, output) + delattr(ctx.state, "_sentry_invoke_span") + + return result + + CallToolsNode.run = wrapped_call_tools_run # type: ignore diff --git a/sentry_sdk/integrations/pydantic_ai/patches/model_request.py b/sentry_sdk/integrations/pydantic_ai/patches/model_request.py new file mode 100644 index 0000000000..25d09843f3 --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/patches/model_request.py @@ -0,0 +1,41 @@ +from functools import wraps + +from sentry_sdk.integrations import DidNotEnable + +from ..spans import ai_client_span, update_ai_client_span + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Callable + +try: + import pydantic_ai +except ImportError: + raise DidNotEnable("pydantic-ai not installed") + + +def _patch_model_request(): + # type: () -> None + """ + Patches model request execution to create AI client spans. + + In pydantic-ai, model requests are handled through the Model interface. + We need to patch the request method on models to create spans. + """ + from pydantic_ai import models + + # Patch the base Model class's request method + if hasattr(models, "Model"): + original_request = models.Model.request + + @wraps(original_request) + async def wrapped_request(self, messages, *args, **kwargs): + # type: (Any, Any, *Any, **Any) -> Any + # Pass all messages (full conversation history) + with ai_client_span(messages, None, self, None) as span: + result = await original_request(self, messages, *args, **kwargs) + update_ai_client_span(span, result) + return result + + models.Model.request = wrapped_request # type: ignore diff --git a/sentry_sdk/integrations/pydantic_ai/patches/tools.py b/sentry_sdk/integrations/pydantic_ai/patches/tools.py new file mode 100644 index 0000000000..f9f457c20e --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/patches/tools.py @@ -0,0 +1,67 @@ +from functools import wraps + +from sentry_sdk.integrations import DidNotEnable + +from ..spans import execute_tool_span, update_execute_tool_span + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Callable + +try: + import pydantic_ai +except ImportError: + raise DidNotEnable("pydantic-ai not installed") + + +def _patch_tool_execution(): + # type: () -> None + """ + Patches tool execution to create execute_tool spans. + + In pydantic-ai, tools are managed through the ToolManager. + We patch the toolset.call_tool method to create spans around tool execution. + + Note: pydantic-ai has built-in OpenTelemetry instrumentation for tools. + Our patching adds Sentry-specific span data on top of that. + """ + from pydantic_ai.toolsets.abstract import AbstractToolset + from pydantic_ai.toolsets.function import FunctionToolset + + def create_wrapped_call_tool(original_call_tool): + # type: (Callable[..., Any]) -> Callable[..., Any] + """Create a wrapped call_tool method.""" + + @wraps(original_call_tool) + async def wrapped_call_tool(self, name, args_dict, ctx, tool): + # type: (Any, str, Any, Any, Any) -> Any + # Always create span if we're in a Sentry transaction context + import sentry_sdk + + current_span = sentry_sdk.get_current_span() + should_create_span = current_span is not None + + if should_create_span: + # Get agent from Sentry scope + agent_data = ( + sentry_sdk.get_current_scope()._contexts.get("pydantic_ai_agent") + or {} + ) + agent = agent_data.get("_agent") + + with execute_tool_span(name, args_dict, agent) as span: + result = await original_call_tool(self, name, args_dict, ctx, tool) + update_execute_tool_span(span, result) + return result + else: + result = await original_call_tool(self, name, args_dict, ctx, tool) + return result + + return wrapped_call_tool + + # Patch AbstractToolset's call_tool method + AbstractToolset.call_tool = create_wrapped_call_tool(AbstractToolset.call_tool) # type: ignore + + # Also patch FunctionToolset specifically since it overrides call_tool + FunctionToolset.call_tool = create_wrapped_call_tool(FunctionToolset.call_tool) # type: ignore diff --git a/sentry_sdk/integrations/pydantic_ai/spans/__init__.py b/sentry_sdk/integrations/pydantic_ai/spans/__init__.py new file mode 100644 index 0000000000..4cc2373191 --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/spans/__init__.py @@ -0,0 +1,4 @@ +from .agent_workflow import agent_workflow_span # noqa: F401 +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/agent_workflow.py b/sentry_sdk/integrations/pydantic_ai/spans/agent_workflow.py new file mode 100644 index 0000000000..e73632fef9 --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/spans/agent_workflow.py @@ -0,0 +1,28 @@ +import sentry_sdk +from sentry_sdk.ai.utils import get_start_span_function +from sentry_sdk.consts import OP + +from ..consts import SPAN_ORIGIN + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + + +def agent_workflow_span(agent): + # type: (Any) -> sentry_sdk.tracing.Span + """Create a root span for the entire agent workflow.""" + start_span_function = get_start_span_function() + + agent_name = "agent" + if agent and hasattr(agent, "name") and agent.name: + agent_name = agent.name + + span = start_span_function( + op=OP.GEN_AI_PIPELINE, + name=f"agent workflow {agent_name}", + origin=SPAN_ORIGIN, + ) + + return span diff --git a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py new file mode 100644 index 0000000000..8d3ca195a1 --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py @@ -0,0 +1,110 @@ +import sentry_sdk +from sentry_sdk.consts import OP, SPANDATA + +from ..consts import SPAN_ORIGIN +from ..utils import ( + _get_model_name, + _set_agent_data, + _set_model_data, + _set_usage_data, + _set_input_messages, + _set_output_data, +) + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + + +def ai_client_span(messages, agent, model, model_settings): + # type: (Any, Any, Any, Any) -> sentry_sdk.tracing.Span + """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 + """ + # Determine model name for span name + model_obj = model + if agent and hasattr(agent, "model"): + model_obj = agent.model + + model_name = _get_model_name(model_obj) or "unknown" + + 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_agent_data(span, agent) + _set_model_data(span, model, model_settings) + + # Set streaming flag + agent_data = sentry_sdk.get_current_scope()._contexts.get("pydantic_ai_agent") or {} + is_streaming = agent_data.get("_streaming", False) + span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, is_streaming) + + # Add available tools if agent is available + agent_obj = agent + if not agent_obj: + # Try to get from Sentry scope + agent_data = ( + sentry_sdk.get_current_scope()._contexts.get("pydantic_ai_agent") or {} + ) + agent_obj = agent_data.get("_agent") + + if agent_obj and hasattr(agent_obj, "_function_toolset"): + try: + from sentry_sdk.utils import safe_serialize + + tools = [] + # Get tools from the function toolset + if hasattr(agent_obj._function_toolset, "tools"): + for tool_name, tool in agent_obj._function_toolset.tools.items(): + tool_info = {"name": tool_name} + + # Add description from function_schema if available + if hasattr(tool, "function_schema"): + schema = tool.function_schema + if hasattr(schema, "description") and schema.description: + tool_info["description"] = schema.description + + # Add parameters from json_schema + if hasattr(schema, "json_schema") and schema.json_schema: + tool_info["parameters"] = schema.json_schema + + tools.append(tool_info) + + if tools: + span.set_data( + SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools) + ) + except Exception: + # If we can't extract tools, just skip it + pass + + # Set input messages (full conversation history) + if messages: + _set_input_messages(span, messages) + + return span + + +def update_ai_client_span(span, model_response): + # type: (sentry_sdk.tracing.Span, Any) -> 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 new file mode 100644 index 0000000000..5746ce9f13 --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py @@ -0,0 +1,42 @@ +import sentry_sdk +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.utils import safe_serialize + +from ..consts import SPAN_ORIGIN +from ..utils import _set_agent_data, _should_send_prompts + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + + +def execute_tool_span(tool_name, tool_args, agent): + # type: (str, Any, Any) -> sentry_sdk.tracing.Span + """Create a span for tool execution.""" + 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_TYPE, "function") + span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool_name) + + _set_agent_data(span, agent) + + if _should_send_prompts() and tool_args is not None: + span.set_data(SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_args)) + + return span + + +def update_execute_tool_span(span, result): + # type: (sentry_sdk.tracing.Span, Any) -> None + """Update the execute tool span with the result.""" + if not span: + return + + if _should_send_prompts() and result is not None: + span.set_data(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 new file mode 100644 index 0000000000..2f415280d3 --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py @@ -0,0 +1,144 @@ +import sentry_sdk +from sentry_sdk.ai.utils import get_start_span_function, set_data_normalized +from sentry_sdk.consts import OP, SPANDATA + +from ..consts import SPAN_ORIGIN +from ..utils import _set_agent_data, _set_model_data, _should_send_prompts + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + + +def invoke_agent_span(user_prompt, agent, model, model_settings): + # type: (Any, Any, Any, Any) -> sentry_sdk.tracing.Span + """Create a span for invoking the agent.""" + start_span_function = get_start_span_function() + + # Determine agent name for span + name = "agent" + if agent and hasattr(agent, "name") and agent.name: + name = agent.name + + span = start_span_function( + op=OP.GEN_AI_INVOKE_AGENT, + name=f"invoke_agent {name}", + origin=SPAN_ORIGIN, + ) + span.__enter__() + + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") + + _set_agent_data(span, agent) + _set_model_data(span, model, model_settings) + + # Add available tools if present + if agent and hasattr(agent, "_function_toolset"): + try: + from sentry_sdk.utils import safe_serialize + + tools = [] + # Get tools from the function toolset + if hasattr(agent._function_toolset, "tools"): + for tool_name, tool in agent._function_toolset.tools.items(): + tool_info = {"name": tool_name} + + # Add description from function_schema if available + if hasattr(tool, "function_schema"): + schema = tool.function_schema + if hasattr(schema, "description") and schema.description: + tool_info["description"] = schema.description + + # Add parameters from json_schema + if hasattr(schema, "json_schema") and schema.json_schema: + tool_info["parameters"] = schema.json_schema + + tools.append(tool_info) + + if tools: + span.set_data( + SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools) + ) + except Exception: + # If we can't extract tools, just skip it + pass + + # Add user prompt and system prompts if available and prompts are enabled + if _should_send_prompts(): + messages = [] + + # Add system prompts (both instructions and system_prompt) + system_texts = [] + + if agent: + # Check for system_prompt + if hasattr(agent, "_system_prompts") and agent._system_prompts: + for prompt in agent._system_prompts: + if isinstance(prompt, str): + system_texts.append(prompt) + + # Check for instructions (stored in _instructions) + if hasattr(agent, "_instructions") and agent._instructions: + instructions = agent._instructions + if isinstance(instructions, str): + system_texts.append(instructions) + elif isinstance(instructions, (list, tuple)): + for instr in instructions: + if isinstance(instr, str): + system_texts.append(instr) + elif callable(instr): + # Skip dynamic/callable instructions + pass + + # Add all system texts as system messages + for system_text in system_texts: + messages.append( + { + "content": [{"text": system_text, "type": "text"}], + "role": "system", + } + ) + + # Add user prompt + if user_prompt: + if isinstance(user_prompt, str): + messages.append( + { + "content": [{"text": user_prompt, "type": "text"}], + "role": "user", + } + ) + elif isinstance(user_prompt, list): + # Handle list of user content + content = [] + for item in user_prompt: + if isinstance(item, str): + content.append({"text": item, "type": "text"}) + if content: + messages.append( + { + "content": content, + "role": "user", + } + ) + + if messages: + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages, unpack=False + ) + + return span + + +def update_invoke_agent_span(span, output): + # type: (sentry_sdk.tracing.Span, Any) -> None + """Update and close the invoke agent span.""" + if span and _should_send_prompts() and output: + output_text = str(output) if not isinstance(output, str) else output + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TEXT, output_text, unpack=False + ) + + if span: + span.__exit__(None, None, None) diff --git a/sentry_sdk/integrations/pydantic_ai/utils.py b/sentry_sdk/integrations/pydantic_ai/utils.py new file mode 100644 index 0000000000..f92b148f34 --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/utils.py @@ -0,0 +1,308 @@ +import sentry_sdk +from sentry_sdk.ai.utils import set_data_normalized +from sentry_sdk.consts import SPANDATA +from sentry_sdk.integrations import DidNotEnable +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing_utils import set_span_errored +from sentry_sdk.utils import event_from_exception, safe_serialize + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, List, Dict + from pydantic_ai.usage import RequestUsage + +try: + import pydantic_ai + +except ImportError: + raise DidNotEnable("pydantic-ai not installed") + + +def _should_send_prompts(): + # type: () -> 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 + + # Get the integration instance from the client + integration = sentry_sdk.get_client().get_integration( + pydantic_ai.__name__.split(".")[0] + ) + + if integration is None: + return False + + return getattr(integration, "include_prompts", True) + + +def _capture_exception(exc): + # type: (Any) -> None + set_span_errored() + + event, hint = event_from_exception( + exc, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "pydantic_ai", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + +def _get_model_name(model_obj): + # type: (Any) -> str | None + """Extract model name from a model object. + + Args: + model_obj: Model object to extract name from + + Returns: + Model name string or None if not found + """ + if not model_obj: + return None + + if hasattr(model_obj, "model_name"): + return model_obj.model_name + elif hasattr(model_obj, "name"): + try: + return model_obj.name() + except Exception: + return str(model_obj) + elif isinstance(model_obj, str): + return model_obj + else: + return str(model_obj) + + +def _set_agent_data(span, agent): + # type: (sentry_sdk.tracing.Span, 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 Sentry scope if not provided) + """ + # Extract agent name from agent object or Sentry scope + agent_obj = agent + if not agent_obj: + # Try to get from Sentry scope + agent_data = ( + sentry_sdk.get_current_scope()._contexts.get("pydantic_ai_agent") or {} + ) + agent_obj = agent_data.get("_agent") + + if agent_obj and hasattr(agent_obj, "name") and agent_obj.name: + span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_obj.name) + + +def _set_model_data(span, model, model_settings): + # type: (sentry_sdk.tracing.Span, Any, Any) -> 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) + """ + # Try to get agent from Sentry scope if we need it + agent_data = sentry_sdk.get_current_scope()._contexts.get("pydantic_ai_agent") or {} + agent_obj = agent_data.get("_agent") + + # Extract model information + model_obj = model + if not model_obj and agent_obj and hasattr(agent_obj, "model"): + model_obj = agent_obj.model + + if model_obj: + # Set system from model + if hasattr(model_obj, "system"): + span.set_data(SPANDATA.GEN_AI_SYSTEM, model_obj.system) + + # Set model name + model_name = _get_model_name(model_obj) + if model_name: + span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) + + # Extract model settings + settings = model_settings + if not settings and agent_obj and hasattr(agent_obj, "model_settings"): + settings = agent_obj.model_settings + + if settings: + settings_map = { + "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, + "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE, + "top_p": SPANDATA.GEN_AI_REQUEST_TOP_P, + "frequency_penalty": SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY, + "presence_penalty": SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY, + } + + # ModelSettings is a TypedDict (dict at runtime), so use dict access + if isinstance(settings, dict): + for setting_name, spandata_key in settings_map.items(): + value = settings.get(setting_name) + if value is not None: + span.set_data(spandata_key, value) + else: + # Fallback for object-style settings + for setting_name, spandata_key in settings_map.items(): + if hasattr(settings, setting_name): + value = getattr(settings, setting_name) + if value is not None: + span.set_data(spandata_key, value) + + +def _set_usage_data(span, usage): + # type: (sentry_sdk.tracing.Span, RequestUsage) -> None + """Set token usage data on a span.""" + if usage is None: + return + + if hasattr(usage, "input_tokens") and usage.input_tokens is not None: + span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens) + + if hasattr(usage, "output_tokens") and usage.output_tokens is not None: + span.set_data(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage.output_tokens) + + if hasattr(usage, "total_tokens") and usage.total_tokens is not None: + span.set_data(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens) + + +def _set_input_messages(span, messages): + # type: (sentry_sdk.tracing.Span, Any) -> None + """Set input messages data on a span.""" + if not _should_send_prompts(): + return + + if not messages: + return + + try: + formatted_messages = [] + system_prompt = None + + # Extract system prompt from any ModelRequest with instructions + for msg in messages: + if hasattr(msg, "instructions") and msg.instructions: + system_prompt = msg.instructions + break + + # Add system prompt as first message if present + if system_prompt: + formatted_messages.append( + {"role": "system", "content": [{"type": "text", "text": system_prompt}]} + ) + + for msg in messages: + if hasattr(msg, "parts"): + for part in msg.parts: + role = "user" + if hasattr(part, "__class__"): + if "System" in part.__class__.__name__: + role = "system" + elif ( + "Assistant" in part.__class__.__name__ + or "Text" in part.__class__.__name__ + or "ToolCall" in part.__class__.__name__ + ): + role = "assistant" + elif "ToolReturn" in part.__class__.__name__: + role = "tool" + + content = [] # type: List[Dict[str, Any] | str] + tool_calls = None + tool_call_id = None + + # Handle ToolCallPart (assistant requesting tool use) + if "ToolCall" in part.__class__.__name__: + tool_call_data = {} + if hasattr(part, "tool_name"): + tool_call_data["name"] = part.tool_name + if hasattr(part, "args"): + tool_call_data["arguments"] = safe_serialize(part.args) + if tool_call_data: + tool_calls = [tool_call_data] + # Handle ToolReturnPart (tool result) + elif "ToolReturn" in part.__class__.__name__: + if hasattr(part, "tool_name"): + tool_call_id = part.tool_name + if hasattr(part, "content"): + content.append({"type": "text", "text": str(part.content)}) + # Handle regular content + elif hasattr(part, "content"): + if isinstance(part.content, str): + content.append({"type": "text", "text": part.content}) + elif isinstance(part.content, list): + for item in part.content: + if isinstance(item, str): + content.append({"type": "text", "text": item}) + else: + content.append(safe_serialize(item)) + else: + content.append({"type": "text", "text": str(part.content)}) + + # Add message if we have content or tool calls + if content or tool_calls: + message = {"role": role} # type: Dict[str, Any] + if content: + message["content"] = content + if tool_calls: + message["tool_calls"] = tool_calls + if tool_call_id: + message["tool_call_id"] = tool_call_id + formatted_messages.append(message) + + if formatted_messages: + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_MESSAGES, formatted_messages, unpack=False + ) + except Exception: + # If we fail to format messages, just skip it + pass + + +def _set_output_data(span, response): + # type: (sentry_sdk.tracing.Span, Any) -> None + """Set output data on a span.""" + if not _should_send_prompts(): + return + + if not response: + return + + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name) + try: + # Extract text from ModelResponse + if hasattr(response, "parts"): + texts = [] + tool_calls = [] + + for part in response.parts: + if hasattr(part, "__class__"): + if "Text" in part.__class__.__name__ and hasattr(part, "content"): + texts.append(part.content) + elif "ToolCall" in part.__class__.__name__: + tool_call_data = { + "type": "function", + } + if hasattr(part, "tool_name"): + tool_call_data["name"] = part.tool_name + if hasattr(part, "args"): + tool_call_data["arguments"] = safe_serialize(part.args) + tool_calls.append(tool_call_data) + + if texts: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, texts) + + if tool_calls: + span.set_data( + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls) + ) + + except Exception: + # If we fail to format output, just skip it + pass diff --git a/setup.py b/setup.py index c6e391d27a..29f85f6d0a 100644 --- a/setup.py +++ b/setup.py @@ -73,6 +73,7 @@ def get_file_text(file_name): "opentelemetry": ["opentelemetry-distro>=0.35b0"], "opentelemetry-experimental": ["opentelemetry-distro"], "pure-eval": ["pure_eval", "executing", "asttokens"], + "pydantic_ai": ["pydantic-ai>=1.0.0"], "pymongo": ["pymongo>=3.1"], "pyspark": ["pyspark>=2.4.4"], "quart": ["quart>=0.16.1", "blinker>=1.1"], diff --git a/tests/integrations/pydantic_ai/__init__.py b/tests/integrations/pydantic_ai/__init__.py new file mode 100644 index 0000000000..3a2ad11c0c --- /dev/null +++ b/tests/integrations/pydantic_ai/__init__.py @@ -0,0 +1,3 @@ +import pytest + +pytest.importorskip("pydantic_ai") diff --git a/tests/integrations/pydantic_ai/test_pydantic_ai.py b/tests/integrations/pydantic_ai/test_pydantic_ai.py new file mode 100644 index 0000000000..7460f28df1 --- /dev/null +++ b/tests/integrations/pydantic_ai/test_pydantic_ai.py @@ -0,0 +1,677 @@ +import asyncio +import pytest + +from sentry_sdk.integrations.pydantic_ai import PydanticAIIntegration + +from pydantic_ai import Agent +from pydantic_ai.models.test import TestModel + + +@pytest.fixture +def test_agent(): + """Create a test agent with model settings.""" + return Agent( + "test", + name="test_agent", + system_prompt="You are a helpful test assistant.", + ) + + +@pytest.fixture +def test_agent_with_settings(): + """Create a test agent with explicit model settings.""" + from pydantic_ai import ModelSettings + + return Agent( + "test", + name="test_agent_settings", + system_prompt="You are a test assistant with settings.", + model_settings=ModelSettings( + temperature=0.7, + max_tokens=100, + top_p=0.9, + ), + ) + + +@pytest.mark.asyncio +async def test_agent_run_async(sentry_init, capture_events, test_agent): + """ + Test that the integration creates spans for async agent runs. + """ + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + events = capture_events() + + result = await test_agent.run("Test input") + + assert result is not None + assert result.output is not None + + (transaction,) = events + spans = transaction["spans"] + + # Verify transaction + assert transaction["transaction"] == "agent workflow test_agent" + assert transaction["contexts"]["trace"]["origin"] == "auto.ai.pydantic_ai" + + # Find span types + invoke_agent_spans = [s for s in spans if s["op"] == "gen_ai.invoke_agent"] + chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + + assert len(invoke_agent_spans) == 1 + assert len(chat_spans) >= 1 + + # Check invoke_agent span + invoke_agent_span = invoke_agent_spans[0] + assert invoke_agent_span["description"] == "invoke_agent test_agent" + assert invoke_agent_span["data"]["gen_ai.operation.name"] == "invoke_agent" + assert invoke_agent_span["data"]["gen_ai.agent.name"] == "test_agent" + assert "gen_ai.request.messages" in invoke_agent_span["data"] + assert "gen_ai.response.text" in invoke_agent_span["data"] + + # Check chat span + chat_span = chat_spans[0] + assert "chat" in chat_span["description"] + assert chat_span["data"]["gen_ai.operation.name"] == "chat" + assert chat_span["data"]["gen_ai.response.streaming"] is False + assert "gen_ai.request.messages" in chat_span["data"] + assert "gen_ai.usage.input_tokens" in chat_span["data"] + assert "gen_ai.usage.output_tokens" in chat_span["data"] + + +def test_agent_run_sync(sentry_init, capture_events, test_agent): + """ + Test that the integration creates spans for sync agent runs. + """ + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + events = capture_events() + + result = test_agent.run_sync("Test input") + + assert result is not None + assert result.output is not None + + (transaction,) = events + spans = transaction["spans"] + + # Verify transaction + assert transaction["transaction"] == "agent workflow test_agent" + assert transaction["contexts"]["trace"]["origin"] == "auto.ai.pydantic_ai" + + # Find span types + chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + assert len(chat_spans) >= 1 + + # Verify streaming flag is False for sync + for chat_span in chat_spans: + assert chat_span["data"]["gen_ai.response.streaming"] is False + + +@pytest.mark.asyncio +async def test_agent_run_stream(sentry_init, capture_events, test_agent): + """ + Test that the integration creates spans for streaming agent runs. + """ + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + events = capture_events() + + async with test_agent.run_stream("Test input") as result: + # Consume the stream + async for _ in result.stream_output(): + pass + + (transaction,) = events + spans = transaction["spans"] + + # Verify transaction + assert transaction["transaction"] == "agent workflow test_agent" + assert transaction["contexts"]["trace"]["origin"] == "auto.ai.pydantic_ai" + + # Find chat spans + chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + assert len(chat_spans) >= 1 + + # Verify streaming flag is True for streaming + for chat_span in chat_spans: + assert chat_span["data"]["gen_ai.response.streaming"] is True + assert "gen_ai.request.messages" in chat_span["data"] + assert "gen_ai.usage.input_tokens" in chat_span["data"] + # Streaming responses should still have output data + assert ( + "gen_ai.response.text" in chat_span["data"] + or "gen_ai.response.model" in chat_span["data"] + ) + + +@pytest.mark.asyncio +async def test_agent_run_stream_events(sentry_init, capture_events, test_agent): + """ + Test that run_stream_events creates spans (it uses run internally, so non-streaming). + """ + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + events = capture_events() + + # Consume all events + async for _ in test_agent.run_stream_events("Test input"): + pass + + (transaction,) = events + + # Verify transaction + assert transaction["transaction"] == "agent workflow test_agent" + + # Find chat spans + spans = transaction["spans"] + chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + assert len(chat_spans) >= 1 + + # run_stream_events uses run() internally, so streaming should be False + for chat_span in chat_spans: + assert chat_span["data"]["gen_ai.response.streaming"] is False + + +@pytest.mark.asyncio +async def test_agent_with_tools(sentry_init, capture_events, test_agent): + """ + Test that tool execution creates execute_tool spans. + """ + + @test_agent.tool_plain + def add_numbers(a: int, b: int) -> int: + """Add two numbers together.""" + return a + b + + sentry_init( + integrations=[PydanticAIIntegration()], + 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 + + (transaction,) = events + spans = transaction["spans"] + + # Find span types + invoke_agent_spans = [s for s in spans if s["op"] == "gen_ai.invoke_agent"] + 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 span + tool_span = tool_spans[0] + 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 invoke_agent span has available_tools + invoke_agent_span = invoke_agent_spans[0] + assert "gen_ai.request.available_tools" in invoke_agent_span["data"] + available_tools_str = invoke_agent_span["data"]["gen_ai.request.available_tools"] + # Available tools is serialized as a string + assert "add_numbers" in available_tools_str + + # Check chat spans also have available_tools + for chat_span in chat_spans: + assert "gen_ai.request.available_tools" in chat_span["data"] + + +@pytest.mark.asyncio +async def test_agent_with_tools_streaming(sentry_init, capture_events, test_agent): + """ + Test that tool execution works correctly with streaming. + """ + + @test_agent.tool_plain + def multiply(a: int, b: int) -> int: + """Multiply two numbers.""" + return a * b + + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + events = capture_events() + + async with test_agent.run_stream("What is 7 times 8?") as result: + async for _ in result.stream_output(): + pass + + (transaction,) = events + spans = transaction["spans"] + + # Find span types + 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 + + # Verify streaming flag is True + for chat_span in chat_spans: + assert chat_span["data"]["gen_ai.response.streaming"] is True + + # Check tool span + tool_span = tool_spans[0] + assert tool_span["data"]["gen_ai.tool.name"] == "multiply" + assert "gen_ai.tool.input" in tool_span["data"] + assert "gen_ai.tool.output" in tool_span["data"] + + +@pytest.mark.asyncio +async def test_model_settings(sentry_init, capture_events, test_agent_with_settings): + """ + Test that model settings are captured in spans. + """ + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + ) + + events = capture_events() + + await test_agent_with_settings.run("Test input") + + (transaction,) = events + spans = transaction["spans"] + + # Find chat span + chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + assert len(chat_spans) >= 1 + + chat_span = chat_spans[0] + # Check that model settings are captured + assert chat_span["data"].get("gen_ai.request.temperature") == 0.7 + assert chat_span["data"].get("gen_ai.request.max_tokens") == 100 + assert chat_span["data"].get("gen_ai.request.top_p") == 0.9 + + +@pytest.mark.asyncio +async def test_system_prompt_in_messages(sentry_init, capture_events): + """ + Test that system prompts are included as the first message. + """ + agent = Agent( + "test", + name="test_system", + system_prompt="You are a helpful assistant specialized in testing.", + ) + + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + events = capture_events() + + await agent.run("Hello") + + (transaction,) = events + spans = transaction["spans"] + + # Find invoke_agent span + invoke_agent_spans = [s for s in spans if s["op"] == "gen_ai.invoke_agent"] + assert len(invoke_agent_spans) == 1 + + invoke_agent_span = invoke_agent_spans[0] + messages_str = invoke_agent_span["data"]["gen_ai.request.messages"] + + # Messages is serialized as a string + # Should contain system role and helpful assistant text + assert "system" in messages_str + assert "helpful assistant" in messages_str + + +@pytest.mark.asyncio +async def test_error_handling(sentry_init, capture_events): + """ + Test error handling in agent execution. + """ + # Use a simpler test that doesn't cause tool failures + # as pydantic-ai has complex error handling for tool errors + agent = Agent( + "test", + name="test_error", + ) + + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + ) + + events = capture_events() + + # Simple run that should succeed + await agent.run("Hello") + + # At minimum, we should have a transaction + assert len(events) >= 1 + transaction = [e for e in events if e.get("type") == "transaction"][0] + assert transaction["transaction"] == "agent workflow test_error" + # Transaction should complete successfully (status key may not exist if no error) + trace_status = transaction["contexts"]["trace"].get("status") + assert trace_status != "error" # Could be None or some other status + + +@pytest.mark.asyncio +async def test_without_pii(sentry_init, capture_events, test_agent): + """ + Test that PII is not captured when send_default_pii is False. + """ + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + send_default_pii=False, + ) + + events = capture_events() + + await test_agent.run("Sensitive input") + + (transaction,) = events + spans = transaction["spans"] + + # Find spans + invoke_agent_spans = [s for s in spans if s["op"] == "gen_ai.invoke_agent"] + chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + + # Verify that messages and response text are not captured + for span in invoke_agent_spans + chat_spans: + assert "gen_ai.request.messages" not in span["data"] + assert "gen_ai.response.text" not in span["data"] + + +@pytest.mark.asyncio +async def test_without_pii_tools(sentry_init, capture_events, test_agent): + """ + Test that tool input/output are not captured when send_default_pii is False. + """ + + @test_agent.tool_plain + def sensitive_tool(data: str) -> str: + """A tool with sensitive data.""" + return f"Processed: {data}" + + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + send_default_pii=False, + ) + + events = capture_events() + + await test_agent.run("Use sensitive tool with private data") + + (transaction,) = events + spans = transaction["spans"] + + # Find tool spans + tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] + + # If tool was executed, verify input/output are not captured + for tool_span in tool_spans: + assert "gen_ai.tool.input" not in tool_span["data"] + assert "gen_ai.tool.output" not in tool_span["data"] + + +@pytest.mark.asyncio +async def test_multiple_agents_concurrent(sentry_init, capture_events, test_agent): + """ + Test that multiple agents can run concurrently without interfering. + """ + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + ) + + events = capture_events() + + async def run_agent(input_text): + return await test_agent.run(input_text) + + # Run 3 agents concurrently + results = await asyncio.gather(*[run_agent(f"Input {i}") for i in range(3)]) + + assert len(results) == 3 + assert len(events) == 3 + + # Verify each transaction is separate + for i, transaction in enumerate(events): + assert transaction["type"] == "transaction" + assert transaction["transaction"] == "agent workflow test_agent" + # Each should have its own spans + assert len(transaction["spans"]) >= 1 + + +@pytest.mark.asyncio +async def test_message_history(sentry_init, capture_events): + """ + Test that full conversation history is captured in chat spans. + """ + agent = Agent( + "test", + name="test_history", + ) + + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + events = capture_events() + + # First message + await agent.run("Hello, I'm Alice") + + # Second message with history + from pydantic_ai import messages + + history = [ + messages.UserPromptPart(content="Hello, I'm Alice"), + messages.ModelResponse( + parts=[messages.TextPart(content="Hello Alice! How can I help you?")], + model_name="test", + ), + ] + + await agent.run("What is my name?", message_history=history) + + # We should have 2 transactions + assert len(events) >= 2 + + # Check the second transaction has the full history + second_transaction = events[1] + spans = second_transaction["spans"] + chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + + if chat_spans: + chat_span = chat_spans[0] + if "gen_ai.request.messages" in chat_span["data"]: + messages_data = chat_span["data"]["gen_ai.request.messages"] + # Should have multiple messages including history + assert len(messages_data) > 1 + + +@pytest.mark.asyncio +async def test_gen_ai_system(sentry_init, capture_events, test_agent): + """ + Test that gen_ai.system is set from the model. + """ + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + ) + + events = capture_events() + + await test_agent.run("Test input") + + (transaction,) = events + spans = transaction["spans"] + + # Find chat span + chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + assert len(chat_spans) >= 1 + + chat_span = chat_spans[0] + # gen_ai.system should be set from the model (TestModel -> 'test') + assert "gen_ai.system" in chat_span["data"] + assert chat_span["data"]["gen_ai.system"] == "test" + + +@pytest.mark.asyncio +async def test_include_prompts_false(sentry_init, capture_events, test_agent): + """ + Test that prompts are not captured when include_prompts=False. + """ + sentry_init( + integrations=[PydanticAIIntegration(include_prompts=False)], + traces_sample_rate=1.0, + send_default_pii=True, # Even with PII enabled, prompts should not be captured + ) + + events = capture_events() + + await test_agent.run("Sensitive prompt") + + (transaction,) = events + spans = transaction["spans"] + + # Find spans + invoke_agent_spans = [s for s in spans if s["op"] == "gen_ai.invoke_agent"] + chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + + # Verify that messages and response text are not captured + for span in invoke_agent_spans + chat_spans: + assert "gen_ai.request.messages" not in span["data"] + assert "gen_ai.response.text" not in span["data"] + + +@pytest.mark.asyncio +async def test_include_prompts_true(sentry_init, capture_events, test_agent): + """ + Test that prompts are captured when include_prompts=True (default). + """ + sentry_init( + integrations=[PydanticAIIntegration(include_prompts=True)], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + events = capture_events() + + await test_agent.run("Test prompt") + + (transaction,) = events + spans = transaction["spans"] + + # Find spans + invoke_agent_spans = [s for s in spans if s["op"] == "gen_ai.invoke_agent"] + chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + + # Verify that messages are captured + assert len(invoke_agent_spans) >= 1 + invoke_agent_span = invoke_agent_spans[0] + assert "gen_ai.request.messages" in invoke_agent_span["data"] + assert "gen_ai.response.text" in invoke_agent_span["data"] + + # Chat spans should also have messages + assert len(chat_spans) >= 1 + for chat_span in chat_spans: + assert "gen_ai.request.messages" in chat_span["data"] + + +@pytest.mark.asyncio +async def test_include_prompts_false_with_tools( + sentry_init, capture_events, test_agent +): + """ + Test that tool input/output are not captured when include_prompts=False. + """ + + @test_agent.tool_plain + def test_tool(value: int) -> int: + """A test tool.""" + return value * 2 + + sentry_init( + integrations=[PydanticAIIntegration(include_prompts=False)], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + events = capture_events() + + await test_agent.run("Use the test tool with value 5") + + (transaction,) = events + spans = transaction["spans"] + + # Find tool spans + tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"] + + # If tool was executed, verify input/output are not captured + for tool_span in tool_spans: + assert "gen_ai.tool.input" not in tool_span["data"] + assert "gen_ai.tool.output" not in tool_span["data"] + + +@pytest.mark.asyncio +async def test_include_prompts_requires_pii(sentry_init, capture_events, test_agent): + """ + Test that include_prompts requires send_default_pii=True. + """ + sentry_init( + integrations=[PydanticAIIntegration(include_prompts=True)], + traces_sample_rate=1.0, + send_default_pii=False, # PII disabled + ) + + events = capture_events() + + await test_agent.run("Test prompt") + + (transaction,) = events + spans = transaction["spans"] + + # Find spans + invoke_agent_spans = [s for s in spans if s["op"] == "gen_ai.invoke_agent"] + chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"] + + # Even with include_prompts=True, if PII is disabled, messages should not be captured + for span in invoke_agent_spans + chat_spans: + assert "gen_ai.request.messages" not in span["data"] + assert "gen_ai.response.text" not in span["data"] diff --git a/tox.ini b/tox.ini index 5fb05f01bc..7eda8953d4 100644 --- a/tox.ini +++ b/tox.ini @@ -94,6 +94,8 @@ envlist = {py3.10,py3.12,py3.13}-openai_agents-v0.2.11 {py3.10,py3.12,py3.13}-openai_agents-v0.3.3 + {py3.10,py3.12,py3.13}-pydantic_ai-v1.0.17 + # ~~~ Cloud ~~~ {py3.6,py3.7}-boto3-v1.12.49 @@ -406,6 +408,9 @@ deps = openai_agents-v0.3.3: openai-agents==0.3.3 openai_agents: pytest-asyncio + pydantic_ai-v1.0.17: pydantic-ai==1.0.17 + pydantic_ai: pytest-asyncio + # ~~~ Cloud ~~~ boto3-v1.12.49: boto3==1.12.49 @@ -777,6 +782,7 @@ setenv = openai_agents: TESTPATH=tests/integrations/openai_agents openfeature: TESTPATH=tests/integrations/openfeature pure_eval: TESTPATH=tests/integrations/pure_eval + pydantic_ai: TESTPATH=tests/integrations/pydantic_ai pymongo: TESTPATH=tests/integrations/pymongo pyramid: TESTPATH=tests/integrations/pyramid quart: TESTPATH=tests/integrations/quart