Create transaction for runner.run

Author: antonpirker

sentry_sdk/integrations/openai_agents.py

@@ -1,10 +1,13 @@
 import sentry_sdk
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.utils import event_from_exception
+from functools import wraps
+import asyncio
 
 from typing import Any
 
 try:
+    import agents
     from agents import (
         Agent,
         RunContextWrapper,
@@ -82,14 +85,48 @@ async def on_handoff(
             current_span.__exit__(None, None, None)
 
 
+def _get_span_function():
+    current_span = sentry_sdk.get_current_span()
+    is_transaction = (
+        current_span is not None and current_span.containing_transaction == current_span
+    )
+    return sentry_sdk.start_span if is_transaction else sentry_sdk.start_transaction
+
+
+def _create_wrapper(original_func):
+    is_async = asyncio.iscoroutinefunction(original_func)
+
+    @classmethod
+    @wraps(original_func)
+    async def async_wrapper(cls, *args, **kwargs):
+        agent = args[0]
+        with _get_span_function()(name=agent.name):
+            result = await original_func(*args, **kwargs)
+            return result
+
+    @classmethod
+    @wraps(original_func)
+    def sync_wrapper(cls, *args, **kwargs):
+        agent = args[0]
+        with _get_span_function()(name=agent.name):
+            result = original_func(*args, **kwargs)
+            return result
+
+    return async_wrapper if is_async else sync_wrapper
+
+
+def _patch_runner():
+    agents.Runner.run = _create_wrapper(agents.Runner.run)
+    agents.Runner.run_sync = _create_wrapper(agents.Runner.run_sync)
+    agents.Runner.run_streamed = _create_wrapper(agents.Runner.run_streamed)
+
+
 class OpenAIAgentsIntegration(Integration):
     identifier = "openai_agents"
     origin = f"auto.ai.{identifier}"
 
-    # def __init__(self):
-    #     pass
-
     @staticmethod
     def setup_once():
         # type: () -> None
+        _patch_runner()
         pass