-
Notifications
You must be signed in to change notification settings - Fork 4.3k
Expand file tree
/
Copy pathtool_guardrails.py
More file actions
181 lines (147 loc) · 6.45 KB
/
Copy pathtool_guardrails.py
File metadata and controls
181 lines (147 loc) · 6.45 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
import asyncio
import json
from agents import (
Agent,
Runner,
ToolGuardrailFunctionOutput,
ToolInputGuardrailData,
ToolOutputGuardrailData,
ToolOutputGuardrailTripwireTriggered,
function_tool,
tool_input_guardrail,
tool_output_guardrail,
)
@function_tool
def send_email(to: str, subject: str, body: str) -> str:
"""Send an email to the specified recipient."""
return f"Email sent to {to} with subject '{subject}'"
@function_tool
def get_user_data(user_id: str) -> dict[str, str]:
"""Get user data by ID."""
# Simulate returning sensitive data
return {
"user_id": user_id,
"name": "John Doe",
"email": "[email protected]",
"ssn": "123-45-6789", # Sensitive data that should be blocked!
"phone": "555-1234",
}
@function_tool
def get_contact_info(user_id: str) -> dict[str, str]:
"""Get contact info by ID."""
return {
"user_id": user_id,
"name": "Jane Smith",
"email": "[email protected]",
"phone": "555-1234",
}
@tool_input_guardrail
def reject_sensitive_words(data: ToolInputGuardrailData) -> ToolGuardrailFunctionOutput:
"""Reject tool calls that contain sensitive words in arguments."""
try:
args = json.loads(data.context.tool_arguments) if data.context.tool_arguments else {}
except json.JSONDecodeError:
return ToolGuardrailFunctionOutput(output_info="Invalid JSON arguments")
# Check for suspicious content
sensitive_words = [
"password",
"hack",
"exploit",
"malware",
"ACME",
]
for key, value in args.items():
value_str = str(value).lower()
for word in sensitive_words:
if word.lower() in value_str:
# Reject tool call and inform the model the function was not called
return ToolGuardrailFunctionOutput.reject_content(
message=f"🚨 Tool call blocked: contains '{word}'",
output_info={"blocked_word": word, "argument": key},
)
return ToolGuardrailFunctionOutput(output_info="Input validated")
@tool_output_guardrail
def block_sensitive_output(data: ToolOutputGuardrailData) -> ToolGuardrailFunctionOutput:
"""Block tool outputs that contain sensitive data."""
output_str = str(data.output).lower()
# Check for sensitive data patterns
if "ssn" in output_str or "123-45-6789" in output_str:
# Use raise_exception to halt execution completely for sensitive data
return ToolGuardrailFunctionOutput.raise_exception(
output_info={"blocked_pattern": "SSN", "tool": data.context.tool_name},
)
return ToolGuardrailFunctionOutput(output_info="Output validated")
@tool_output_guardrail
def reject_phone_numbers(data: ToolOutputGuardrailData) -> ToolGuardrailFunctionOutput:
"""Reject function output containing phone numbers."""
output_str = str(data.output)
if "555-1234" in output_str:
return ToolGuardrailFunctionOutput.reject_content(
message="User data not retrieved as it contains a phone number which is restricted.",
output_info={"redacted": "phone_number"},
)
return ToolGuardrailFunctionOutput(output_info="Phone number check passed")
# Apply guardrails to tools
send_email.tool_input_guardrails = [reject_sensitive_words]
get_user_data.tool_output_guardrails = [block_sensitive_output]
get_contact_info.tool_output_guardrails = [reject_phone_numbers]
agent = Agent(
name="Secure Assistant",
instructions=(
"You are a helpful assistant with access to email and user data tools. "
"When the user provides all required arguments for a requested tool, call it instead of "
"asking a follow-up question."
),
tools=[send_email, get_user_data, get_contact_info],
)
async def main():
print("=== Tool Guardrails Example ===\n")
try:
# Example 1: Normal operation - should work fine
print("1. Normal email sending:")
result = await Runner.run(
agent,
"Send an email to [email protected] with subject 'Welcome' and body "
"'Welcome to our service.'",
)
print(f"✅ Successful tool execution: {result.final_output}\n")
# Example 2: Input guardrail triggers - function tool call is rejected but execution continues
print("2. Attempting to send email with suspicious content:")
result = await Runner.run(
agent,
"Send an email to [email protected] with subject 'Introduction' and body "
"'Introducing ACME corp.'",
)
print(f"❌ Guardrail rejected function tool call: {result.final_output}\n")
except Exception as e:
print(f"Error: {e}\n")
try:
# Example 3: Output guardrail triggers - should raise exception for sensitive data
print("3. Attempting to get user data (contains SSN). Execution blocked:")
result = await Runner.run(agent, "Get the data for user ID user123")
print(f"✅ Successful tool execution: {result.final_output}\n")
except ToolOutputGuardrailTripwireTriggered as e:
print("🚨 Output guardrail triggered: Execution halted for sensitive data")
print(f"Details: {e.output.output_info}\n")
try:
# Example 4: Output guardrail triggers - reject returning function tool output but continue execution
print("4. Rejecting function tool output containing phone numbers:")
result = await Runner.run(agent, "Get contact info for user456")
print(f"❌ Guardrail rejected function tool output: {result.final_output}\n")
except Exception as e:
print(f"Error: {e}\n")
if __name__ == "__main__":
asyncio.run(main())
"""
Example output:
=== Tool Guardrails Example ===
1. Normal email sending:
✅ Successful tool execution: I've sent a welcome email to [email protected] with an appropriate subject and greeting message.
2. Attempting to send email with suspicious content:
❌ Guardrail rejected function tool call: I'm unable to send the email as mentioning ACME Corp. is restricted.
3. Attempting to get user data (contains SSN). Execution blocked:
🚨 Output guardrail triggered: Execution halted for sensitive data
Details: {'blocked_pattern': 'SSN', 'tool': 'get_user_data'}
4. Rejecting function tool output containing sensitive data:
❌ Guardrail rejected function tool output: I'm unable to retrieve the contact info for user456 because it contains restricted information.
"""