Coverage for agentos/guardrails/rules.py: 96%

1""" 

2Built-in guardrail rules — PII detection, keyword blocking, length limits, regex, 

3toxicity heuristics, and code injection detection. 

4""" 

5 

6import re 

7 

8from agentos.guardrails.engine import GuardrailAction, GuardrailCategory, GuardrailRule 

9 

10 

11def PIIRule( # noqa: N802 

12 name: str = "pii_detector", 

13 action: GuardrailAction = GuardrailAction.SANITIZE, 

14 enabled: bool = True, 

15) -> GuardrailRule: 

16 """Detects common PII patterns (email, phone, SSN, credit card) and redacts.""" 

17 

18 _pii_patterns = [ 

19 (r"\b[\w._%+-]+@[\w.-]+\.[a-zA-Z]{2,}\b", "[EMAIL]"), 

20 (r"\b\d{3}[-.]?\d{3}[-.]?\d{4}\b", "[PHONE]"), 

21 (r"\b\d{3}-\d{2}-\d{4}\b", "[SSN]"), 

22 (r"\b(?:\d{4}[- ]?){3}\d{4}\b", "[CARD]"), 

23 ] 

24 

25 def _check(text: str) -> bool: 

26 for pat, _ in _pii_patterns: 

27 if re.search(pat, text): 

28 return True 

29 return False 

30 

31 def _sanitize(text: str) -> str: 

32 for pat, repl in _pii_patterns: 

33 text = re.sub(pat, repl, text) 

34 return text 

35 

36 return GuardrailRule( 

37 name=name, 

38 category=GuardrailCategory.PII, 

39 action=action, 

40 check=_check, 

41 sanitize=_sanitize, 

42 description="Redacts emails, phone numbers, SSNs, and credit card numbers.", 

43 enabled=enabled, 

44 ) 

45 

46 

47def KeywordBlockRule( # noqa: N802 

48 keywords: list[str], 

49 name: str = "keyword_block", 

50 case_sensitive: bool = False, 

51 enabled: bool = True, 

52) -> GuardrailRule: 

53 """Blocks text containing any of the specified keywords.""" 

54 

55 _kw = keywords if case_sensitive else [kw.lower() for kw in keywords] 

56 

57 def _check(text: str) -> bool: 

58 t = text if case_sensitive else text.lower() 

59 return any(kw in t for kw in _kw) 

60 

61 return GuardrailRule( 

62 name=name, 

63 category=GuardrailCategory.KEYWORD, 

64 action=GuardrailAction.BLOCK, 

65 check=_check, 

66 description=f"Blocks content containing: {', '.join(keywords[:5])}", 

67 enabled=enabled, 

68 ) 

69 

70 

71def LengthLimitRule( # noqa: N802 

72 max_input: int = 32_000, 

73 max_output: int = 16_000, 

74 name: str = "length_limit", 

75 enabled: bool = True, 

76) -> GuardrailRule: 

77 """Blocks text exceeding length limits (input or output).""" 

78 

79 def _check(text: str) -> bool: 

80 return len(text) > max(max_input, max_output) 

81 

82 return GuardrailRule( 

83 name=name, 

84 category=GuardrailCategory.LENGTH, 

85 action=GuardrailAction.BLOCK, 

86 check=_check, 

87 description=f"Limits input to {max_input} chars, output to {max_output} chars.", 

88 enabled=enabled, 

89 ) 

90 

91 

92def RegexRule( # noqa: N802 

93 pattern: str, 

94 name: str = "regex_rule", 

95 action: GuardrailAction = GuardrailAction.FLAG, 

96 description: str = "", 

97 enabled: bool = True, 

98) -> GuardrailRule: 

99 """Flags or blocks text matching a custom regex pattern.""" 

100 _pat = re.compile(pattern) 

101 

102 def _check(text: str) -> bool: 

103 return bool(_pat.search(text)) 

104 

105 return GuardrailRule( 

106 name=name, 

107 category=GuardrailCategory.CUSTOM, 

108 action=action, 

109 check=_check, 

110 description=description or f"Regex: {pattern[:40]}", 

111 enabled=enabled, 

112 ) 

113 

114 

115def ToxicityRule( # noqa: N802 

116 name: str = "toxicity_check", 

117 action: GuardrailAction = GuardrailAction.FLAG, 

118 enabled: bool = True, 

119) -> GuardrailRule: 

120 """Heuristic toxicity detection via keyword lists (offline, no API call).""" 

121 

122 _toxic = [ 

123 "kill yourself", 

124 "kys", 

125 "die in a fire", 

126 "i hope you die", 

127 "nigger", 

128 "faggot", 

129 "retard", 

130 "cunt", 

131 "terrorist", 

132 "bomb making", 

133 "how to make a bomb", 

134 "child porn", 

135 "cp ", 

136 "lolicon", 

137 ] 

138 

139 def _check(text: str) -> bool: 

140 t = text.lower() 

141 return any(toxic in t for toxic in _toxic) 

142 

143 return GuardrailRule( 

144 name=name, 

145 category=GuardrailCategory.TOXICITY, 

146 action=action, 

147 check=_check, 

148 description="Flags text containing toxic or harmful language.", 

149 enabled=enabled, 

150 ) 

151 

152 

153def CodeInjectionRule( # noqa: N802 

154 name: str = "code_injection_detector", 

155 action: GuardrailAction = GuardrailAction.BLOCK, 

156 enabled: bool = True, 

157) -> GuardrailRule: 

158 """Detects prompt injection and code injection patterns.""" 

159 

160 _patterns = [ 

161 r"ignore (all )?(previous|above|prior) (instructions?|prompts?)", 

162 r"forget (your|all) (instructions?|rules?|training)", 

163 r"you are now (DAN|developer mode|jailbroken)", 

164 r"system:\s*you are", 

165 r"<\|im_start\|>", 

166 r"<\|system\|>", 

167 r"```.*\b(?:rm\s+-rf|DROP\s+TABLE|DELETE\s+FROM|shutdown)\b", 

168 r"\b(?:DROP\s+TABLE|DELETE\s+FROM|TRUNCATE\s+TABLE|ALTER\s+TABLE)\b", 

169 r"\brm\s+-rf\s+/", 

170 r"\bexec\s*\(.*\)", 

171 r"\beval\s*\(.*\)", 

172 r"\b__import__\s*\(.*\)", 

173 r"\bimportlib\.import_module\b", 

174 ] 

175 

176 _compiled = [re.compile(p, re.IGNORECASE) for p in _patterns] 

177 

178 def _check(text: str) -> bool: 

179 return any(pat.search(text) for pat in _compiled) 

180 

181 return GuardrailRule( 

182 name=name, 

183 category=GuardrailCategory.INJECTION, 

184 action=action, 

185 check=_check, 

186 description="Blocks prompt injection and code injection attempts.", 

187 enabled=enabled, 

188 ) 

189 

190 

191def build_default_rules( 

192 blocked_keywords: list[str] | None = None, 

193 max_input_length: int = 32_000, 

194 max_output_length: int = 16_000, 

195) -> list[GuardrailRule]: 

196 """Build a sensible default rule set for production use.""" 

197 rules: list[GuardrailRule] = [ 

198 CodeInjectionRule(), 

199 PIIRule(), 

200 ToxicityRule(action=GuardrailAction.FLAG), 

201 LengthLimitRule(max_input=max_input_length, max_output=max_output_length), 

202 ] 

203 if blocked_keywords: 

204 rules.append(KeywordBlockRule(keywords=blocked_keywords)) 

205 return rules