sentry-sdk 3.0.0a2__py2.py3-none-any.whl → 3.0.0a4__py2.py3-none-any.whl

sentry_sdk/integrations/openai_agents/__init__.py

@@ -0,0 +1,49 @@
+from sentry_sdk.integrations import DidNotEnable, Integration
+
+from .patches import (
+    _create_get_model_wrapper,
+    _create_get_all_tools_wrapper,
+    _create_run_wrapper,
+    _patch_agent_run,
+)
+
+try:
+    import agents
+
+except ImportError:
+    raise DidNotEnable("OpenAI Agents not installed")
+
+
+def _patch_runner() -> None:
+    # Create the root span for one full agent run (including eventual handoffs)
+    # Note agents.run.DEFAULT_AGENT_RUNNER.run_sync is a wrapper around
+    # agents.run.DEFAULT_AGENT_RUNNER.run. It does not need to be wrapped separately.
+    # TODO-anton: Also patch streaming runner: agents.Runner.run_streamed
+    agents.run.DEFAULT_AGENT_RUNNER.run = _create_run_wrapper(
+        agents.run.DEFAULT_AGENT_RUNNER.run
+    )
+
+    # Creating the actual spans for each agent run.
+    _patch_agent_run()
+
+
+def _patch_model() -> None:
+    agents.run.AgentRunner._get_model = classmethod(
+        _create_get_model_wrapper(agents.run.AgentRunner._get_model),
+    )
+
+
+def _patch_tools() -> None:
+    agents.run.AgentRunner._get_all_tools = classmethod(
+        _create_get_all_tools_wrapper(agents.run.AgentRunner._get_all_tools),
+    )
+
+
+class OpenAIAgentsIntegration(Integration):
+    identifier = "openai_agents"
+
+    @staticmethod
+    def setup_once() -> None:
+        _patch_tools()
+        _patch_model()
+        _patch_runner()

sentry_sdk/integrations/openai_agents/consts.py

@@ -0,0 +1 @@
+SPAN_ORIGIN = "auto.ai.openai_agents"

sentry_sdk/integrations/openai_agents/patches/__init__.py

@@ -0,0 +1,4 @@
+from .models import _create_get_model_wrapper  # noqa: F401
+from .tools import _create_get_all_tools_wrapper  # noqa: F401
+from .runner import _create_run_wrapper  # noqa: F401
+from .agent_run import _patch_agent_run  # noqa: F401

sentry_sdk/integrations/openai_agents/patches/agent_run.py

@@ -0,0 +1,152 @@
+from __future__ import annotations
+
+from functools import wraps
+
+from sentry_sdk.integrations import DidNotEnable
+
+from ..spans import invoke_agent_span, update_invoke_agent_span, handoff_span
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, Optional
+
+
+try:
+    import agents
+except ImportError:
+    raise DidNotEnable("OpenAI Agents not installed")
+
+
+def _patch_agent_run() -> None:
+    """
+    Patches AgentRunner methods to create agent invocation spans.
+    This directly patches the execution flow to track when agents start and stop.
+    """
+
+    # Store original methods
+    original_run_single_turn = agents.run.AgentRunner._run_single_turn
+    original_execute_handoffs = agents._run_impl.RunImpl.execute_handoffs
+    original_execute_final_output = agents._run_impl.RunImpl.execute_final_output
+
+    def _start_invoke_agent_span(
+        context_wrapper: agents.RunContextWrapper, agent: agents.Agent
+    ) -> None:
+        """Start an agent invocation span"""
+        # Store the agent on the context wrapper so we can access it later
+        context_wrapper._sentry_current_agent = agent
+        invoke_agent_span(context_wrapper, agent)
+
+    def _end_invoke_agent_span(
+        context_wrapper: agents.RunContextWrapper,
+        agent: agents.Agent,
+        output: Optional[Any] = None,
+    ) -> None:
+        """End the agent invocation span"""
+        # Clear the stored agent
+        if hasattr(context_wrapper, "_sentry_current_agent"):
+            delattr(context_wrapper, "_sentry_current_agent")
+
+        update_invoke_agent_span(context_wrapper, agent, output)
+
+    def _has_active_agent_span(context_wrapper: agents.RunContextWrapper) -> bool:
+        """Check if there's an active agent span for this context"""
+        return getattr(context_wrapper, "_sentry_current_agent", None) is not None
+
+    def _get_current_agent(
+        context_wrapper: agents.RunContextWrapper,
+    ) -> Optional[agents.Agent]:
+        """Get the current agent from context wrapper"""
+        return getattr(context_wrapper, "_sentry_current_agent", None)
+
+    @wraps(
+        original_run_single_turn.__func__
+        if hasattr(original_run_single_turn, "__func__")
+        else original_run_single_turn
+    )
+    async def patched_run_single_turn(
+        cls: agents.Runner, *args: Any, **kwargs: Any
+    ) -> Any:
+        """Patched _run_single_turn that creates agent invocation spans"""
+        agent = kwargs.get("agent")
+        context_wrapper = kwargs.get("context_wrapper")
+        should_run_agent_start_hooks = kwargs.get("should_run_agent_start_hooks")
+
+        # Start agent span when agent starts (but only once per agent)
+        if should_run_agent_start_hooks and agent and context_wrapper:
+            # End any existing span for a different agent
+            if _has_active_agent_span(context_wrapper):
+                current_agent = _get_current_agent(context_wrapper)
+                if current_agent and current_agent != agent:
+                    _end_invoke_agent_span(context_wrapper, current_agent)
+
+            _start_invoke_agent_span(context_wrapper, agent)
+
+        # Call original method with all the correct parameters
+        try:
+            result = await original_run_single_turn(*args, **kwargs)
+        finally:
+            if agent and context_wrapper and _has_active_agent_span(context_wrapper):
+                _end_invoke_agent_span(context_wrapper, agent)
+
+        return result
+
+    @wraps(
+        original_execute_handoffs.__func__
+        if hasattr(original_execute_handoffs, "__func__")
+        else original_execute_handoffs
+    )
+    async def patched_execute_handoffs(
+        cls: agents.Runner, *args: Any, **kwargs: Any
+    ) -> Any:
+        """Patched execute_handoffs that creates handoff spans and ends agent span for handoffs"""
+        context_wrapper = kwargs.get("context_wrapper")
+        run_handoffs = kwargs.get("run_handoffs")
+        agent = kwargs.get("agent")
+
+        # Create Sentry handoff span for the first handoff (agents library only processes the first one)
+        if run_handoffs:
+            first_handoff = run_handoffs[0]
+            handoff_agent_name = first_handoff.handoff.agent_name
+            handoff_span(context_wrapper, agent, handoff_agent_name)
+
+        # Call original method with all parameters
+        try:
+            result = await original_execute_handoffs(*args, **kwargs)
+
+        finally:
+            # End span for current agent after handoff processing is complete
+            if agent and context_wrapper and _has_active_agent_span(context_wrapper):
+                _end_invoke_agent_span(context_wrapper, agent)
+
+        return result
+
+    @wraps(
+        original_execute_final_output.__func__
+        if hasattr(original_execute_final_output, "__func__")
+        else original_execute_final_output
+    )
+    async def patched_execute_final_output(
+        cls: agents.Runner, *args: Any, **kwargs: Any
+    ) -> Any:
+        """Patched execute_final_output that ends agent span for final outputs"""
+        agent = kwargs.get("agent")
+        context_wrapper = kwargs.get("context_wrapper")
+        final_output = kwargs.get("final_output")
+
+        # Call original method with all parameters
+        try:
+            result = await original_execute_final_output(*args, **kwargs)
+        finally:
+            # End span for current agent after final output processing is complete
+            if agent and context_wrapper and _has_active_agent_span(context_wrapper):
+                _end_invoke_agent_span(context_wrapper, agent, final_output)
+
+        return result
+
+    # Apply patches
+    agents.run.AgentRunner._run_single_turn = classmethod(patched_run_single_turn)
+    agents._run_impl.RunImpl.execute_handoffs = classmethod(patched_execute_handoffs)
+    agents._run_impl.RunImpl.execute_final_output = classmethod(
+        patched_execute_final_output
+    )

sentry_sdk/integrations/openai_agents/patches/models.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+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 agents
+except ImportError:
+    raise DidNotEnable("OpenAI Agents not installed")
+
+
+def _create_get_model_wrapper(
+    original_get_model: Callable[..., Any],
+) -> Callable[..., Any]:
+    """
+    Wraps the agents.Runner._get_model method to wrap the get_response method of the model to create a AI client span.
+    """
+
+    @wraps(
+        original_get_model.__func__
+        if hasattr(original_get_model, "__func__")
+        else original_get_model
+    )
+    def wrapped_get_model(
+        cls: agents.Runner, agent: agents.Agent, run_config: agents.RunConfig
+    ) -> agents.Model:
+        model = original_get_model(agent, run_config)
+        original_get_response = model.get_response
+
+        @wraps(original_get_response)
+        async def wrapped_get_response(*args: Any, **kwargs: Any) -> Any:
+            with ai_client_span(agent, kwargs) as span:
+                result = await original_get_response(*args, **kwargs)
+
+                update_ai_client_span(span, agent, kwargs, result)
+
+            return result
+
+        model.get_response = wrapped_get_response
+
+        return model
+
+    return wrapped_get_model

sentry_sdk/integrations/openai_agents/patches/runner.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from functools import wraps
+
+import sentry_sdk
+
+from ..spans import agent_workflow_span
+from ..utils import _capture_exception
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, Callable
+
+
+def _create_run_wrapper(original_func: Callable[..., Any]) -> Callable[..., Any]:
+    """
+    Wraps the agents.Runner.run methods to create a root span for the agent workflow runs.
+
+    Note agents.Runner.run_sync() is a wrapper around agents.Runner.run(),
+    so it does not need to be wrapped separately.
+    """
+
+    @wraps(original_func)
+    async def wrapper(*args: Any, **kwargs: Any) -> Any:
+        agent = args[0]
+        with agent_workflow_span(agent):
+            result = None
+            try:
+                result = await original_func(*args, **kwargs)
+                return result
+            except Exception as exc:
+                _capture_exception(exc)
+
+                # It could be that there is a "invoke agent" span still open
+                span = sentry_sdk.get_current_span()
+                if span is not None and span.timestamp is None:
+                    span.__exit__(None, None, None)
+
+                raise exc from None
+
+    return wrapper

sentry_sdk/integrations/openai_agents/patches/tools.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+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 agents
+except ImportError:
+    raise DidNotEnable("OpenAI Agents not installed")
+
+
+def _create_get_all_tools_wrapper(
+    original_get_all_tools: Callable[..., Any],
+) -> Callable[..., Any]:
+    """
+    Wraps the agents.Runner._get_all_tools method of the Runner class to wrap all function tools with Sentry instrumentation.
+    """
+
+    @wraps(
+        original_get_all_tools.__func__
+        if hasattr(original_get_all_tools, "__func__")
+        else original_get_all_tools
+    )
+    async def wrapped_get_all_tools(
+        cls: agents.Runner,
+        agent: agents.Agent,
+        context_wrapper: agents.RunContextWrapper,
+    ) -> list[agents.Tool]:
+        # Get the original tools
+        tools = await original_get_all_tools(agent, context_wrapper)
+
+        wrapped_tools = []
+        for tool in tools:
+            # Wrap only the function tools (for now)
+            if tool.__class__.__name__ != "FunctionTool":
+                wrapped_tools.append(tool)
+                continue
+
+            # Create a new FunctionTool with our wrapped invoke method
+            original_on_invoke = tool.on_invoke_tool
+
+            def create_wrapped_invoke(
+                current_tool: agents.Tool, current_on_invoke: Callable[..., Any]
+            ) -> Callable[..., Any]:
+                @wraps(current_on_invoke)
+                async def sentry_wrapped_on_invoke_tool(
+                    *args: Any, **kwargs: Any
+                ) -> Any:
+                    with execute_tool_span(current_tool, *args, **kwargs) as span:
+                        # We can not capture exceptions in tool execution here because
+                        # `_on_invoke_tool` is swallowing the exception here:
+                        # https://github.com/openai/openai-agents-python/blob/main/src/agents/tool.py#L409-L422
+                        # And because function_tool is a decorator with `default_tool_error_function` set as a default parameter
+                        # I was unable to monkey patch it because those are evaluated at module import time
+                        # and the SDK is too late to patch it. I was also unable to patch `_on_invoke_tool_impl`
+                        # because it is nested inside this import time code. As if they made it hard to patch on purpose...
+                        result = await current_on_invoke(*args, **kwargs)
+                        update_execute_tool_span(span, agent, current_tool, result)
+
+                    return result
+
+                return sentry_wrapped_on_invoke_tool
+
+            wrapped_tool = agents.FunctionTool(
+                name=tool.name,
+                description=tool.description,
+                params_json_schema=tool.params_json_schema,
+                on_invoke_tool=create_wrapped_invoke(tool, original_on_invoke),
+                strict_json_schema=tool.strict_json_schema,
+                is_enabled=tool.is_enabled,
+            )
+            wrapped_tools.append(wrapped_tool)
+
+        return wrapped_tools
+
+    return wrapped_get_all_tools

sentry_sdk/integrations/openai_agents/spans/__init__.py

@@ -0,0 +1,5 @@
+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 .handoff import handoff_span  # noqa: F401
+from .invoke_agent import invoke_agent_span, update_invoke_agent_span  # noqa: F401

sentry_sdk/integrations/openai_agents/spans/agent_workflow.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+import sentry_sdk
+
+from ..consts import SPAN_ORIGIN
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import agents
+
+
+def agent_workflow_span(agent: agents.Agent) -> sentry_sdk.tracing.Span:
+    # Create a transaction or a span if an transaction is already active
+    span = sentry_sdk.start_span(
+        name=f"{agent.name} workflow",
+        origin=SPAN_ORIGIN,
+    )
+
+    return span

sentry_sdk/integrations/openai_agents/spans/ai_client.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import sentry_sdk
+from sentry_sdk.consts import OP, SPANDATA
+
+from ..consts import SPAN_ORIGIN
+from ..utils import (
+    _set_agent_data,
+    _set_input_data,
+    _set_output_data,
+    _set_usage_data,
+)
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from agents import Agent
+    from typing import Any
+
+
+def ai_client_span(
+    agent: Agent, get_response_kwargs: dict[str, Any]
+) -> sentry_sdk.tracing.Span:
+    model_name = agent.model.model if hasattr(agent.model, "model") else agent.model
+    # TODO-anton: implement other types of operations. Now "chat" is hardcoded.
+    span = sentry_sdk.start_span(
+        op=OP.GEN_AI_CHAT,
+        description=f"chat {model_name}",
+        origin=SPAN_ORIGIN,
+    )
+    # TODO-anton: remove hardcoded stuff and replace something that also works for embedding and so on
+    span.set_attribute(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
+
+    return span
+
+
+def update_ai_client_span(
+    span: sentry_sdk.tracing.Span,
+    agent: Agent,
+    get_response_kwargs: dict[str, Any],
+    result: Any,
+) -> None:
+    _set_agent_data(span, agent)
+    _set_usage_data(span, result.usage)
+    _set_input_data(span, get_response_kwargs)
+    _set_output_data(span, result)

sentry_sdk/integrations/openai_agents/spans/execute_tool.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import sentry_sdk
+from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.scope import should_send_default_pii
+
+from ..consts import SPAN_ORIGIN
+from ..utils import _set_agent_data
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import agents
+    from typing import Any
+
+
+def execute_tool_span(
+    tool: agents.Tool, *args: Any, **kwargs: Any
+) -> sentry_sdk.tracing.Span:
+    span = sentry_sdk.start_span(
+        op=OP.GEN_AI_EXECUTE_TOOL,
+        name=f"execute_tool {tool.name}",
+        origin=SPAN_ORIGIN,
+    )
+
+    span.set_attribute(SPANDATA.GEN_AI_OPERATION_NAME, "execute_tool")
+
+    if tool.__class__.__name__ == "FunctionTool":
+        span.set_attribute(SPANDATA.GEN_AI_TOOL_TYPE, "function")
+
+    span.set_attribute(SPANDATA.GEN_AI_TOOL_NAME, tool.name)
+    span.set_attribute(SPANDATA.GEN_AI_TOOL_DESCRIPTION, tool.description)
+
+    if should_send_default_pii():
+        input = args[1]
+        span.set_attribute(SPANDATA.GEN_AI_TOOL_INPUT, input)
+
+    return span
+
+
+def update_execute_tool_span(
+    span: sentry_sdk.tracing.Span, agent: agents.Agent, tool: agents.Tool, result: Any
+) -> None:
+    _set_agent_data(span, agent)
+
+    if should_send_default_pii():
+        span.set_attribute(SPANDATA.GEN_AI_TOOL_OUTPUT, result)

sentry_sdk/integrations/openai_agents/spans/handoff.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+import sentry_sdk
+from sentry_sdk.consts import OP, SPANDATA
+
+from ..consts import SPAN_ORIGIN
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import agents
+
+
+def handoff_span(
+    context_wrapper: agents.RunContextWrapper,
+    from_agent: agents.Agent,
+    to_agent_name: str,
+) -> None:
+    with sentry_sdk.start_span(
+        op=OP.GEN_AI_HANDOFF,
+        name=f"handoff from {from_agent.name} to {to_agent_name}",
+        origin=SPAN_ORIGIN,
+    ) as span:
+        span.set_attribute(SPANDATA.GEN_AI_OPERATION_NAME, "handoff")

sentry_sdk/integrations/openai_agents/spans/invoke_agent.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+import sentry_sdk
+from sentry_sdk.consts import OP, SPANDATA
+
+from ..consts import SPAN_ORIGIN
+from ..utils import _set_agent_data
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import agents
+    from typing import Any
+
+
+def invoke_agent_span(
+    context_wrapper: agents.RunContextWrapper, agent: agents.Agent
+) -> sentry_sdk.tracing.Span:
+    span = sentry_sdk.start_span(
+        op=OP.GEN_AI_INVOKE_AGENT,
+        name=f"invoke_agent {agent.name}",
+        origin=SPAN_ORIGIN,
+    )
+    span.__enter__()
+
+    span.set_attribute(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
+
+    _set_agent_data(span, agent)
+
+    context_wrapper._sentry_invoke_agent_span = span
+
+    return span
+
+
+def update_invoke_agent_span(
+    context_wrapper: agents.RunContextWrapper, agent: agents.Agent, output: Any
+) -> None:
+    span = getattr(context_wrapper, "_sentry_invoke_agent_span", None)
+    if span is not None:
+        span.__exit__(None, None, None)
+        del context_wrapper._sentry_invoke_agent_span