|
1 | 1 | import sentry_sdk
|
2 | 2 | from sentry_sdk.integrations import DidNotEnable, Integration
|
3 | 3 | from sentry_sdk.utils import event_from_exception
|
| 4 | +from functools import wraps |
| 5 | +import asyncio |
4 | 6 |
|
5 | 7 | from typing import Any
|
6 | 8 |
|
7 | 9 | try:
|
| 10 | + import agents |
8 | 11 | from agents import (
|
9 | 12 | Agent,
|
10 | 13 | RunContextWrapper,
|
@@ -82,14 +85,48 @@ async def on_handoff(
|
82 | 85 | current_span.__exit__(None, None, None)
|
83 | 86 |
|
84 | 87 |
|
| 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 | + |
85 | 124 | class OpenAIAgentsIntegration(Integration):
|
86 | 125 | identifier = "openai_agents"
|
87 | 126 | origin = f"auto.ai.{identifier}"
|
88 | 127 |
|
89 |
| - # def __init__(self): |
90 |
| - # pass |
91 |
| - |
92 | 128 | @staticmethod
|
93 | 129 | def setup_once():
|
94 | 130 | # type: () -> None
|
| 131 | + _patch_runner() |
95 | 132 | pass
|
0 commit comments