fix(openai-agents): Patch tool functions following library refactor

sentry_sdk/integrations/openai_agents/__init__.py

@@ -1,8 +1,10 @@
 from sentry_sdk.integrations import DidNotEnable, Integration
+from sentry_sdk.utils import parse_version
 
 from .patches import (
     _create_get_model_wrapper,
-    _create_get_all_tools_wrapper,
+    _create_runner_get_all_tools_wrapper,
+    _create_run_loop_get_all_tools_wrapper,
     _create_run_wrapper,
     _create_run_streamed_wrapper,
     _patch_agent_run,
@@ -17,11 +19,21 @@
     # after it, even if we don't use it.
     import agents
     from agents.run import DEFAULT_AGENT_RUNNER
+    from agents.version import __version__ as OPENAI_AGENTS_VERSION
 
 except ImportError:
     raise DidNotEnable("OpenAI Agents not installed")
 
 
+try:
+    # AgentRunner methods moved in v0.8
+    # https://github.com/openai/openai-agents-python/commit/3ce7c24d349b77bb750062b7e0e856d9ff48a5d5#diff-7470b3a5c5cbe2fcbb2703dc24f326f45a5819d853be2b1f395d122d278cd911
+    from agents.run_internal import run_loop, turn_preparation
+except ImportError:
+    run_loop = None
+    turn_preparation = None
+
+
 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
@@ -45,9 +57,15 @@ def _patch_model() -> None:
     )
 
 
-def _patch_tools() -> None:
+def _patch_agent_runner_get_all_tools() -> None:
     agents.run.AgentRunner._get_all_tools = classmethod(
-        _create_get_all_tools_wrapper(agents.run.AgentRunner._get_all_tools),
+        _create_runner_get_all_tools_wrapper(agents.run.AgentRunner._get_all_tools),
+    )
+
+
+def _patch_run_get_all_tools() -> None:
+    agents.run.get_all_tools = _create_run_loop_get_all_tools_wrapper(
+        run_loop.get_all_tools
     )
 
 
@@ -57,6 +75,15 @@ class OpenAIAgentsIntegration(Integration):
     @staticmethod
     def setup_once() -> None:
         _patch_error_tracing()
-        _patch_tools()
         _patch_model()
         _patch_runner()
+
+        library_version = parse_version(OPENAI_AGENTS_VERSION)
+        if library_version is not None and library_version >= (
+            0,
+            8,
+        ):
+            _patch_run_get_all_tools()
+            return
+
+        _patch_agent_runner_get_all_tools()

sentry_sdk/integrations/openai_agents/patches/__init__.py

@@ -1,5 +1,8 @@
 from .models import _create_get_model_wrapper  # noqa: F401
-from .tools import _create_get_all_tools_wrapper  # noqa: F401
+from .tools import (
+    _create_runner_get_all_tools_wrapper,
+    _create_run_loop_get_all_tools_wrapper,
+)  # noqa: F401
 from .runner import _create_run_wrapper, _create_run_streamed_wrapper  # noqa: F401
 from .agent_run import _patch_agent_run  # noqa: F401
 from .error_tracing import _patch_error_tracing  # noqa: F401

sentry_sdk/integrations/openai_agents/patches/tools.py

@@ -7,21 +7,70 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any, Callable
+    from typing import Any, Callable, Awaitable
 
 try:
     import agents
 except ImportError:
     raise DidNotEnable("OpenAI Agents not installed")
 
 
-def _create_get_all_tools_wrapper(
+async def _get_all_tools(
+    original_get_all_tools: "Callable[..., Awaitable[list[agents.Tool]]]",
+    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
+
+
+def _create_runner_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__")
@@ -32,51 +81,19 @@ async def wrapped_get_all_tools(
         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 await _get_all_tools(original_get_all_tools, agent, context_wrapper)
+
+    return wrapped_get_all_tools
+
+
+def _create_run_loop_get_all_tools_wrapper(
+    original_get_all_tools: "Callable[..., Any]",
+) -> "Callable[..., Any]":
+    @wraps(original_get_all_tools)
+    async def wrapped_get_all_tools(
+        agent: "agents.Agent",
+        context_wrapper: "agents.RunContextWrapper",
+    ) -> "list[agents.Tool]":
+        return await _get_all_tools(original_get_all_tools, agent, context_wrapper)
 
     return wrapped_get_all_tools