Coverage for agentos/security/guardrails.py: 52%
154 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 10:59 +0800
1"""
2AgentOS Guardrails — Content Safety & Policy Enforcement Layer
3━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
5Production-grade guardrails system with pluggable rules, LLM-based
6content moderation, and policy enforcement pipeline.
8Architecture:
9 GuardrailsPipeline
10 ├─ InputGuard (validate user input before agent processing)
11 ├─ OutputGuard (validate agent output before returning to user)
12 ├─ ToolGuard (validate tool calls for safety)
13 └─ PolicyEngine (RBAC + content policy rules)
15Key Features:
16 - Pluggable rule engine with hot-reload
17 - LLM-based content moderation (PII/Safety/Toxicity)
18 - Regex-based pattern matching for fast-path checks
19 - Policy violation audit trail
20 - Configurable block/warn/redact actions
21"""
23from __future__ import annotations
25import re
26import json
27from dataclasses import dataclass, field
28from enum import Enum
29from typing import Any, Callable, Dict, List, Optional, Pattern
31# ---------------------------------------------------------------------------
32# Enums & Data Classes
33# ---------------------------------------------------------------------------
36class ViolationSeverity(str, Enum):
37 """Severity level of a guardrail violation."""
38 CRITICAL = "critical" # Immediate block, alert ops
39 HIGH = "high" # Block the request
40 MEDIUM = "medium" # Warn but allow (with redaction)
41 LOW = "low" # Log only
44class GuardAction(str, Enum):
45 """Action to take when a guardrail is triggered."""
46 BLOCK = "block" # Reject the request entirely
47 WARN = "warn" # Allow but flag with warning
48 REDACT = "redact" # Remove offending content, allow rest
49 LOG = "log" # Log only, no user-visible effect
52class Category(str, Enum):
53 """Standard content safety categories."""
54 PII = "pii" # Personally Identifiable Information
55 TOXICITY = "toxicity" # Hate speech, harassment
56 SELF_HARM = "self_harm" # Suicide, self-injury
57 VIOLENCE = "violence" # Graphic violence
58 SEXUAL = "sexual" # Explicit sexual content
59 JAILBREAK = "jailbreak" # Prompt injection / jailbreak attempts
60 DATA_LEAK = "data_leak" # Attempting to leak system prompts / internals
61 MALICIOUS_CODE = "malicious_code" # Code injection, reverse shell, etc.
62 OFF_TOPIC = "off_topic" # Outside defined scope
63 CUSTOM = "custom" # User-defined category
66@dataclass
67class GuardViolation:
68 """A single guardrail violation detected."""
69 category: Category
70 severity: ViolationSeverity
71 action: GuardAction
72 message: str
73 matched_pattern: Optional[str] = None
74 matched_text: Optional[str] = None
75 rule_id: Optional[str] = None
76 metadata: Dict[str, Any] = field(default_factory=dict)
79@dataclass
80class GuardResult:
81 """Result of running guardrails on content."""
82 passed: bool = True
83 violations: List[GuardViolation] = field(default_factory=list)
84 redacted_content: Optional[str] = None
85 warnings: List[str] = field(default_factory=list)
87 @property
88 def blocked(self) -> bool:
89 return any(v.action == GuardAction.BLOCK for v in self.violations)
91 def to_dict(self) -> Dict[str, Any]:
92 return {
93 "passed": self.passed,
94 "blocked": self.blocked,
95 "violations": [
96 {
97 "category": v.category.value,
98 "severity": v.severity.value,
99 "action": v.action.value,
100 "message": v.message,
101 "rule_id": v.rule_id,
102 }
103 for v in self.violations
104 ],
105 "warnings": self.warnings,
106 }
109# ---------------------------------------------------------------------------
110# PII Detection Patterns
111# ---------------------------------------------------------------------------
113PII_PATTERNS: Dict[str, Pattern[str]] = {
114 "email": re.compile(r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}"),
115 "phone_cn": re.compile(r"1[3-9]\d{9}"),
116 "phone_us": re.compile(r"\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}"),
117 "ssn": re.compile(r"\d{3}-\d{2}-\d{4}"),
118 "credit_card": re.compile(r"\b(?:\d{4}[ -]?){3}\d{4}\b"),
119 "ip_address": re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"),
120 "api_key": re.compile(r"(?:api[_-]?key|apikey|token|secret|password)\s*[:=]\s*['\"]?[\w-]{20,}['\"]?", re.IGNORECASE),
121}
124# ---------------------------------------------------------------------------
125# Regex-based Fast-Path Rules
126# ---------------------------------------------------------------------------
128@dataclass
129class RegexRule:
130 """A regex-based guardrail rule for fast-path matching."""
131 rule_id: str
132 category: Category
133 severity: ViolationSeverity
134 action: GuardAction
135 pattern: Pattern[str]
136 message: str
139DEFAULT_RULES: List[RegexRule] = [
140 # PII Rules
141 RegexRule("pii-email", Category.PII, ViolationSeverity.HIGH, GuardAction.REDACT,
142 PII_PATTERNS["email"], "Email address detected"),
143 RegexRule("pii-phone-cn", Category.PII, ViolationSeverity.MEDIUM, GuardAction.REDACT,
144 PII_PATTERNS["phone_cn"], "Chinese phone number detected"),
145 RegexRule("pii-ssn", Category.PII, ViolationSeverity.CRITICAL, GuardAction.BLOCK,
146 PII_PATTERNS["ssn"], "SSN detected"),
147 RegexRule("pii-cc", Category.PII, ViolationSeverity.CRITICAL, GuardAction.BLOCK,
148 PII_PATTERNS["credit_card"], "Credit card number detected"),
149 RegexRule("pii-apikey", Category.PII, ViolationSeverity.CRITICAL, GuardAction.BLOCK,
150 PII_PATTERNS["api_key"], "Potential API key in text"),
152 # Jailbreak patterns
153 RegexRule("jb-ignore", Category.JAILBREAK, ViolationSeverity.CRITICAL, GuardAction.BLOCK,
154 re.compile(r"(?:ignore|forget|disregard)\s+(?:all\s+)?(?:previous|above|prior)\s+(?:instructions?|prompts?|rules?)", re.IGNORECASE),
155 "Jailbreak attempt: ignore instructions"),
156 RegexRule("jb-dan", Category.JAILBREAK, ViolationSeverity.CRITICAL, GuardAction.BLOCK,
157 re.compile(r"\bDAN\s*(?:mode|jailbreak)?\b", re.IGNORECASE),
158 "Jailbreak attempt: DAN mode"),
159 RegexRule("jb-roleplay", Category.JAILBREAK, ViolationSeverity.HIGH, GuardAction.BLOCK,
160 re.compile(r"(?:pretend|act\s+as\s+if|imagine)\s+you\s+(?:are|were)\s+(?:an?\s+)?(?:unfiltered|unrestricted|evil|dark|malicious)", re.IGNORECASE),
161 "Jailbreak attempt: roleplay escalation"),
163 # Malicious code
164 RegexRule("mc-reverse-shell", Category.MALICIOUS_CODE, ViolationSeverity.CRITICAL, GuardAction.BLOCK,
165 re.compile(r"(?:bash|sh|nc|netcat|ncat)\s+.*(?:>&?\s*/dev/(?:tcp|udp)|-e\s+/bin/(?:bash|sh))", re.IGNORECASE),
166 "Reverse shell attempt detected"),
167 RegexRule("mc-rm-rf", Category.MALICIOUS_CODE, ViolationSeverity.HIGH, GuardAction.BLOCK,
168 re.compile(r"(?:rm\s+-rf|del\s+/[fsq])\s+(?:/|~|\*)", re.IGNORECASE),
169 "Destructive file operation detected"),
170]
173# ---------------------------------------------------------------------------
174# Guardrail Engine
175# ---------------------------------------------------------------------------
177class RegexGuard:
178 """Fast-path regex-based guard for common patterns."""
180 def __init__(self, rules: Optional[List[RegexRule]] = None):
181 self._rules: Dict[str, RegexRule] = {}
182 for rule in (rules or DEFAULT_RULES):
183 self._rules[rule.rule_id] = rule
185 def add_rule(self, rule: RegexRule) -> None:
186 self._rules[rule.rule_id] = rule
188 def remove_rule(self, rule_id: str) -> None:
189 self._rules.pop(rule_id, None)
191 def scan(self, content: str) -> List[GuardViolation]:
192 violations: List[GuardViolation] = []
193 for rule in self._rules.values():
194 for match in rule.pattern.finditer(content):
195 violations.append(GuardViolation(
196 category=rule.category,
197 severity=rule.severity,
198 action=rule.action,
199 message=rule.message,
200 matched_pattern=rule.pattern.pattern,
201 matched_text=match.group(),
202 rule_id=rule.rule_id,
203 ))
204 return violations
206 def redact(self, content: str, violations: List[GuardViolation]) -> str:
207 """Redact PII from content based on matched violations."""
208 result = content
209 for v in violations:
210 if v.action == GuardAction.REDACT and v.matched_text:
211 category = v.category.value
212 replacement = f"[REDACTED_{category.upper()}]"
213 result = result.replace(v.matched_text, replacement)
214 return result
217class ContentSafetyGuard:
218 """
219 LLM-based content safety assessment.
221 Uses a lightweight classifier prompt to assess content safety
222 across multiple categories. Can be backed by any LLM provider.
223 """
225 SAFETY_PROMPT = """Assess the following content for safety violations.
226Categories: toxicity, self_harm, violence, sexual, jailbreak, data_leak.
228Respond with ONLY a JSON object:
229{
230 "safe": true/false,
231 "categories": [{"category": "<name>", "severity": "low|medium|high|critical", "reason": "<brief>"}]
232}
234Content to assess:
235---
236{content}
237---"""
239 def __init__(self, llm_call: Optional[Callable] = None):
240 self._llm_call = llm_call
242 async def assess(self, content: str) -> List[GuardViolation]:
243 if self._llm_call is None:
244 return [] # No LLM backend configured, skip
246 prompt = self.SAFETY_PROMPT.format(content=content[:4000])
247 try:
248 response = await self._llm_call(prompt)
249 result = json.loads(response)
250 except Exception:
251 return []
253 if result.get("safe", True):
254 return []
256 violations = []
257 severity_map = {
258 "low": ViolationSeverity.LOW,
259 "medium": ViolationSeverity.MEDIUM,
260 "high": ViolationSeverity.HIGH,
261 "critical": ViolationSeverity.CRITICAL,
262 }
263 for cat in result.get("categories", []):
264 cat_name = cat.get("category", "custom")
265 try:
266 cat_enum = Category(cat_name)
267 except ValueError:
268 cat_enum = Category.CUSTOM
270 violations.append(GuardViolation(
271 category=cat_enum,
272 severity=severity_map.get(cat.get("severity", "medium"), ViolationSeverity.MEDIUM),
273 action=GuardAction.BLOCK,
274 message=cat.get("reason", f"Content safety violation: {cat_name}"),
275 metadata={"llm_assessment": cat},
276 ))
278 return violations
281# ---------------------------------------------------------------------------
282# Guardrails Pipeline
283# ---------------------------------------------------------------------------
285class GuardrailsPipeline:
286 """
287 Production guardrails pipeline combining regex fast-path and LLM-based
288 content safety assessment.
290 Usage:
291 pipeline = GuardrailsPipeline()
292 pipeline.add_regex_rule(...)
294 # Input validation
295 result = await pipeline.check_input(user_message)
296 if not result.passed:
297 raise GuardViolationError(result)
299 # Output validation
300 result = await pipeline.check_output(agent_response)
301 """
303 def __init__(
304 self,
305 regex_guard: Optional[RegexGuard] = None,
306 safety_guard: Optional[ContentSafetyGuard] = None,
307 enable_regex: bool = True,
308 enable_safety: bool = True,
309 ):
310 self._regex = regex_guard or RegexGuard()
311 self._safety = safety_guard or ContentSafetyGuard()
312 self._enable_regex = enable_regex
313 self._enable_safety = enable_safety
314 self._audit_log: List[GuardResult] = []
316 def add_regex_rule(self, rule: RegexRule) -> None:
317 self._regex.add_rule(rule)
319 def remove_regex_rule(self, rule_id: str) -> None:
320 self._regex.remove_rule(rule_id)
322 async def check_input(self, content: str) -> GuardResult:
323 """Validate user input before agent processing."""
324 return await self._check(content, stage="input")
326 async def check_output(self, content: str) -> GuardResult:
327 """Validate agent output before returning to user."""
328 return await self._check(content, stage="output")
330 async def check_tool_call(self, tool_name: str, arguments: Dict[str, Any]) -> GuardResult:
331 """Validate tool calls for safety."""
332 content = f"Tool: {tool_name}\nArgs: {json.dumps(arguments)}"
333 return await self._check(content, stage="tool_call")
335 async def _check(self, content: str, stage: str = "unknown") -> GuardResult:
336 violations: List[GuardViolation] = []
338 # Fast-path: regex scanning
339 if self._enable_regex:
340 violations.extend(self._regex.scan(content))
342 # Deep check: LLM safety assessment
343 if self._enable_safety and content.strip():
344 safety_violations = await self._safety.assess(content)
345 violations.extend(safety_violations)
347 # Determine result
348 if not violations:
349 result = GuardResult(passed=True)
350 else:
351 redacted = self._regex.redact(content, violations) if any(
352 v.action == GuardAction.REDACT for v in violations
353 ) else None
355 result = GuardResult(
356 passed=not any(v.action == GuardAction.BLOCK for v in violations),
357 violations=violations,
358 redacted_content=redacted,
359 warnings=[v.message for v in violations if v.action == GuardAction.WARN],
360 )
362 self._audit_log.append(result)
363 return result
365 def get_audit_log(self) -> List[Dict[str, Any]]:
366 return [r.to_dict() for r in self._audit_log]
368 def get_statistics(self) -> Dict[str, int]:
369 total = len(self._audit_log)
370 blocked = sum(1 for r in self._audit_log if r.blocked)
371 passed = sum(1 for r in self._audit_log if r.passed and not r.violations)
372 warned = total - blocked - passed
373 return {
374 "total_checks": total,
375 "passed": passed,
376 "blocked": blocked,
377 "warned": warned,
378 }
381# ---------------------------------------------------------------------------
382# Exception
383# ---------------------------------------------------------------------------
385class GuardViolationError(Exception):
386 """Raised when guardrails block a request."""
388 def __init__(self, result: GuardResult):
389 self.result = result
390 violations_summary = "; ".join(
391 f"[{v.category.value}] {v.message}" for v in result.violations
392 )
393 super().__init__(f"Guardrail blocked: {violations_summary}")
396# ---------------------------------------------------------------------------
397# Convenience: Pre-built Pipeline
398# ---------------------------------------------------------------------------
400def create_default_pipeline() -> GuardrailsPipeline:
401 """Create a GuardrailsPipeline with sensible defaults."""
402 return GuardrailsPipeline(
403 regex_guard=RegexGuard(rules=DEFAULT_RULES),
404 enable_regex=True,
405 enable_safety=False, # LLM-based safety off by default; opt-in
406 )
409def create_strict_pipeline() -> GuardrailsPipeline:
410 """Create a GuardrailsPipeline with strict rules + LLM safety."""
411 return GuardrailsPipeline(
412 regex_guard=RegexGuard(rules=DEFAULT_RULES),
413 safety_guard=ContentSafetyGuard(),
414 enable_regex=True,
415 enable_safety=True,
416 )