Coverage for agentos/security/guardrails.py: 52%
155 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-05 20:52 +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 hashlib
27import json
28from dataclasses import dataclass, field
29from enum import Enum, auto
30from typing import Any, Callable, Dict, List, Optional, Pattern, Set, Union
32# ---------------------------------------------------------------------------
33# Enums & Data Classes
34# ---------------------------------------------------------------------------
37class ViolationSeverity(str, Enum):
38 """Severity level of a guardrail violation."""
39 CRITICAL = "critical" # Immediate block, alert ops
40 HIGH = "high" # Block the request
41 MEDIUM = "medium" # Warn but allow (with redaction)
42 LOW = "low" # Log only
45class GuardAction(str, Enum):
46 """Action to take when a guardrail is triggered."""
47 BLOCK = "block" # Reject the request entirely
48 WARN = "warn" # Allow but flag with warning
49 REDACT = "redact" # Remove offending content, allow rest
50 LOG = "log" # Log only, no user-visible effect
53class Category(str, Enum):
54 """Standard content safety categories."""
55 PII = "pii" # Personally Identifiable Information
56 TOXICITY = "toxicity" # Hate speech, harassment
57 SELF_HARM = "self_harm" # Suicide, self-injury
58 VIOLENCE = "violence" # Graphic violence
59 SEXUAL = "sexual" # Explicit sexual content
60 JAILBREAK = "jailbreak" # Prompt injection / jailbreak attempts
61 DATA_LEAK = "data_leak" # Attempting to leak system prompts / internals
62 MALICIOUS_CODE = "malicious_code" # Code injection, reverse shell, etc.
63 OFF_TOPIC = "off_topic" # Outside defined scope
64 CUSTOM = "custom" # User-defined category
67@dataclass
68class GuardViolation:
69 """A single guardrail violation detected."""
70 category: Category
71 severity: ViolationSeverity
72 action: GuardAction
73 message: str
74 matched_pattern: Optional[str] = None
75 matched_text: Optional[str] = None
76 rule_id: Optional[str] = None
77 metadata: Dict[str, Any] = field(default_factory=dict)
80@dataclass
81class GuardResult:
82 """Result of running guardrails on content."""
83 passed: bool = True
84 violations: List[GuardViolation] = field(default_factory=list)
85 redacted_content: Optional[str] = None
86 warnings: List[str] = field(default_factory=list)
88 @property
89 def blocked(self) -> bool:
90 return any(v.action == GuardAction.BLOCK for v in self.violations)
92 def to_dict(self) -> Dict[str, Any]:
93 return {
94 "passed": self.passed,
95 "blocked": self.blocked,
96 "violations": [
97 {
98 "category": v.category.value,
99 "severity": v.severity.value,
100 "action": v.action.value,
101 "message": v.message,
102 "rule_id": v.rule_id,
103 }
104 for v in self.violations
105 ],
106 "warnings": self.warnings,
107 }
110# ---------------------------------------------------------------------------
111# PII Detection Patterns
112# ---------------------------------------------------------------------------
114PII_PATTERNS: Dict[str, Pattern[str]] = {
115 "email": re.compile(r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}"),
116 "phone_cn": re.compile(r"1[3-9]\d{9}"),
117 "phone_us": re.compile(r"\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}"),
118 "ssn": re.compile(r"\d{3}-\d{2}-\d{4}"),
119 "credit_card": re.compile(r"\b(?:\d{4}[ -]?){3}\d{4}\b"),
120 "ip_address": re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"),
121 "api_key": re.compile(r"(?:api[_-]?key|apikey|token|secret|password)\s*[:=]\s*['\"]?[\w-]{20,}['\"]?", re.IGNORECASE),
122}
125# ---------------------------------------------------------------------------
126# Regex-based Fast-Path Rules
127# ---------------------------------------------------------------------------
129@dataclass
130class RegexRule:
131 """A regex-based guardrail rule for fast-path matching."""
132 rule_id: str
133 category: Category
134 severity: ViolationSeverity
135 action: GuardAction
136 pattern: Pattern[str]
137 message: str
140DEFAULT_RULES: List[RegexRule] = [
141 # PII Rules
142 RegexRule("pii-email", Category.PII, ViolationSeverity.HIGH, GuardAction.REDACT,
143 PII_PATTERNS["email"], "Email address detected"),
144 RegexRule("pii-phone-cn", Category.PII, ViolationSeverity.MEDIUM, GuardAction.REDACT,
145 PII_PATTERNS["phone_cn"], "Chinese phone number detected"),
146 RegexRule("pii-ssn", Category.PII, ViolationSeverity.CRITICAL, GuardAction.BLOCK,
147 PII_PATTERNS["ssn"], "SSN detected"),
148 RegexRule("pii-cc", Category.PII, ViolationSeverity.CRITICAL, GuardAction.BLOCK,
149 PII_PATTERNS["credit_card"], "Credit card number detected"),
150 RegexRule("pii-apikey", Category.PII, ViolationSeverity.CRITICAL, GuardAction.BLOCK,
151 PII_PATTERNS["api_key"], "Potential API key in text"),
153 # Jailbreak patterns
154 RegexRule("jb-ignore", Category.JAILBREAK, ViolationSeverity.CRITICAL, GuardAction.BLOCK,
155 re.compile(r"(?:ignore|forget|disregard)\s+(?:all\s+)?(?:previous|above|prior)\s+(?:instructions?|prompts?|rules?)", re.IGNORECASE),
156 "Jailbreak attempt: ignore instructions"),
157 RegexRule("jb-dan", Category.JAILBREAK, ViolationSeverity.CRITICAL, GuardAction.BLOCK,
158 re.compile(r"\bDAN\s*(?:mode|jailbreak)?\b", re.IGNORECASE),
159 "Jailbreak attempt: DAN mode"),
160 RegexRule("jb-roleplay", Category.JAILBREAK, ViolationSeverity.HIGH, GuardAction.BLOCK,
161 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),
162 "Jailbreak attempt: roleplay escalation"),
164 # Malicious code
165 RegexRule("mc-reverse-shell", Category.MALICIOUS_CODE, ViolationSeverity.CRITICAL, GuardAction.BLOCK,
166 re.compile(r"(?:bash|sh|nc|netcat|ncat)\s+.*(?:>&?\s*/dev/(?:tcp|udp)|-e\s+/bin/(?:bash|sh))", re.IGNORECASE),
167 "Reverse shell attempt detected"),
168 RegexRule("mc-rm-rf", Category.MALICIOUS_CODE, ViolationSeverity.HIGH, GuardAction.BLOCK,
169 re.compile(r"(?:rm\s+-rf|del\s+/[fsq])\s+(?:/|~|\*)", re.IGNORECASE),
170 "Destructive file operation detected"),
171]
174# ---------------------------------------------------------------------------
175# Guardrail Engine
176# ---------------------------------------------------------------------------
178class RegexGuard:
179 """Fast-path regex-based guard for common patterns."""
181 def __init__(self, rules: Optional[List[RegexRule]] = None):
182 self._rules: Dict[str, RegexRule] = {}
183 for rule in (rules or DEFAULT_RULES):
184 self._rules[rule.rule_id] = rule
186 def add_rule(self, rule: RegexRule) -> None:
187 self._rules[rule.rule_id] = rule
189 def remove_rule(self, rule_id: str) -> None:
190 self._rules.pop(rule_id, None)
192 def scan(self, content: str) -> List[GuardViolation]:
193 violations: List[GuardViolation] = []
194 for rule in self._rules.values():
195 for match in rule.pattern.finditer(content):
196 violations.append(GuardViolation(
197 category=rule.category,
198 severity=rule.severity,
199 action=rule.action,
200 message=rule.message,
201 matched_pattern=rule.pattern.pattern,
202 matched_text=match.group(),
203 rule_id=rule.rule_id,
204 ))
205 return violations
207 def redact(self, content: str, violations: List[GuardViolation]) -> str:
208 """Redact PII from content based on matched violations."""
209 result = content
210 for v in violations:
211 if v.action == GuardAction.REDACT and v.matched_text:
212 category = v.category.value
213 replacement = f"[REDACTED_{category.upper()}]"
214 result = result.replace(v.matched_text, replacement)
215 return result
218class ContentSafetyGuard:
219 """
220 LLM-based content safety assessment.
222 Uses a lightweight classifier prompt to assess content safety
223 across multiple categories. Can be backed by any LLM provider.
224 """
226 SAFETY_PROMPT = """Assess the following content for safety violations.
227Categories: toxicity, self_harm, violence, sexual, jailbreak, data_leak.
229Respond with ONLY a JSON object:
230{
231 "safe": true/false,
232 "categories": [{"category": "<name>", "severity": "low|medium|high|critical", "reason": "<brief>"}]
233}
235Content to assess:
236---
237{content}
238---"""
240 def __init__(self, llm_call: Optional[Callable] = None):
241 self._llm_call = llm_call
243 async def assess(self, content: str) -> List[GuardViolation]:
244 if self._llm_call is None:
245 return [] # No LLM backend configured, skip
247 prompt = self.SAFETY_PROMPT.format(content=content[:4000])
248 try:
249 response = await self._llm_call(prompt)
250 result = json.loads(response)
251 except Exception:
252 return []
254 if result.get("safe", True):
255 return []
257 violations = []
258 severity_map = {
259 "low": ViolationSeverity.LOW,
260 "medium": ViolationSeverity.MEDIUM,
261 "high": ViolationSeverity.HIGH,
262 "critical": ViolationSeverity.CRITICAL,
263 }
264 for cat in result.get("categories", []):
265 cat_name = cat.get("category", "custom")
266 try:
267 cat_enum = Category(cat_name)
268 except ValueError:
269 cat_enum = Category.CUSTOM
271 violations.append(GuardViolation(
272 category=cat_enum,
273 severity=severity_map.get(cat.get("severity", "medium"), ViolationSeverity.MEDIUM),
274 action=GuardAction.BLOCK,
275 message=cat.get("reason", f"Content safety violation: {cat_name}"),
276 metadata={"llm_assessment": cat},
277 ))
279 return violations
282# ---------------------------------------------------------------------------
283# Guardrails Pipeline
284# ---------------------------------------------------------------------------
286class GuardrailsPipeline:
287 """
288 Production guardrails pipeline combining regex fast-path and LLM-based
289 content safety assessment.
291 Usage:
292 pipeline = GuardrailsPipeline()
293 pipeline.add_regex_rule(...)
295 # Input validation
296 result = await pipeline.check_input(user_message)
297 if not result.passed:
298 raise GuardViolationError(result)
300 # Output validation
301 result = await pipeline.check_output(agent_response)
302 """
304 def __init__(
305 self,
306 regex_guard: Optional[RegexGuard] = None,
307 safety_guard: Optional[ContentSafetyGuard] = None,
308 enable_regex: bool = True,
309 enable_safety: bool = True,
310 ):
311 self._regex = regex_guard or RegexGuard()
312 self._safety = safety_guard or ContentSafetyGuard()
313 self._enable_regex = enable_regex
314 self._enable_safety = enable_safety
315 self._audit_log: List[GuardResult] = []
317 def add_regex_rule(self, rule: RegexRule) -> None:
318 self._regex.add_rule(rule)
320 def remove_regex_rule(self, rule_id: str) -> None:
321 self._regex.remove_rule(rule_id)
323 async def check_input(self, content: str) -> GuardResult:
324 """Validate user input before agent processing."""
325 return await self._check(content, stage="input")
327 async def check_output(self, content: str) -> GuardResult:
328 """Validate agent output before returning to user."""
329 return await self._check(content, stage="output")
331 async def check_tool_call(self, tool_name: str, arguments: Dict[str, Any]) -> GuardResult:
332 """Validate tool calls for safety."""
333 content = f"Tool: {tool_name}\nArgs: {json.dumps(arguments)}"
334 return await self._check(content, stage="tool_call")
336 async def _check(self, content: str, stage: str = "unknown") -> GuardResult:
337 violations: List[GuardViolation] = []
339 # Fast-path: regex scanning
340 if self._enable_regex:
341 violations.extend(self._regex.scan(content))
343 # Deep check: LLM safety assessment
344 if self._enable_safety and content.strip():
345 safety_violations = await self._safety.assess(content)
346 violations.extend(safety_violations)
348 # Determine result
349 if not violations:
350 result = GuardResult(passed=True)
351 else:
352 redacted = self._regex.redact(content, violations) if any(
353 v.action == GuardAction.REDACT for v in violations
354 ) else None
356 result = GuardResult(
357 passed=not any(v.action == GuardAction.BLOCK for v in violations),
358 violations=violations,
359 redacted_content=redacted,
360 warnings=[v.message for v in violations if v.action == GuardAction.WARN],
361 )
363 self._audit_log.append(result)
364 return result
366 def get_audit_log(self) -> List[Dict[str, Any]]:
367 return [r.to_dict() for r in self._audit_log]
369 def get_statistics(self) -> Dict[str, int]:
370 total = len(self._audit_log)
371 blocked = sum(1 for r in self._audit_log if r.blocked)
372 passed = sum(1 for r in self._audit_log if r.passed and not r.violations)
373 warned = total - blocked - passed
374 return {
375 "total_checks": total,
376 "passed": passed,
377 "blocked": blocked,
378 "warned": warned,
379 }
382# ---------------------------------------------------------------------------
383# Exception
384# ---------------------------------------------------------------------------
386class GuardViolationError(Exception):
387 """Raised when guardrails block a request."""
389 def __init__(self, result: GuardResult):
390 self.result = result
391 violations_summary = "; ".join(
392 f"[{v.category.value}] {v.message}" for v in result.violations
393 )
394 super().__init__(f"Guardrail blocked: {violations_summary}")
397# ---------------------------------------------------------------------------
398# Convenience: Pre-built Pipeline
399# ---------------------------------------------------------------------------
401def create_default_pipeline() -> GuardrailsPipeline:
402 """Create a GuardrailsPipeline with sensible defaults."""
403 return GuardrailsPipeline(
404 regex_guard=RegexGuard(rules=DEFAULT_RULES),
405 enable_regex=True,
406 enable_safety=False, # LLM-based safety off by default; opt-in
407 )
410def create_strict_pipeline() -> GuardrailsPipeline:
411 """Create a GuardrailsPipeline with strict rules + LLM safety."""
412 return GuardrailsPipeline(
413 regex_guard=RegexGuard(rules=DEFAULT_RULES),
414 safety_guard=ContentSafetyGuard(),
415 enable_regex=True,
416 enable_safety=True,
417 )