1
1
from functools import wraps
2
2
3
3
from sentry_sdk .integrations import DidNotEnable
4
+ from sentry_sdk .integrations .openai_agents .spans .guardrail import (
5
+ guardrail_span ,
6
+ update_guardrail_span ,
7
+ )
4
8
from ..spans import invoke_agent_span , update_invoke_agent_span , handoff_span
5
9
6
10
from typing import TYPE_CHECKING
@@ -26,6 +30,13 @@ def _patch_agent_run():
26
30
original_execute_handoffs = agents ._run_impl .RunImpl .execute_handoffs
27
31
original_execute_final_output = agents ._run_impl .RunImpl .execute_final_output
28
32
33
+ original_run_single_input_guardrail = (
34
+ agents ._run_impl .RunImpl .run_single_input_guardrail
35
+ )
36
+ original_run_single_output_guardrail = (
37
+ agents ._run_impl .RunImpl .run_single_output_guardrail
38
+ )
39
+
29
40
def _start_invoke_agent_span (context_wrapper , agent , kwargs ):
30
41
# type: (agents.RunContextWrapper, agents.Agent, dict[str, Any]) -> None
31
42
"""Start an agent invocation span"""
@@ -132,9 +143,47 @@ async def patched_execute_final_output(cls, *args, **kwargs):
132
143
133
144
return result
134
145
146
+ @wraps (
147
+ original_run_single_input_guardrail .__func__
148
+ if hasattr (original_run_single_input_guardrail , "__func__" )
149
+ else original_run_single_input_guardrail
150
+ )
151
+ async def patched_run_single_input_guardrail (cls , * args , ** kwargs ):
152
+ # type: (agents.Runner, *Any, **Any) -> Any
153
+ agent = args [0 ]
154
+ guardrail = args [1 ]
155
+
156
+ with guardrail_span (guardrail , "input" , args , kwargs ) as span :
157
+ result = await original_run_single_input_guardrail (* args , ** kwargs )
158
+ update_guardrail_span (span , agent , guardrail , "input" , result )
159
+
160
+ return result
161
+
162
+ @wraps (
163
+ original_run_single_output_guardrail .__func__
164
+ if hasattr (original_run_single_output_guardrail , "__func__" )
165
+ else original_run_single_output_guardrail
166
+ )
167
+ async def patched_run_single_output_guardrail (cls , * args , ** kwargs ):
168
+ # type: (agents.Runner, *Any, **Any) -> Any
169
+ guardrail = args [0 ]
170
+ agent = args [1 ]
171
+
172
+ with guardrail_span (guardrail , "output" , args , kwargs ) as span :
173
+ result = await original_run_single_output_guardrail (* args , ** kwargs )
174
+ update_guardrail_span (span , agent , guardrail , "output" , result )
175
+
176
+ return result
177
+
135
178
# Apply patches
136
179
agents .run .AgentRunner ._run_single_turn = classmethod (patched_run_single_turn )
137
180
agents ._run_impl .RunImpl .execute_handoffs = classmethod (patched_execute_handoffs )
138
181
agents ._run_impl .RunImpl .execute_final_output = classmethod (
139
182
patched_execute_final_output
140
183
)
184
+ agents ._run_impl .RunImpl .run_single_input_guardrail = classmethod (
185
+ patched_run_single_input_guardrail
186
+ )
187
+ agents ._run_impl .RunImpl .run_single_output_guardrail = classmethod (
188
+ patched_run_single_output_guardrail
189
+ )
0 commit comments