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

1""" 

2Guardrail engine — rule registry, evaluation, and result aggregation. 

3""" 

4 

5from collections.abc import Callable 

6from dataclasses import dataclass, field 

7from enum import StrEnum 

8from typing import Any 

9 

10 

11class GuardrailAction(StrEnum): 

12 """Guardrail disposition for a single rule match.""" 

13 

14 BLOCK = "block" 

15 FLAG = "flag" 

16 SANITIZE = "sanitize" 

17 PASS = "pass" 

18 

19 

20class GuardrailCategory(StrEnum): 

21 """Semantic category of a guardrail rule.""" 

22 

23 PII = "pii" 

24 TOXICITY = "toxicity" 

25 INJECTION = "injection" 

26 KEYWORD = "keyword" 

27 LENGTH = "length" 

28 CUSTOM = "custom" 

29 

30 

31@dataclass 

32class GuardrailResult: 

33 """Aggregate result after all guardrails have been evaluated.""" 

34 

35 passed: bool 

36 action: GuardrailAction 

37 violations: list[str] = field(default_factory=list) 

38 sanitized_text: str | None = None 

39 metadata: dict[str, Any] = field(default_factory=dict) 

40 

41 @property 

42 def blocked(self) -> bool: 

43 return self.action == GuardrailAction.BLOCK 

44 

45 

46@dataclass 

47class GuardrailRule: 

48 """A single guardrail rule definition.""" 

49 

50 name: str 

51 category: GuardrailCategory 

52 action: GuardrailAction 

53 check: Callable[[str], bool] 

54 sanitize: Callable[[str], str] | None = None 

55 description: str = "" 

56 enabled: bool = True 

57 metadata: dict[str, Any] = field(default_factory=dict) 

58 

59 

60class InputGuardrail: 

61 """Validates user prompts before they reach the LLM.""" 

62 

63 def __init__(self, rules: list[GuardrailRule] | None = None): 

64 self._rules: dict[str, GuardrailRule] = {} 

65 if rules: 

66 for r in rules: 

67 self.add_rule(r) 

68 

69 def add_rule(self, rule: GuardrailRule) -> None: 

70 self._rules[rule.name] = rule 

71 

72 def remove_rule(self, name: str) -> None: 

73 self._rules.pop(name, None) 

74 

75 def evaluate(self, text: str) -> GuardrailResult: 

76 """Run all enabled input rules against the text.""" 

77 violations: list[str] = [] 

78 worst_action = GuardrailAction.PASS 

79 sanitized = text 

80 for rule in self._rules.values(): 

81 if not rule.enabled: 

82 continue 

83 if rule.check(sanitized): 

84 violations.append(f"{rule.name}: {rule.description or rule.category.value}") 

85 if rule.sanitize: 

86 sanitized = rule.sanitize(sanitized) 

87 if _action_priority(rule.action) > _action_priority(worst_action): 

88 worst_action = rule.action 

89 

90 passed = worst_action != GuardrailAction.BLOCK 

91 return GuardrailResult( 

92 passed=passed, 

93 action=worst_action, 

94 violations=violations, 

95 sanitized_text=sanitized if sanitized != text else None, 

96 ) 

97 

98 

99class OutputGuardrail: 

100 """Validates LLM outputs before they reach the user.""" 

101 

102 def __init__(self, rules: list[GuardrailRule] | None = None): 

103 self._rules: dict[str, GuardrailRule] = {} 

104 if rules: 

105 for r in rules: 

106 self.add_rule(r) 

107 

108 def add_rule(self, rule: GuardrailRule) -> None: 

109 self._rules[rule.name] = rule 

110 

111 def remove_rule(self, name: str) -> None: 

112 self._rules.pop(name, None) 

113 

114 def evaluate(self, text: str) -> GuardrailResult: 

115 """Run all enabled output rules against the text.""" 

116 violations: list[str] = [] 

117 worst_action = GuardrailAction.PASS 

118 sanitized = text 

119 for rule in self._rules.values(): 

120 if not rule.enabled: 

121 continue 

122 if rule.check(sanitized): 

123 violations.append(f"{rule.name}: {rule.description or rule.category.value}") 

124 if rule.sanitize: 

125 sanitized = rule.sanitize(sanitized) 

126 if _action_priority(rule.action) > _action_priority(worst_action): 

127 worst_action = rule.action 

128 

129 passed = worst_action != GuardrailAction.BLOCK 

130 return GuardrailResult( 

131 passed=passed, 

132 action=worst_action, 

133 violations=violations, 

134 sanitized_text=sanitized if sanitized != text else None, 

135 ) 

136 

137 

138class GuardrailEngine: 

139 """Unified guardrail engine managing both input and output pipelines.""" 

140 

141 def __init__( 

142 self, 

143 input_rules: list[GuardrailRule] | None = None, 

144 output_rules: list[GuardrailRule] | None = None, 

145 ): 

146 self.input = InputGuardrail(input_rules) 

147 self.output = OutputGuardrail(output_rules) 

148 

149 def check_input(self, prompt: str) -> GuardrailResult: 

150 return self.input.evaluate(prompt) 

151 

152 def check_output(self, response: str) -> GuardrailResult: 

153 return self.output.evaluate(response) 

154 

155 def check(self, prompt: str, response: str = "") -> tuple[GuardrailResult, GuardrailResult]: 

156 """Evaluate input and output guardrails. Empty response skips output check.""" 

157 inp = self.input.evaluate(prompt) 

158 out = ( 

159 self.output.evaluate(response) 

160 if response 

161 else GuardrailResult(passed=True, action=GuardrailAction.PASS) 

162 ) 

163 return inp, out 

164 

165 

166def _action_priority(action: GuardrailAction) -> int: 

167 return { 

168 GuardrailAction.PASS: 0, 

169 GuardrailAction.FLAG: 1, 

170 GuardrailAction.SANITIZE: 2, 

171 GuardrailAction.BLOCK: 3, 

172 }[action]