8000 Create transaction for runner.run · getsentry/sentry-python@0385775 · GitHub
[go: up one dir, main page]

Skip to content

Commit 0385775

Browse files
committed
Create transaction for runner.run
1 parent 55f3ea8 commit 0385775

File tree

1 file changed

+40
-3
lines changed

1 file changed

+40
-3
lines changed

sentry_sdk/integrations/openai_agents.py

Lines changed: 40 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,13 @@
11
import sentry_sdk
22
from sentry_sdk.integrations import DidNotEnable, Integration
33
from sentry_sdk.utils import event_from_exception
4+
from functools import wraps
5+
import asyncio
46

57
from typing import Any
68

79
try:
10+
import agents
811
from agents import (
912
Agent,
1013
RunContextWrapper,
@@ -82,14 +85,48 @@ async def on_handoff(
8285
current_span.__exit__(None, None, None)
8386

8487

88+
def _get_span_function():
89+
current_span = sentry_sdk.get_current_span()
90+
is_transaction = (
91+
current_span is not None and current_span.containing_transaction == current_span
92+
)
93+
return sentry_sdk.start_span if is_transaction else sentry_sdk.start_transaction
94+
95+
96+
def _create_wrapper(original_func):
97+
is_async = asyncio.iscoroutinefunction(original_func)
98+
99+
@classmethod
100+
@wraps(original_func)
101+
async def async_wrapper(cls, *args, **kwargs):
102+
agent = args[0]
103+
with _get_span_function()(name=agent.name):
104+
result = await original_func(*args, **kwargs)
105+
return result
106+
107+
@classmethod
108+
@wraps(original_func)
109+
def sync_wrapper(cls, *args, **kwargs):
110+
agent = args[0]
111+
with _get_span_function()(name=agent.name):
112+
result = original_func(*args, **kwargs)
113+
return result
114+
115+
return async_wrapper if is_async else sync_wrapper
116+
117+
118+
def _patch_runner():
119+
agents.Runner.run = _create_wrapper(agents.Runner.run)
120+
agents.Runner.run_sync = _create_wrapper(agents.Runner.run_sync)
121+
agents.Runner.run_streamed = _create_wrapper(agents.Runner.run_streamed)
122+
123+
85124
class OpenAIAgentsIntegration(Integration):
86125
identifier = "openai_agents"
87126
origin = f"auto.ai.{identifier}"
88127

89-
# def __init__(self):
90-
# pass
91-
92128
@staticmethod
93129
def setup_once():
94130
# type: () -> None
131+
_patch_runner()
95132
pass

0 commit comments

Comments
 (0)
0