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

Author: alexander-alderman-webb

sentry_sdk/integrations/openai_agents/__init__.py

@@ -1,11 +1,21 @@
 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_model_wrapper,
+    _create_run_internal_get_model_wrapper,
+    _create_runner_get_all_tools_wrapper,
+    _create_run_internal_get_all_tools_wrapper,
     _create_run_wrapper,
     _create_run_streamed_wrapper,
-    _patch_agent_run,
+    _patch_agent_runner_run_single_turn,
+    _patch_run_loop_run_single_turn,
+    _patch_agent_runner_run_single_turn_streamed,
+    _patch_run_loop_run_single_turn_streamed,
+    _patch_run_impl_execute_handoffs,
+    _patch_turn_resolution_execute_handoffs,
+    _patch_run_impl_execute_final_output,
+    _patch_turn_resolution_execute_final_output,
     _patch_error_tracing,
 )
 
@@ -17,37 +27,42 @@
     # 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")
 
 
-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.
-    agents.run.DEFAULT_AGENT_RUNNER.run = _create_run_wrapper(
-        agents.run.DEFAULT_AGENT_RUNNER.run
-    )
+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
 
-    # Patch streaming runner
-    agents.run.DEFAULT_AGENT_RUNNER.run_streamed = _create_run_streamed_wrapper(
-        agents.run.DEFAULT_AGENT_RUNNER.run_streamed
-    )
 
-    # Creating the actual spans for each agent run (works for both streaming and non-streaming).
-    _patch_agent_run()
+def _patch_agent_runner_get_model() -> None:
+    agents.run.AgentRunner._get_model = classmethod(
+        _create_runner_get_model_wrapper(agents.run.AgentRunner._get_model),
+    )
 
 
-def _patch_model() -> None:
-    agents.run.AgentRunner._get_model = classmethod(
-        _create_get_model_wrapper(agents.run.AgentRunner._get_model),
+def _patch_run_internal_get_model():
+    agents.run_internal.run_loop.get_model = _create_run_internal_get_model_wrapper(
+        turn_preparation.get_model
     )
 
 
-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_internal_get_all_tools_wrapper(
+        run_loop.get_all_tools
     )
 
 
@@ -56,7 +71,38 @@ class OpenAIAgentsIntegration(Integration):
 
     @staticmethod
     def setup_once() -> 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
+        )
+
+        # Patch streaming runner
+        agents.run.DEFAULT_AGENT_RUNNER.run_streamed = _create_run_streamed_wrapper(
+            agents.run.DEFAULT_AGENT_RUNNER.run_streamed
+        )
+
+        if parse_version(OPENAI_AGENTS_VERSION) >= (
+            0,
+            8,
+        ):
+            _patch_error_tracing()
+            _patch_run_get_all_tools()
+            _patch_run_internal_get_model()
+
+            _patch_run_loop_run_single_turn()
+            _patch_run_loop_run_single_turn_streamed()
+            _patch_turn_resolution_execute_handoffs()
+            _patch_turn_resolution_execute_final_output()
+
+            return
+
         _patch_error_tracing()
-        _patch_tools()
-        _patch_model()
-        _patch_runner()
+        _patch_agent_runner_get_all_tools()
+        _patch_agent_runner_get_model()
+
+        _patch_agent_runner_run_single_turn()
+        _patch_agent_runner_run_single_turn_streamed()
+        _patch_run_impl_execute_handoffs()
+        _patch_run_impl_execute_final_output()

sentry_sdk/integrations/openai_agents/patches/__init__.py

@@ -1,5 +1,20 @@
-from .models import _create_get_model_wrapper  # noqa: F401
-from .tools import _create_get_all_tools_wrapper  # noqa: F401
+from .models import (
+    _create_runner_get_model_wrapper,
+    _create_run_internal_get_model_wrapper,
+)  # noqa: F401
+from .tools import (
+    _create_runner_get_all_tools_wrapper,
+    _create_run_internal_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 .agent_run import (
+    _patch_agent_runner_run_single_turn,
+    _patch_run_loop_run_single_turn,
+    _patch_agent_runner_run_single_turn_streamed,
+    _patch_run_loop_run_single_turn_streamed,
+    _patch_run_impl_execute_handoffs,
+    _patch_turn_resolution_execute_handoffs,
+    _patch_run_impl_execute_final_output,
+    _patch_turn_resolution_execute_final_output,
+)  # noqa: F401
 from .error_tracing import _patch_error_tracing  # noqa: F401