ref(openai-agents): Use first class tool hooks when available

sentry_sdk/integrations/openai_agents/__init__.py

@@ -46,17 +46,19 @@
     from agents.run_internal.run_steps import SingleStepResult
 
 
-def _patch_runner() -> None:
+def _patch_runner(use_run_hooks: "bool") -> 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.
     agents.run.DEFAULT_AGENT_RUNNER.run = _create_run_wrapper(
-        agents.run.DEFAULT_AGENT_RUNNER.run
+        agents.run.DEFAULT_AGENT_RUNNER.run,
+        use_run_hooks=use_run_hooks,
     )
 
     # Patch streaming runner
     agents.run.DEFAULT_AGENT_RUNNER.run_streamed = _create_run_streamed_wrapper(
-        agents.run.DEFAULT_AGENT_RUNNER.run_streamed
+        agents.run.DEFAULT_AGENT_RUNNER.run_streamed,
+        use_run_hooks=use_run_hooks,
     )
 
 
@@ -92,26 +94,19 @@ class OpenAIAgentsIntegration(Integration):
     @staticmethod
     def setup_once() -> None:
         _patch_error_tracing()
-        _patch_runner()
 
         library_version = parse_version(OPENAI_AGENTS_VERSION)
+        # ToolContext.tool_arguments added in https://github.com/openai/openai-agents-python/commit/5e1db14da542c77f8fdd5e2e26017977ae415813
+        use_run_hooks = library_version is not None and library_version >= (0, 3, 2)
+
+        _patch_runner(use_run_hooks=use_run_hooks)
+
         if library_version is not None and library_version >= (
             0,
             8,
         ):
             if run_loop is not None:
 
-                @wraps(run_loop.get_all_tools)
-                async def new_wrapped_get_all_tools(
-                    agent: "agents.Agent",
-                    context_wrapper: "agents.RunContextWrapper",
-                ) -> "list[agents.Tool]":
-                    return await _get_all_tools(
-                        run_loop.get_all_tools, agent, context_wrapper
-                    )
-
-                agents.run.get_all_tools = new_wrapped_get_all_tools
-
                 @wraps(run_loop.run_single_turn)
                 async def new_wrapped_run_single_turn(
                     *args: "Any", **kwargs: "Any"
@@ -175,17 +170,22 @@ async def new_wrapped_final_output(
 
             return
 
-        original_get_all_tools = AgentRunner._get_all_tools
-
-        @wraps(AgentRunner._get_all_tools.__func__)
-        async def old_wrapped_get_all_tools(
-            cls: "agents.Runner",
-            agent: "agents.Agent",
-            context_wrapper: "agents.RunContextWrapper",
-        ) -> "list[agents.Tool]":
-            return await _get_all_tools(original_get_all_tools, agent, context_wrapper)
+        if not use_run_hooks:
+            original_get_all_tools = AgentRunner._get_all_tools
+
+            @wraps(AgentRunner._get_all_tools.__func__)
+            async def old_wrapped_get_all_tools(
+                cls: "agents.Runner",
+                agent: "agents.Agent",
+                context_wrapper: "agents.RunContextWrapper",
+            ) -> "list[agents.Tool]":
+                return await _get_all_tools(
+                    original_get_all_tools, agent, context_wrapper
+                )
 
-        agents.run.AgentRunner._get_all_tools = classmethod(old_wrapped_get_all_tools)
+            agents.run.AgentRunner._get_all_tools = classmethod(
+                old_wrapped_get_all_tools
+            )
 
         original_get_model = AgentRunner._get_model
 

sentry_sdk/integrations/openai_agents/patches/runner.py

@@ -4,24 +4,122 @@
 import sentry_sdk
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable
+from sentry_sdk.scope import should_send_default_pii
 from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.utils import capture_internal_exceptions, reraise
 
-from ..spans import agent_workflow_span, update_invoke_agent_span
+from ..spans import (
+    agent_workflow_span,
+    execute_tool_span,
+    update_execute_tool_span,
+    update_invoke_agent_span,
+)
 from ..utils import _capture_exception
 
 try:
+    from agents import FunctionTool, RunHooks
     from agents.exceptions import AgentsException
 except ImportError:
     raise DidNotEnable("OpenAI Agents not installed")
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypeVar
 
 if TYPE_CHECKING:
     from typing import Any, AsyncIterator, Callable
 
+    from agents import Agent, Tool, ToolContext
 
-def _create_run_wrapper(original_func: "Callable[..., Any]") -> "Callable[..., Any]":
+
+TContext = TypeVar("TContext")
+
+
+class _SentryRunHooks(RunHooks[TContext]):  # type: ignore[misc]
+    """
+    Responsible for creating and managing Execute Tool spans. These spans are
+    stored on the ToolContext reference that is shared between `on_tool_start()`
+    and `on_tool_end()`
+    """
+
+    async def on_tool_start(
+        self,
+        context: "ToolContext[TContext]",
+        agent: "Agent[TContext]",
+        tool: "Tool",
+    ) -> "None":
+        if not isinstance(tool, FunctionTool):
+            return
+
+        span = execute_tool_span(tool, agent)
+        span.__enter__()
+        context._sentry_execute_tool_span = span
+
+        if not should_send_default_pii():
+            return
+
+        if isinstance(span, StreamedSpan):
+            span.set_attribute(SPANDATA.GEN_AI_TOOL_INPUT, context.tool_arguments)
+        else:
+            span.set_data(SPANDATA.GEN_AI_TOOL_INPUT, context.tool_arguments)
+
+    async def on_tool_end(
+        self,
+        context: "ToolContext[TContext]",
+        agent: "Agent[TContext]",
+        tool: "Tool",
+        result: "object",
+    ) -> "None":
+        if not isinstance(tool, FunctionTool):
+            return
+
+        span = getattr(context, "_sentry_execute_tool_span", None)
+        if span is not None:
+            del context._sentry_execute_tool_span
+            update_execute_tool_span(span, agent, tool, result)
+            span.__exit__(None, None, None)
+
+
+def _patch_run_hooks(hooks: "RunHooks[TContext]") -> None:
+    """
+    Patch a RunHooks instance. This is used when the user have themselves provided
+    a RunHooks instance, as only one instance can be passed to `AgentRunner.run()`
+    and `AgentRunner.run_streamed()` functions.
+    """
+    is_already_patched = getattr(hooks, "_sentry_is_patched", False)
+    if is_already_patched:
+        return
+
+    original_on_tool_start = hooks.on_tool_start
+    original_on_tool_end = hooks.on_tool_end
+
+    sentry_hooks = _SentryRunHooks()  # type: ignore[var-annotated]
+
+    @wraps(original_on_tool_start)
+    async def on_tool_start(
+        context: "ToolContext[TContext]", agent: "Agent[TContext]", tool: "Tool"
+    ) -> "None":
+        with capture_internal_exceptions():
+            await sentry_hooks.on_tool_start(context, agent, tool)
+        await original_on_tool_start(context, agent, tool)
+
+    @wraps(original_on_tool_end)
+    async def on_tool_end(
+        context: "ToolContext[TContext]",
+        agent: "Agent[TContext]",
+        tool: "Tool",
+        result: "object",
+    ) -> "None":
+        with capture_internal_exceptions():
+            await sentry_hooks.on_tool_end(context, agent, tool, result)
+        await original_on_tool_end(context, agent, tool, result)
+
+    hooks._sentry_is_patched = True
+    hooks.on_tool_start = on_tool_start
+    hooks.on_tool_end = on_tool_end
+
+
+def _create_run_wrapper(
+    original_func: "Callable[..., Any]", use_run_hooks: "bool"
+) -> "Callable[..., Any]":
     """
     Wraps the agents.Runner.run methods to
     - create and manage a root span for the agent workflow runs.
@@ -33,6 +131,13 @@ def _create_run_wrapper(original_func: "Callable[..., Any]") -> "Callable[..., A
 
     @wraps(original_func)
     async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
+        if use_run_hooks:
+            hooks = kwargs.get("hooks")
+            if hooks is not None:
+                _patch_run_hooks(hooks=hooks)
+            else:
+                kwargs["hooks"] = _SentryRunHooks()
+
         # 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():
@@ -123,7 +228,7 @@ async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
 
 
 def _create_run_streamed_wrapper(
-    original_func: "Callable[..., Any]",
+    original_func: "Callable[..., Any]", use_run_hooks: "bool"
 ) -> "Callable[..., Any]":
     """
     Wraps the agents.Runner.run_streamed method to
@@ -173,6 +278,14 @@ def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
         else:
             args = (agent, *args[1:])
 
+        if use_run_hooks:
+            sentry_hooks = _SentryRunHooks()  # type: ignore[var-annotated]
+            hooks = kwargs.get("hooks")
+            if hooks is not None:
+                _patch_run_hooks(hooks=hooks)
+            else:
+                kwargs["hooks"] = sentry_hooks
+
         try:
             # Call original function to get RunResultStreaming
             run_result = original_func(*args, **kwargs)

sentry_sdk/integrations/openai_agents/patches/tools.py

@@ -1,7 +1,10 @@
 from functools import wraps
 from typing import TYPE_CHECKING
 
+from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable
+from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.traces import StreamedSpan
 
 from ..spans import execute_tool_span, update_execute_tool_span
 
@@ -53,6 +56,14 @@ async def sentry_wrapped_on_invoke_tool(
                     result = await current_on_invoke(*args, **kwargs)
                     update_execute_tool_span(span, agent, current_tool, result)
 
+                    if not should_send_default_pii():
+                        return result
+
+                    if isinstance(span, StreamedSpan):
+                        span.set_attribute(SPANDATA.GEN_AI_TOOL_INPUT, args[1])
+                    else:
+                        span.set_data(SPANDATA.GEN_AI_TOOL_INPUT, args[1])
+
                 return result
 
             return sentry_wrapped_on_invoke_tool

sentry_sdk/integrations/openai_agents/spans/execute_tool.py

@@ -30,8 +30,6 @@ def execute_tool_span(
                 SPANDATA.GEN_AI_TOOL_DESCRIPTION: tool.description,
             },
         )
-
-        set_on_span = span.set_attribute
     else:
         span = sentry_sdk.start_span(
             op=OP.GEN_AI_EXECUTE_TOOL,
@@ -44,12 +42,6 @@ def execute_tool_span(
         span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool.name)
         span.set_data(SPANDATA.GEN_AI_TOOL_DESCRIPTION, tool.description)
 
-        set_on_span = span.set_data
-
-    if should_send_default_pii():
-        input = args[1]
-        set_on_span(SPANDATA.GEN_AI_TOOL_INPUT, input)
-
     return span