Adapt AI Agent Monitoring for potel (#4551)

Port the `openai-agents` integration to `potel-base`.

---------

Co-authored-by: Ivana Kellyer <[email protected]>
Co-authored-by: Anton Pirker <[email protected]>

Author: 3 people

.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,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