Skip to content

Commit 9238c28

Browse files
rm-openairm-openai
committed
Allow previous_response_id to be passed to input guardrails
1 parent af80e3a commit 9238c28

File tree

6 files changed

+281
-25
lines changed

6 files changed

+281
-25
lines changed

‎src/agents/__init__.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,10 @@
1818
)
1919
from .guardrail import (
2020
GuardrailFunctionOutput,
21+
InputGuardailInputs,
2122
InputGuardrail,
23+
InputGuardrailFunction,
24+
InputGuardrailFunctionLegacy,
2225
InputGuardrailResult,
2326
OutputGuardrail,
2427
OutputGuardrailResult,
@@ -174,6 +177,9 @@ def enable_verbose_stdout_logging():
174177
"OutputGuardrail",
175178
"OutputGuardrailResult",
176179
"GuardrailFunctionOutput",
180+
"InputGuardailInputs",
181+
"InputGuardrailFunction",
182+
"InputGuardrailFunctionLegacy",
177183
"input_guardrail",
178184
"output_guardrail",
179185
"handoff",

‎src/agents/_run_impl.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -688,10 +688,11 @@ async def run_single_input_guardrail(
688688
agent: Agent[Any],
689689
guardrail: InputGuardrail[TContext],
690690
input: str | list[TResponseInputItem],
691+
previous_response_id: str | None,
691692
context: RunContextWrapper[TContext],
692693
) -> InputGuardrailResult:
693694
with guardrail_span(guardrail.get_name()) as span_guardrail:
694-
result = await guardrail.run(agent, input, context)
695+
result = await guardrail.run(agent, input, context, previous_response_id)
695696
span_guardrail.span_data.triggered = result.output.tripwire_triggered
696697
return result
697698

‎src/agents/guardrail.py

Lines changed: 43 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,9 @@
33
import inspect
44
from collections.abc import Awaitable
55
from dataclasses import dataclass
6-
from typing import TYPE_CHECKING, Any, Callable, Generic, Union, overload
6+
from typing import TYPE_CHECKING, Any, Callable, Generic, Union, cast, overload
77

8-
from typing_extensions import TypeVar
8+
from typing_extensions import TypeAlias, TypeVar
99

1010
from .exceptions import UserError
1111
from .items import TResponseInputItem
@@ -68,6 +68,31 @@ class OutputGuardrailResult:
6868
"""The output of the guardrail function."""
6969

7070

71+
InputGuardrailFunctionLegacy: TypeAlias = Callable[
72+
[RunContextWrapper[TContext], "Agent[Any]", Union[str, list[TResponseInputItem]]],
73+
MaybeAwaitable[GuardrailFunctionOutput],
74+
]
75+
"""The legacy guardrail function signature, retained for backwards compatibility. Of the form:
76+
def my_guardrail(ctx, agent, input)
77+
"""
78+
79+
80+
@dataclass
81+
class InputGuardailInputs:
82+
agent: Agent[Any]
83+
input: str | list[TResponseInputItem]
84+
previous_response_id: str | None
85+
86+
87+
InputGuardrailFunction: TypeAlias = Callable[
88+
[RunContextWrapper[TContext], InputGuardailInputs],
89+
MaybeAwaitable[GuardrailFunctionOutput],
90+
]
91+
"""The new guardrail function signature, of the form:
92+
def my_guardrail(ctx, inputs)
93+
"""
94+
95+
7196
@dataclass
7297
class InputGuardrail(Generic[TContext]):
7398
"""Input guardrails are checks that run in parallel to the agent's execution.
@@ -82,10 +107,7 @@ class InputGuardrail(Generic[TContext]):
82107
execution will immediately stop and a `InputGuardrailTripwireTriggered` exception will be raised
83108
"""
84109

85-
guardrail_function: Callable[
86-
[RunContextWrapper[TContext], Agent[Any], str | list[TResponseInputItem]],
87-
MaybeAwaitable[GuardrailFunctionOutput],
88-
]
110+
guardrail_function: InputGuardrailFunction[TContext] | InputGuardrailFunctionLegacy[TContext]
89111
"""A function that receives the agent input and the context, and returns a
90112
`GuardrailResult`. The result marks whether the tripwire was triggered, and can optionally
91113
include information about the guardrail's output.
@@ -107,11 +129,21 @@ async def run(
107129
agent: Agent[Any],
108130
input: str | list[TResponseInputItem],
109131
context: RunContextWrapper[TContext],
132+
previous_response_id: str | None,
110133
) -> InputGuardrailResult:
111134
if not callable(self.guardrail_function):
112135
raise UserError(f"Guardrail function must be callable, got {self.guardrail_function}")
113136

114-
output = self.guardrail_function(context, agent, input)
137+
sig = inspect.signature(self.guardrail_function)
138+
if len(sig.parameters) == 3:
139+
# Legacy guardrail function
140+
legacy_function = cast(InputGuardrailFunctionLegacy[TContext], self.guardrail_function)
141+
output = legacy_function(context, agent, input)
142+
else:
143+
# New guardrail function
144+
new_function = cast(InputGuardrailFunction[TContext], self.guardrail_function)
145+
output = new_function(context, InputGuardailInputs(agent, input, previous_response_id))
146+
115147
if inspect.isawaitable(output):
116148
return InputGuardrailResult(
117149
guardrail=self,
@@ -182,13 +214,11 @@ async def run(
182214
TContext_co = TypeVar("TContext_co", bound=Any, covariant=True)
183215

184216
# For InputGuardrail
185-
_InputGuardrailFuncSync = Callable[
186-
[RunContextWrapper[TContext_co], "Agent[Any]", Union[str, list[TResponseInputItem]]],
187-
GuardrailFunctionOutput,
217+
_InputGuardrailFuncSync = Union[
218+
InputGuardrailFunctionLegacy[TContext_co], InputGuardrailFunction[TContext_co]
188219
]
189-
_InputGuardrailFuncAsync = Callable[
190-
[RunContextWrapper[TContext_co], "Agent[Any]", Union[str, list[TResponseInputItem]]],
191-
Awaitable[GuardrailFunctionOutput],
220+
_InputGuardrailFuncAsync = Union[
221+
InputGuardrailFunctionLegacy[TContext_co], InputGuardrailFunction[TContext_co]
192222
]
193223

194224

‎src/agents/run.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,7 @@ async def run(
221221
starting_agent.input_guardrails
222222
+ (run_config.input_guardrails or []),
223223
copy.deepcopy(input),
224+
previous_response_id,
224225
context_wrapper,
225226
),
226227
cls._run_single_turn(
@@ -446,6 +447,7 @@ async def _run_input_guardrails_with_queue(
446447
agent: Agent[Any],
447448
guardrails: list[InputGuardrail[TContext]],
448449
input: str | list[TResponseInputItem],
450+
previous_response_id: str | None,
449451
context: RunContextWrapper[TContext],
450452
streamed_result: RunResultStreaming,
451453
parent_span: Span[Any],
@@ -455,7 +457,9 @@ async def _run_input_guardrails_with_queue(
455457
# We'll run the guardrails and push them onto the queue as they complete
456458
guardrail_tasks = [
457459
asyncio.create_task(
458-
RunImpl.run_single_input_guardrail(agent, guardrail, input, context)
460+
RunImpl.run_single_input_guardrail(
461+
agent, guardrail, input, previous_response_id, context
462+
)
459463
)
460464
for guardrail in guardrails
461465
]
@@ -551,6 +555,7 @@ async def _run_streamed_impl(
551555
starting_agent,
552556
starting_agent.input_guardrails + (run_config.input_guardrails or []),
553557
copy.deepcopy(ItemHelpers.input_to_new_input_list(starting_input)),
558+
previous_response_id,
554559
context_wrapper,
555560
streamed_result,
556561
current_span,
@@ -825,14 +830,17 @@ async def _run_input_guardrails(
825830
agent: Agent[Any],
826831
guardrails: list[InputGuardrail[TContext]],
827832
input: str | list[TResponseInputItem],
833+
previous_response_id: str | None,
828834
context: RunContextWrapper[TContext],
829835
) -> list[InputGuardrailResult]:
830836
if not guardrails:
831837
return []
832838

833839
guardrail_tasks = [
834840
asyncio.create_task(
835-
RunImpl.run_single_input_guardrail(agent, guardrail, input, context)
841+
RunImpl.run_single_input_guardrail(
842+
agent, guardrail, input, previous_response_id, context
843+
)
836844
)
837845
for guardrail in guardrails
838846
]

‎tests/test_guardrails.py renamed to ‎tests/guardrails/test_guardrails.py

Lines changed: 36 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -32,14 +32,20 @@ def sync_guardrail(
3232
async def test_sync_input_guardrail():
3333
guardrail = InputGuardrail(guardrail_function=get_sync_guardrail(triggers=False))
3434
result = await guardrail.run(
35-
agent=Agent(name="test"), input="test", context=RunContextWrapper(context=None)
35+
agent=Agent(name="test"),
36+
input="test",
37+
context=RunContextWrapper(context=None),
38+
previous_response_id=None,
3639
)
3740
assert not result.output.tripwire_triggered
3841
assert result.output.output_info is None
3942

4043
guardrail = InputGuardrail(guardrail_function=get_sync_guardrail(triggers=True))
4144
result = await guardrail.run(
42-
agent=Agent(name="test"), input="test", context=RunContextWrapper(context=None)
45+
agent=Agent(name="test"),
46+
input="test",
47+
context=RunContextWrapper(context=None),
48+
previous_response_id=None,
4349
)
4450
assert result.output.tripwire_triggered
4551
assert result.output.output_info is None
@@ -48,7 +54,10 @@ async def test_sync_input_guardrail():
4854
guardrail_function=get_sync_guardrail(triggers=True, output_info="test")
4955
)
5056
result = await guardrail.run(
51-
agent=Agent(name="test"), input="test", context=RunContextWrapper(context=None)
57+
agent=Agent(name="test"),
58+
input="test",
59+
context=RunContextWrapper(context=None),
60+
previous_response_id=None,
5261
)
5362
assert result.output.tripwire_triggered
5463
assert result.output.output_info == "test"
@@ -70,14 +79,20 @@ async def async_guardrail(
7079
async def test_async_input_guardrail():
7180
guardrail = InputGuardrail(guardrail_function=get_async_input_guardrail(triggers=False))
7281
result = await guardrail.run(
73-
agent=Agent(name="test"), input="test", context=RunContextWrapper(context=None)
82+
agent=Agent(name="test"),
83+
input="test",
84+
context=RunContextWrapper(context=None),
85+
previous_response_id=None,
7486
)
7587
assert not result.output.tripwire_triggered
7688
assert result.output.output_info is None
7789

7890
guardrail = InputGuardrail(guardrail_function=get_async_input_guardrail(triggers=True))
7991
result = await guardrail.run(
80-
agent=Agent(name="test"), input="test", context=RunContextWrapper(context=None)
92+
agent=Agent(name="test"),
93+
input="test",
94+
context=RunContextWrapper(context=None),
95+
previous_response_id=None,
8196
)
8297
assert result.output.tripwire_triggered
8398
assert result.output.output_info is None
@@ -86,7 +101,10 @@ async def test_async_input_guardrail():
86101
guardrail_function=get_async_input_guardrail(triggers=True, output_info="test")
87102
)
88103
result = await guardrail.run(
89-
agent=Agent(name="test"), input="test", context=RunContextWrapper(context=None)
104+
agent=Agent(name="test"),
105+
input="test",
106+
context=RunContextWrapper(context=None),
107+
previous_response_id=None,
90108
)
91109
assert result.output.tripwire_triggered
92110
assert result.output.output_info == "test"
@@ -98,7 +116,10 @@ async def test_invalid_input_guardrail_raises_user_error():
98116
# Purposely ignoring type error
99117
guardrail = InputGuardrail(guardrail_function="foo") # type: ignore
100118
await guardrail.run(
101-
agent=Agent(name="test"), input="test", context=RunContextWrapper(context=None)
119+
agent=Agent(name="test"),
120+
input="test",
121+
context=RunContextWrapper(context=None),
122+
previous_response_id=None,
102123
)
103124

104125

@@ -210,14 +231,20 @@ def decorated_named_input_guardrail(
210231
async def test_input_guardrail_decorators():
211232
guardrail = decorated_input_guardrail
212233
result = await guardrail.run(
213-
agent=Agent(name="test"), input="test", context=RunContextWrapper(context=None)
234+
agent=Agent(name="test"),
235+
input="test",
236+
previous_response_id=None,
237+
context=RunContextWrapper(context=None),
214238
)
215239
assert not result.output.tripwire_triggered
216240
assert result.output.output_info == "test_1"
217241

218242
guardrail = decorated_named_input_guardrail
219243
result = await guardrail.run(
220-
agent=Agent(name="test"), input="test", context=RunContextWrapper(context=None)
244+
agent=Agent(name="test"),
245+
input="test",
246+
previous_response_id=None,
247+
context=RunContextWrapper(context=None),
221248
)
222249
assert not result.output.tripwire_triggered
223250
assert result.output.output_info == "test_2"

0 commit comments

Comments
 (0)