Skip to content

Commit eb34e26

Browse files
michi-okahatadot-agi
michi-okahata
and
dot-agi
authored
input/output guardrail decorator (#1039)
* input/output guardrail decorator #1003 * changed span_kinds name, confirmed no existing semconv for guardrails * ruff fixes * ruff fix 2 * combined input/output into a guardrail decorator with parameter 'spec' * quick fix to comment * fix init * guardrail examples * add guardrail docs * ruff fix * pc fix * fix agentops init in example notebook * ruff again --------- Co-authored-by: Pratyush Shukla <ps4534@nyu.edu>
1 parent f3fd730 commit eb34e26

File tree

13 files changed

+315
-25
lines changed

13 files changed

+315
-25
lines changed

‎agentops/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
from typing import List, Optional, Union, Dict, Any
1616
from agentops.client import Client
1717
from agentops.sdk.core import TraceContext, tracer
18-
from agentops.sdk.decorators import trace, session, agent, task, workflow, operation, tool
18+
from agentops.sdk.decorators import trace, session, agent, task, workflow, operation, tool, guardrail
1919
from agentops.enums import TraceState, SUCCESS, ERROR, UNSET
2020
from opentelemetry.trace.status import StatusCode
2121

@@ -265,6 +265,7 @@ def end_trace(
265265
"task",
266266
"workflow",
267267
"operation",
268+
"guardrail",
268269
"tracer",
269270
"tool",
270271
# Trace state enums

‎agentops/sdk/decorators/__init__.py‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
trace = create_entity_decorator(SpanKind.SESSION)
1919
tool = create_entity_decorator(SpanKind.TOOL)
2020
operation = task
21+
guardrail = create_entity_decorator(SpanKind.GUARDRAIL)
2122

2223

2324
# For backward compatibility: @session decorator calls @trace decorator
@@ -37,4 +38,13 @@ def session(*args, **kwargs): # noqa: F811
3738
# For now, keeping the alias as it was, assuming it was intentional for `operation` to be `task`.
3839
operation = task
3940

40-
__all__ = ["agent", "task", "workflow", "trace", "session", "operation", "tool"]
41+
__all__ = [
42+
"agent",
43+
"task",
44+
"workflow",
45+
"trace",
46+
"session",
47+
"operation",
48+
"tool",
49+
"guardrail",
50+
]

‎agentops/sdk/decorators/factory.py‎

Lines changed: 24 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -33,9 +33,10 @@ def decorator(
3333
version: Optional[Any] = None,
3434
tags: Optional[Union[list, dict]] = None,
3535
cost=None,
36+
spec=None,
3637
) -> Callable[..., Any]:
3738
if wrapped is None:
38-
return functools.partial(decorator, name=name, version=version, tags=tags, cost=cost)
39+
return functools.partial(decorator, name=name, version=version, tags=tags, cost=cost, spec=spec)
3940

4041
if inspect.isclass(wrapped):
4142
# Class decoration wraps __init__ and aenter/aexit for context management.
@@ -168,10 +169,13 @@ async def _wrapped_session_async() -> Any:
168169
attributes={CoreAttributes.TAGS: tags} if tags else None,
169170
)
170171
try:
171-
_record_entity_input(span, args, kwargs)
172+
_record_entity_input(span, args, kwargs, entity_kind=entity_kind)
172173
# Set cost attribute if tool
173174
if entity_kind == "tool" and cost is not None:
174175
span.set_attribute(SpanAttributes.LLM_USAGE_TOOL_COST, cost)
176+
# Set spec attribute if guardrail
177+
if entity_kind == "guardrail" and (spec == "input" or spec == "output"):
178+
span.set_attribute(SpanAttributes.AGENTOPS_DECORATOR_SPEC.format(entity_kind=entity_kind), spec)
175179
except Exception as e:
176180
logger.warning(f"Input recording failed for '{operation_name}': {e}")
177181
result = wrapped_func(*args, **kwargs)
@@ -184,10 +188,13 @@ async def _wrapped_session_async() -> Any:
184188
attributes={CoreAttributes.TAGS: tags} if tags else None,
185189
)
186190
try:
187-
_record_entity_input(span, args, kwargs)
191+
_record_entity_input(span, args, kwargs, entity_kind=entity_kind)
188192
# Set cost attribute if tool
189193
if entity_kind == "tool" and cost is not None:
190194
span.set_attribute(SpanAttributes.LLM_USAGE_TOOL_COST, cost)
195+
# Set spec attribute if guardrail
196+
if entity_kind == "guardrail" and (spec == "input" or spec == "output"):
197+
span.set_attribute(SpanAttributes.AGENTOPS_DECORATOR_SPEC.format(entity_kind=entity_kind), spec)
191198
except Exception as e:
192199
logger.warning(f"Input recording failed for '{operation_name}': {e}")
193200
result = wrapped_func(*args, **kwargs)
@@ -202,16 +209,21 @@ async def _wrapped_async() -> Any:
202209
attributes={CoreAttributes.TAGS: tags} if tags else None,
203210
) as span:
204211
try:
205-
_record_entity_input(span, args, kwargs)
212+
_record_entity_input(span, args, kwargs, entity_kind=entity_kind)
206213
# Set cost attribute if tool
207214
if entity_kind == "tool" and cost is not None:
208215
span.set_attribute(SpanAttributes.LLM_USAGE_TOOL_COST, cost)
216+
# Set spec attribute if guardrail
217+
if entity_kind == "guardrail" and (spec == "input" or spec == "output"):
218+
span.set_attribute(
219+
SpanAttributes.AGENTOPS_DECORATOR_SPEC.format(entity_kind=entity_kind), spec
220+
)
209221
except Exception as e:
210222
logger.warning(f"Input recording failed for '{operation_name}': {e}")
211223
try:
212224
result = await wrapped_func(*args, **kwargs)
213225
try:
214-
_record_entity_output(span, result)
226+
_record_entity_output(span, result, entity_kind=entity_kind)
215227
except Exception as e:
216228
logger.warning(f"Output recording failed for '{operation_name}': {e}")
217229
return result
@@ -229,16 +241,21 @@ async def _wrapped_async() -> Any:
229241
attributes={CoreAttributes.TAGS: tags} if tags else None,
230242
) as span:
231243
try:
232-
_record_entity_input(span, args, kwargs)
244+
_record_entity_input(span, args, kwargs, entity_kind=entity_kind)
233245
# Set cost attribute if tool
234246
if entity_kind == "tool" and cost is not None:
235247
span.set_attribute(SpanAttributes.LLM_USAGE_TOOL_COST, cost)
248+
# Set spec attribute if guardrail
249+
if entity_kind == "guardrail" and (spec == "input" or spec == "output"):
250+
span.set_attribute(
251+
SpanAttributes.AGENTOPS_DECORATOR_SPEC.format(entity_kind=entity_kind), spec
252+
)
236253
except Exception as e:
237254
logger.warning(f"Input recording failed for '{operation_name}': {e}")
238255
try:
239256
result = wrapped_func(*args, **kwargs)
240257
try:
241-
_record_entity_output(span, result)
258+
_record_entity_output(span, result, entity_kind=entity_kind)
242259
except Exception as e:
243260
logger.warning(f"Output recording failed for '{operation_name}': {e}")
244261
return result

‎agentops/sdk/decorators/utility.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -135,27 +135,27 @@ def _create_as_current_span(
135135
logger.debug(f"[DEBUG] AFTER {operation_name}.{span_kind} - Returned to context: {after_span}")
136136

137137

138-
def _record_entity_input(span: trace.Span, args: tuple, kwargs: Dict[str, Any]) -> None:
138+
def _record_entity_input(span: trace.Span, args: tuple, kwargs: Dict[str, Any], entity_kind: str = "entity") -> None:
139139
"""Record operation input parameters to span if content tracing is enabled"""
140140
try:
141141
input_data = {"args": args, "kwargs": kwargs}
142142
json_data = safe_serialize(input_data)
143143

144144
if _check_content_size(json_data):
145-
span.set_attribute(SpanAttributes.AGENTOPS_ENTITY_INPUT, json_data)
145+
span.set_attribute(SpanAttributes.AGENTOPS_DECORATOR_INPUT.format(entity_kind=entity_kind), json_data)
146146
else:
147147
logger.debug("Operation input exceeds size limit, not recording")
148148
except Exception as err:
149149
logger.warning(f"Failed to serialize operation input: {err}")
150150

151151

152-
def _record_entity_output(span: trace.Span, result: Any) -> None:
152+
def _record_entity_output(span: trace.Span, result: Any, entity_kind: str = "entity") -> None:
153153
"""Record operation output value to span if content tracing is enabled"""
154154
try:
155155
json_data = safe_serialize(result)
156156

157157
if _check_content_size(json_data):
158-
span.set_attribute(SpanAttributes.AGENTOPS_ENTITY_OUTPUT, json_data)
158+
span.set_attribute(SpanAttributes.AGENTOPS_DECORATOR_OUTPUT.format(entity_kind=entity_kind), json_data)
159159
else:
160160
logger.debug("Operation output exceeds size limit, not recording")
161161
except Exception as err:

‎agentops/semconv/span_attributes.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,9 @@ class SpanAttributes:
8888
AGENTOPS_ENTITY_INPUT = "agentops.entity.input"
8989
AGENTOPS_SPAN_KIND = "agentops.span.kind"
9090
AGENTOPS_ENTITY_NAME = "agentops.entity.name"
91+
AGENTOPS_DECORATOR_SPEC = "agentops.{entity_kind}.spec"
92+
AGENTOPS_DECORATOR_INPUT = "agentops.{entity_kind}.input"
93+
AGENTOPS_DECORATOR_OUTPUT = "agentops.{entity_kind}.output"
9194

9295
# Operation attributes
9396
OPERATION_NAME = "operation.name"

‎agentops/semconv/span_kinds.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ class SpanKind:
2727
UNKNOWN = "unknown"
2828
CHAIN = "chain"
2929
TEXT = "text"
30+
GUARDRAIL = "guardrail"
3031

3132

3233
class AgentOpsSpanKindValues(Enum):

‎docs/v2/concepts/decorators.mdx‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ AgentOps provides the following decorators:
1414
| `@workflow` | Track a sequence of operations | WORKFLOW span |
1515
| `@task` | Track smaller units of work (similar to operations) | TASK span |
1616
| `@tool` | Track tool usage and cost in agent operations | TOOL span |
17+
| `@guardrail` | Track guardrail input and output | GUARDRAIL span |
1718

1819
## Decorator Hierarchy
1920

@@ -235,6 +236,27 @@ The tool decorator provides:
235236
- Support for all function types (sync, async, generator, async generator)
236237
- Cost accumulation in generator and async generator operations
237238

239+
### @guardrail
240+
241+
The `@guardrail` decorator tracks guardrail input and output. You can specify the guardrail type (`"input"` or `"output"`) with the `spec` parameter.
242+
243+
```python
244+
from agentops.sdk.decorators import guardrail
245+
import agentops
246+
import re
247+
248+
# Initialize AgentOps
249+
agentops.init(api_key="YOUR_API_KEY")
250+
251+
@guardrail(spec="input")
252+
def secret_key_guardrail(input):
253+
pattern = r'\bsk-[a-zA-Z0-9]{10,}\b'
254+
result = True if re.search(pattern, input) else False
255+
return {
256+
"tripwire_triggered" : result
257+
}
258+
```
259+
238260
## Decorator Attributes
239261

240262
You can pass additional attributes to decorators:

‎examples/google_adk/human_approval.ipynb‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,7 @@
183183
" \"\"\"\n",
184184
" Prompts for human approval and returns the decision as a JSON string.\n",
185185
" \"\"\"\n",
186-
" print(f\"🔔 HUMAN APPROVAL REQUIRED:\")\n",
186+
" print(\"🔔 HUMAN APPROVAL REQUIRED:\")\n",
187187
" print(f\" Amount: ${amount:,.2f}\")\n",
188188
" print(f\" Reason: {reason}\")\n",
189189
" decision = \"\"\n",

‎examples/openai_agents/agent_guardrails.ipynb‎

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"id": "f68ce4af",
6+
"metadata": {},
7+
"source": [
8+
"# OpenAI Agents Guardrails Demonstration\n",
9+
"\n",
10+
"This notebook demonstrates guardrails using the Agents SDK and how one can observe them using the AgentOps platform."
11+
]
12+
},
13+
{
14+
"cell_type": "code",
15+
"execution_count": null,
16+
"id": "10bcf29b",
17+
"metadata": {},
18+
"outputs": [],
19+
"source": [
20+
"# Install required packages\n",
21+
"%pip install agentops\n",
22+
"%pip install openai-agents\n",
23+
"%pip install dotenv pydantic"
24+
]
25+
},
26+
{
27+
"cell_type": "code",
28+
"execution_count": null,
29+
"id": "a3be4e68",
30+
"metadata": {},
31+
"outputs": [],
32+
"source": [
33+
"# Import dependencies\n",
34+
"from pydantic import BaseModel\n",
35+
"from agents import (\n",
36+
" Agent,\n",
37+
" GuardrailFunctionOutput,\n",
38+
" InputGuardrailTripwireTriggered,\n",
39+
" RunContextWrapper,\n",
40+
" Runner,\n",
41+
" TResponseInputItem,\n",
42+
" input_guardrail,\n",
43+
")"
44+
]
45+
},
46+
{
47+
"cell_type": "code",
48+
"execution_count": null,
49+
"id": "2d0dddb6",
50+
"metadata": {},
51+
"outputs": [],
52+
"source": [
53+
"# Load API keys\n",
54+
"import os\n",
55+
"from dotenv import load_dotenv\n",
56+
"\n",
57+
"load_dotenv()\n",
58+
"\n",
59+
"os.environ[\"AGENTOPS_API_KEY\"] = os.getenv(\"AGENTOPS_API_KEY\", \"your_api_key_here\")\n",
60+
"os.environ[\"OPENAI_API_KEY\"] = os.getenv(\"OPENAI_API_KEY\", \"your_openai_api_key_here\")"
61+
]
62+
},
63+
{
64+
"cell_type": "code",
65+
"execution_count": null,
66+
"id": "114e216b",
67+
"metadata": {},
68+
"outputs": [],
69+
"source": [
70+
"# Initialize agentops and import the guardrail decorator\n",
71+
"import agentops\n",
72+
"from agentops import guardrail\n",
73+
"\n",
74+
"agentops.init(api_key=os.environ[\"AGENTOPS_API_KEY\"], tags=[\"agentops-example\"], auto_start_session=False)\n",
75+
"tracer = agentops.start_trace(trace_name=\"OpenAI Agents Guardrail Example\")"
76+
]
77+
},
78+
{
79+
"cell_type": "code",
80+
"execution_count": null,
81+
"id": "0bf8b54d",
82+
"metadata": {},
83+
"outputs": [],
84+
"source": [
85+
"# OpenAI Agents SDK guardrail example with agentops guardrails decorator for observability\n",
86+
"class MathHomeworkOutput(BaseModel):\n",
87+
" is_math_homework: bool\n",
88+
" reasoning: str\n",
89+
"\n",
90+
"\n",
91+
"guardrail_agent = Agent(\n",
92+
" name=\"Guardrail check\",\n",
93+
" instructions=\"Check if the user is asking you to do their math homework.\",\n",
94+
" output_type=MathHomeworkOutput,\n",
95+
")\n",
96+
"\n",
97+
"\n",
98+
"@input_guardrail\n",
99+
"@guardrail(spec=\"input\") # Specify guardrail type as input or output\n",
100+
"async def math_guardrail(\n",
101+
" ctx: RunContextWrapper[None], agent: Agent, input: str | list[TResponseInputItem]\n",
102+
") -> GuardrailFunctionOutput:\n",
103+
" result = await Runner.run(guardrail_agent, input, context=ctx.context)\n",
104+
"\n",
105+
" return GuardrailFunctionOutput(\n",
106+
" output_info=result.final_output,\n",
107+
" tripwire_triggered=result.final_output.is_math_homework,\n",
108+
" )\n",
109+
"\n",
110+
"\n",
111+
"agent = Agent(\n",
112+
" name=\"Customer support agent\",\n",
113+
" instructions=\"You are a customer support agent. You help customers with their questions.\",\n",
114+
" input_guardrails=[math_guardrail],\n",
115+
")\n",
116+
"\n",
117+
"\n",
118+
"async def main():\n",
119+
" # This should trip the guardrail\n",
120+
" try:\n",
121+
" await Runner.run(agent, \"Hello, can you help me solve for x: 2x + 3 = 11?\")\n",
122+
" print(\"Guardrail didn't trip - this is unexpected\")\n",
123+
"\n",
124+
" except InputGuardrailTripwireTriggered:\n",
125+
" print(\"Math homework guardrail tripped\")\n",
126+
"\n",
127+
"\n",
128+
"await main()"
129+
]
130+
},
131+
{
132+
"cell_type": "code",
133+
"execution_count": null,
134+
"id": "63bf8e09",
135+
"metadata": {},
136+
"outputs": [],
137+
"source": [
138+
"agentops.end_trace(tracer, end_state=\"Success\")"
139+
]
140+
}
141+
],
142+
"metadata": {
143+
"kernelspec": {
144+
"display_name": "agentops (3.11.11)",
145+
"language": "python",
146+
"name": "python3"
147+
},
148+
"language_info": {
149+
"codemirror_mode": {
150+
"name": "ipython",
151+
"version": 3
152+
},
153+
"file_extension": ".py",
154+
"mimetype": "text/x-python",
155+
"name": "python",
156+
"nbconvert_exporter": "python",
157+
"pygments_lexer": "ipython3",
158+
"version": "3.11.11"
159+
}
160+
},
161+
"nbformat": 4,
162+
"nbformat_minor": 5
163+
}

0 commit comments

Comments
 (0)