Coverage for agentos/guardrails/policy.py: 0%

1""" 

2Guardrail policy enforcement — cumulative violation tracking, rate limiting, 

3and session-scoped policy decisions. 

4""" 

5 

6from collections.abc import Callable 

7from dataclasses import dataclass, field 

8from enum import StrEnum 

9from typing import Any 

10 

11from agentos.guardrails.engine import GuardrailAction, GuardrailCategory, GuardrailResult 

12 

13 

14class PolicyViolation(StrEnum): 

15 """Policy-level violation reasons.""" 

16 

17 SESSION_BLOCKED = "session_blocked" 

18 RATE_LIMITED = "rate_limited" 

19 CUMULATIVE_VIOLATIONS = "cumulative_violations" 

20 CATEGORY_BANNED = "category_banned" 

21 

22 

23@dataclass 

24class GuardrailPolicy: 

25 """Session-scoped policy configuration.""" 

26 

27 max_total_violations: int = 5 

28 max_violations_per_category: dict[str, int] = field(default_factory=dict) 

29 window_seconds: int = 300 

30 auto_block_on: set[GuardrailCategory] = field(default_factory=set) 

31 on_session_block: str = "reject" # reject / warn 

32 monitoring_callback: Callable[[str, dict[str, Any]], None] | None = None 

33 

34 def __post_init__(self): 

35 if not self.max_violations_per_category: 

36 self.max_violations_per_category = { 

37 GuardrailCategory.INJECTION.value: 2, 

38 GuardrailCategory.TOXICITY.value: 3, 

39 GuardrailCategory.KEYWORD.value: 3, 

40 } 

41 

42 

43class PolicyEnforcer: 

44 """Tracks violations per session and enforces cumulative policy.""" 

45 

46 def __init__(self, policy: GuardrailPolicy | None = None): 

47 self.policy = policy or GuardrailPolicy() 

48 self._violation_count: int = 0 

49 self._category_counts: dict[str, int] = {} 

50 self._session_blocked: bool = False 

51 self._violation_log: list[tuple[float, str, str]] = [] 

52 

53 def evaluate(self, result: GuardrailResult, category: str = "") -> PolicyViolation | None: 

54 """Evaluate a guardrail result against the current policy. 

55 

56 Returns None if no policy violation, or the reason for violation. 

57 """ 

58 if self._session_blocked: 

59 return PolicyViolation.SESSION_BLOCKED 

60 

61 if result.action == GuardrailAction.PASS: 

62 return None 

63 

64 import time 

65 

66 now = time.time() 

67 

68 self._violation_count += 1 

69 if category: 

70 self._category_counts[category] = self._category_counts.get(category, 0) + 1 

71 self._violation_log.append((now, category, result.action.value)) 

72 

73 # Clean old entries outside window 

74 cutoff = now - self.policy.window_seconds 

75 self._violation_log = [(t, c, a) for t, c, a in self._violation_log if t > cutoff] 

76 

77 # Check cumulative violations 

78 if self._violation_count >= self.policy.max_total_violations: 

79 self._session_blocked = True 

80 self._emit("session_blocked", {"total_violations": self._violation_count}) 

81 return PolicyViolation.CUMULATIVE_VIOLATIONS 

82 

83 # Check per-category limits 

84 if category and category in self.policy.max_violations_per_category: 

85 limit = self.policy.max_violations_per_category[category] 

86 if self._category_counts[category] >= limit: 

87 self._emit( 

88 "category_blocked", 

89 {"category": category, "count": self._category_counts[category]}, 

90 ) 

91 return PolicyViolation.CATEGORY_BANNED 

92 

93 return None 

94 

95 def reset(self) -> None: 

96 """Reset all violation counters for a new session.""" 

97 self._violation_count = 0 

98 self._category_counts.clear() 

99 self._session_blocked = False 

100 self._violation_log.clear() 

101 

102 @property 

103 def is_blocked(self) -> bool: 

104 return self._session_blocked 

105 

106 @property 

107 def total_violations(self) -> int: 

108 return self._violation_count 

109 

110 def _emit(self, event: str, data: dict[str, Any]) -> None: 

111 if self.policy.monitoring_callback: 

112 try: 

113 self.policy.monitoring_callback(event, data) 

114 except Exception: 

115 pass