Coverage for agentos/security/guardrails.py: 0%
156 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 12:57 +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 json
26import re
27from collections.abc import Callable
28from dataclasses import dataclass, field
29from enum import StrEnum
30from re import Pattern
31from typing import Any
33# ---------------------------------------------------------------------------
34# Enums & Data Classes
35# ---------------------------------------------------------------------------
38class ViolationSeverity(StrEnum):
39 """Severity level of a guardrail violation."""
41 CRITICAL = "critical" # Immediate block, alert ops
42 HIGH = "high" # Block the request
43 MEDIUM = "medium" # Warn but allow (with redaction)
44 LOW = "low" # Log only
47class GuardAction(StrEnum):
48 """Action to take when a guardrail is triggered."""
50 BLOCK = "block" # Reject the request entirely
51 WARN = "warn" # Allow but flag with warning
52 REDACT = "redact" # Remove offending content, allow rest
53 LOG = "log" # Log only, no user-visible effect
56class Category(StrEnum):
57 """Standard content safety categories."""
59 PII = "pii" # Personally Identifiable Information
60 TOXICITY = "toxicity" # Hate speech, harassment
61 SELF_HARM = "self_harm" # Suicide, self-injury
62 VIOLENCE = "violence" # Graphic violence
63 SEXUAL = "sexual" # Explicit sexual content
64 JAILBREAK = "jailbreak" # Prompt injection / jailbreak attempts
65 DATA_LEAK = "data_leak" # Attempting to leak system prompts / internals
66 MALICIOUS_CODE = "malicious_code" # Code injection, reverse shell, etc.
67 OFF_TOPIC = "off_topic" # Outside defined scope
68 CUSTOM = "custom" # User-defined category
71@dataclass
72class GuardViolation:
73 """A single guardrail violation detected."""
75 category: Category
76 severity: ViolationSeverity
77 action: GuardAction
78 message: str
79 matched_pattern: str | None = None
80 matched_text: str | None = None
81 rule_id: str | None = None
82 metadata: dict[str, Any] = field(default_factory=dict)
85@dataclass
86class GuardResult:
87 """Result of running guardrails on content."""
89 passed: bool = True
90 violations: list[GuardViolation] = field(default_factory=list)
91 redacted_content: str | None = None
92 warnings: list[str] = field(default_factory=list)
94 @property
95 def blocked(self) -> bool:
96 return any(v.action == GuardAction.BLOCK for v in self.violations)
98 def to_dict(self) -> dict[str, Any]:
99 return {
100 "passed": self.passed,
101 "blocked": self.blocked,
102 "violations": [
103 {
104 "category": v.category.value,
105 "severity": v.severity.value,
106 "action": v.action.value,
107 "message": v.message,
108 "rule_id": v.rule_id,
109 }
110 for v in self.violations
111 ],
112 "warnings": self.warnings,
113 }
116# ---------------------------------------------------------------------------
117# PII Detection Patterns
118# ---------------------------------------------------------------------------
120PII_PATTERNS: dict[str, Pattern[str]] = {
121 "email": re.compile(r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}"),
122 "phone_cn": re.compile(r"1[3-9]\d{9}"),
123 "phone_us": re.compile(r"\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}"),
124 "ssn": re.compile(r"\d{3}-\d{2}-\d{4}"),
125 "credit_card": re.compile(r"\b(?:\d{4}[ -]?){3}\d{4}\b"),
126 "ip_address": re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"),
127 "api_key": re.compile(
128 r"(?:api[_-]?key|apikey|token|secret|password)\s*[:=]\s*['\"]?[\w-]{20,}['\"]?",
129 re.IGNORECASE,
130 ),
131}
134# ---------------------------------------------------------------------------
135# Regex-based Fast-Path Rules
136# ---------------------------------------------------------------------------
139@dataclass
140class RegexRule:
141 """A regex-based guardrail rule for fast-path matching."""
143 rule_id: str
144 category: Category
145 severity: ViolationSeverity
146 action: GuardAction
147 pattern: Pattern[str]
148 message: str
151DEFAULT_RULES: list[RegexRule] = [
152 # PII Rules
153 RegexRule(
154 "pii-email",
155 Category.PII,
156 ViolationSeverity.HIGH,
157 GuardAction.REDACT,
158 PII_PATTERNS["email"],
159 "Email address detected",
160 ),
161 RegexRule(
162 "pii-phone-cn",
163 Category.PII,
164 ViolationSeverity.MEDIUM,
165 GuardAction.REDACT,
166 PII_PATTERNS["phone_cn"],
167 "Chinese phone number detected",
168 ),
169 RegexRule(
170 "pii-ssn",
171 Category.PII,
172 ViolationSeverity.CRITICAL,
173 GuardAction.BLOCK,
174 PII_PATTERNS["ssn"],
175 "SSN detected",
176 ),
177 RegexRule(
178 "pii-cc",
179 Category.PII,
180 ViolationSeverity.CRITICAL,
181 GuardAction.BLOCK,
182 PII_PATTERNS["credit_card"],
183 "Credit card number detected",
184 ),
185 RegexRule(
186 "pii-apikey",
187 Category.PII,
188 ViolationSeverity.CRITICAL,
189 GuardAction.BLOCK,
190 PII_PATTERNS["api_key"],
191 "Potential API key in text",
192 ),
193 # Jailbreak patterns
194 RegexRule(
195 "jb-ignore",
196 Category.JAILBREAK,
197 ViolationSeverity.CRITICAL,
198 GuardAction.BLOCK,
199 re.compile(
200 r"(?:ignore|forget|disregard)\s+(?:all\s+)?(?:previous|above|prior)\s+(?:instructions?|prompts?|rules?)",
201 re.IGNORECASE,
202 ),
203 "Jailbreak attempt: ignore instructions",
204 ),
205 RegexRule(
206 "jb-dan",
207 Category.JAILBREAK,
208 ViolationSeverity.CRITICAL,
209 GuardAction.BLOCK,
210 re.compile(r"\bDAN\s*(?:mode|jailbreak)?\b", re.IGNORECASE),
211 "Jailbreak attempt: DAN mode",
212 ),
213 RegexRule(
214 "jb-roleplay",
215 Category.JAILBREAK,
216 ViolationSeverity.HIGH,
217 GuardAction.BLOCK,
218 re.compile(
219 r"(?:pretend|act\s+as\s+if|imagine)\s+you\s+(?:are|were)\s+(?:an?\s+)?(?:unfiltered|unrestricted|evil|dark|malicious)",
220 re.IGNORECASE,
221 ),
222 "Jailbreak attempt: roleplay escalation",
223 ),
224 # Malicious code
225 RegexRule(
226 "mc-reverse-shell",
227 Category.MALICIOUS_CODE,
228 ViolationSeverity.CRITICAL,
229 GuardAction.BLOCK,
230 re.compile(
231 r"(?:bash|sh|nc|netcat|ncat)\s+.*(?:>&?\s*/dev/(?:tcp|udp)|-e\s+/bin/(?:bash|sh))",
232 re.IGNORECASE,
233 ),
234 "Reverse shell attempt detected",
235 ),
236 RegexRule(
237 "mc-rm-rf",
238 Category.MALICIOUS_CODE,
239 ViolationSeverity.HIGH,
240 GuardAction.BLOCK,
241 re.compile(r"(?:rm\s+-rf|del\s+/[fsq])\s+(?:/|~|\*)", re.IGNORECASE),
242 "Destructive file operation detected",
243 ),
244]
247# ---------------------------------------------------------------------------
248# Guardrail Engine
249# ---------------------------------------------------------------------------
252class RegexGuard:
253 """Fast-path regex-based guard for common patterns."""
255 def __init__(self, rules: list[RegexRule] | None = None):
256 self._rules: dict[str, RegexRule] = {}
257 for rule in rules or DEFAULT_RULES:
258 self._rules[rule.rule_id] = rule
260 def add_rule(self, rule: RegexRule) -> None:
261 self._rules[rule.rule_id] = rule
263 def remove_rule(self, rule_id: str) -> None:
264 self._rules.pop(rule_id, None)
266 def scan(self, content: str) -> list[GuardViolation]:
267 violations: list[GuardViolation] = []
268 for rule in self._rules.values():
269 for match in rule.pattern.finditer(content):
270 violations.append(
271 GuardViolation(
272 category=rule.category,
273 severity=rule.severity,
274 action=rule.action,
275 message=rule.message,
276 matched_pattern=rule.pattern.pattern,
277 matched_text=match.group(),
278 rule_id=rule.rule_id,
279 )
280 )
281 return violations
283 def redact(self, content: str, violations: list[GuardViolation]) -> str:
284 """Redact PII from content based on matched violations."""
285 result = content
286 for v in violations:
287 if v.action == GuardAction.REDACT and v.matched_text:
288 category = v.category.value
289 replacement = f"[REDACTED_{category.upper()}]"
290 result = result.replace(v.matched_text, replacement)
291 return result
294class ContentSafetyGuard:
295 """
296 LLM-based content safety assessment.
298 Uses a lightweight classifier prompt to assess content safety
299 across multiple categories. Can be backed by any LLM provider.
300 """
302 SAFETY_PROMPT = """Assess the following content for safety violations.
303Categories: toxicity, self_harm, violence, sexual, jailbreak, data_leak.
305Respond with ONLY a JSON object:
306{
307 "safe": true/false,
308 "categories": [{"category": "<name>", "severity": "low|medium|high|critical", "reason": "<brief>"}]
309}
311Content to assess:
312---
313{content}
314---"""
316 def __init__(self, llm_call: Callable | None = None):
317 self._llm_call = llm_call
319 async def assess(self, content: str) -> list[GuardViolation]:
320 if self._llm_call is None:
321 return [] # No LLM backend configured, skip
323 prompt = self.SAFETY_PROMPT.format(content=content[:4000])
324 try:
325 response = await self._llm_call(prompt)
326 result = json.loads(response)
327 except Exception:
328 return []
330 if result.get("safe", True):
331 return []
333 violations = []
334 severity_map = {
335 "low": ViolationSeverity.LOW,
336 "medium": ViolationSeverity.MEDIUM,
337 "high": ViolationSeverity.HIGH,
338 "critical": ViolationSeverity.CRITICAL,
339 }
340 for cat in result.get("categories", []):
341 cat_name = cat.get("category", "custom")
342 try:
343 cat_enum = Category(cat_name)
344 except ValueError:
345 cat_enum = Category.CUSTOM
347 violations.append(
348 GuardViolation(
349 category=cat_enum,
350 severity=severity_map.get(
351 cat.get("severity", "medium"), ViolationSeverity.MEDIUM
352 ),
353 action=GuardAction.BLOCK,
354 message=cat.get("reason", f"Content safety violation: {cat_name}"),
355 metadata={"llm_assessment": cat},
356 )
357 )
359 return violations
362# ---------------------------------------------------------------------------
363# Guardrails Pipeline
364# ---------------------------------------------------------------------------
367class GuardrailsPipeline:
368 """
369 Production guardrails pipeline combining regex fast-path and LLM-based
370 content safety assessment.
372 Usage:
373 pipeline = GuardrailsPipeline()
374 pipeline.add_regex_rule(...)
376 # Input validation
377 result = await pipeline.check_input(user_message)
378 if not result.passed:
379 raise GuardViolationError(result)
381 # Output validation
382 result = await pipeline.check_output(agent_response)
383 """
385 def __init__(
386 self,
387 regex_guard: RegexGuard | None = None,
388 safety_guard: ContentSafetyGuard | None = None,
389 enable_regex: bool = True,
390 enable_safety: bool = True,
391 ):
392 self._regex = regex_guard or RegexGuard()
393 self._safety = safety_guard or ContentSafetyGuard()
394 self._enable_regex = enable_regex
395 self._enable_safety = enable_safety
396 self._audit_log: list[GuardResult] = []
398 def add_regex_rule(self, rule: RegexRule) -> None:
399 self._regex.add_rule(rule)
401 def remove_regex_rule(self, rule_id: str) -> None:
402 self._regex.remove_rule(rule_id)
404 async def check_input(self, content: str) -> GuardResult:
405 """Validate user input before agent processing."""
406 return await self._check(content, stage="input")
408 async def check_output(self, content: str) -> GuardResult:
409 """Validate agent output before returning to user."""
410 return await self._check(content, stage="output")
412 async def check_tool_call(self, tool_name: str, arguments: dict[str, Any]) -> GuardResult:
413 """Validate tool calls for safety."""
414 content = f"Tool: {tool_name}\nArgs: {json.dumps(arguments)}"
415 return await self._check(content, stage="tool_call")
417 async def _check(self, content: str, stage: str = "unknown") -> GuardResult:
418 violations: list[GuardViolation] = []
420 # Fast-path: regex scanning
421 if self._enable_regex:
422 violations.extend(self._regex.scan(content))
424 # Deep check: LLM safety assessment
425 if self._enable_safety and content.strip():
426 safety_violations = await self._safety.assess(content)
427 violations.extend(safety_violations)
429 # Determine result
430 if not violations:
431 result = GuardResult(passed=True)
432 else:
433 redacted = (
434 self._regex.redact(content, violations)
435 if any(v.action == GuardAction.REDACT for v in violations)
436 else None
437 )
439 result = GuardResult(
440 passed=not any(v.action == GuardAction.BLOCK for v in violations),
441 violations=violations,
442 redacted_content=redacted,
443 warnings=[v.message for v in violations if v.action == GuardAction.WARN],
444 )
446 self._audit_log.append(result)
447 return result
449 def get_audit_log(self) -> list[dict[str, Any]]:
450 return [r.to_dict() for r in self._audit_log]
452 def get_statistics(self) -> dict[str, int]:
453 total = len(self._audit_log)
454 blocked = sum(1 for r in self._audit_log if r.blocked)
455 passed = sum(1 for r in self._audit_log if r.passed and not r.violations)
456 warned = total - blocked - passed
457 return {
458 "total_checks": total,
459 "passed": passed,
460 "blocked": blocked,
461 "warned": warned,
462 }
465# ---------------------------------------------------------------------------
466# Exception
467# ---------------------------------------------------------------------------
470class GuardViolationError(Exception):
471 """Raised when guardrails block a request."""
473 def __init__(self, result: GuardResult):
474 self.result = result
475 violations_summary = "; ".join(
476 f"[{v.category.value}] {v.message}" for v in result.violations
477 )
478 super().__init__(f"Guardrail blocked: {violations_summary}")
481# ---------------------------------------------------------------------------
482# Convenience: Pre-built Pipeline
483# ---------------------------------------------------------------------------
486def create_default_pipeline() -> GuardrailsPipeline:
487 """Create a GuardrailsPipeline with sensible defaults."""
488 return GuardrailsPipeline(
489 regex_guard=RegexGuard(rules=DEFAULT_RULES),
490 enable_regex=True,
491 enable_safety=False, # LLM-based safety off by default; opt-in
492 )
495def create_strict_pipeline() -> GuardrailsPipeline:
496 """Create a GuardrailsPipeline with strict rules + LLM safety."""
497 return GuardrailsPipeline(
498 regex_guard=RegexGuard(rules=DEFAULT_RULES),
499 safety_guard=ContentSafetyGuard(),
500 enable_regex=True,
501 enable_safety=True,
502 )