Revert "Remove OpenAI agents from Potel (#4523)"

This reverts commit ca42492.

Author: szokeasaurusrex

.github/workflows/test-integrations-ai.yml

@@ -59,6 +59,10 @@ jobs:
         run: |
           set -x # print commands that are executed
           ./scripts/runtox.sh "py${{ matrix.python-version }}-openai-latest"
+      - name: Test openai_agents latest
+        run: |
+          set -x # print commands that are executed
+          ./scripts/runtox.sh "py${{ matrix.python-version }}-openai_agents-latest"
       - name: Test huggingface_hub latest
         run: |
           set -x # print commands that are executed
@@ -121,6 +125,10 @@ jobs:
         run: |
           set -x # print commands that are executed
           ./scripts/runtox.sh --exclude-latest "py${{ matrix.python-version }}-openai"
+      - name: Test openai_agents pinned
+        run: |
+          set -x # print commands that are executed
+          ./scripts/runtox.sh --exclude-latest "py${{ matrix.python-version }}-openai_agents"
       - name: Test huggingface_hub pinned
         run: |
           set -x # print commands that are executed

scripts/populate_tox/config.py

@@ -152,6 +152,12 @@
     "loguru": {
         "package": "loguru",
     },
+    "openai_agents": {
+        "package": "openai-agents",
+        "deps": {
+            "*": ["pytest-asyncio"],
+        },
+    },
     "openfeature": {
         "package": "openfeature-sdk",
     },

scripts/populate_tox/tox.jinja

@@ -372,6 +372,7 @@ setenv =
     litestar: TESTPATH=tests/integrations/litestar
     loguru: TESTPATH=tests/integrations/loguru
     openai: TESTPATH=tests/integrations/openai
+    openai_agents: TESTPATH=tests/integrations/openai_agents
     openfeature: TESTPATH=tests/integrations/openfeature
     pure_eval: TESTPATH=tests/integrations/pure_eval
     pymongo: TESTPATH=tests/integrations/pymongo

scripts/split_tox_gh_actions/split_tox_gh_actions.py

@@ -63,6 +63,7 @@
         "cohere",
         "langchain",
         "openai",
+        "openai_agents",
         "huggingface_hub",
     ],
     "Cloud": [

sentry_sdk/integrations/__init__.py

@@ -147,6 +147,7 @@ def iter_default_integrations(
     "launchdarkly": (9, 8, 0),
     "loguru": (0, 7, 0),
     "openai": (1, 0, 0),
+    "openai_agents": (0, 0, 19),
     "openfeature": (0, 7, 1),
     "quart": (0, 16, 0),
     "ray": (2, 7, 0),

sentry_sdk/integrations/openai_agents/__init__.py

@@ -0,0 +1,53 @@
+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():
+    # type: () -> 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():
+    # type: () -> None
+    agents.run.AgentRunner._get_model = classmethod(
+        _create_get_model_wrapper(agents.run.AgentRunner._get_model),
+    )
+
+
+def _patch_tools():
+    # type: () -> 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():
+        # type: () -> 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,143 @@
+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():
+    # type: () -> 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, agent):
+        # type: (agents.RunContextWrapper, 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, agent, output=None):
+        # type: (agents.RunContextWrapper, agents.Agent, Optional[Any]) -> 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):
+        # type: (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):
+        # type: (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, *args, **kwargs):
+        # type: (agents.Runner, *Any, **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
+        result = await original_run_single_turn(*args, **kwargs)
+
+        return result
+
+    @wraps(
+        original_execute_handoffs.__func__
+        if hasattr(original_execute_handoffs, "__func__")
+        else original_execute_handoffs
+    )
+    async def patched_execute_handoffs(cls, *args, **kwargs):
+        # type: (agents.Runner, *Any, **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, *args, **kwargs):
+        # type: (agents.Runner, *Any, **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,50 @@
+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):
+    # type: (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, agent, run_config):
+        # type: (agents.Runner, agents.Agent, 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, **kwargs):
+            # type: (*Any, **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