Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-guard/src/lexigram/ai/guard/input/pii.py: 41%

1"""PII (Personally IdentifiableProtocol Information) detector for input content. 

2 

3Detects and optionally redacts PII in user input before it is sent to an 

4external LLM provider. Detection uses regex patterns — no external API 

5calls are made. 

6 

7Supported PII entity types: 

8- ``EMAIL`` — email addresses 

9- ``PHONE`` — US/international phone numbers 

10- ``SSN`` — US Social Security Numbers (``xxx-xx-xxxx`` format) 

11- ``CREDIT_CARD`` — major card formats (Visa, MC, Amex, Discover) 

12- ``IP_ADDRESS`` — IPv4 addresses 

13- ``AWS_KEY`` — AWS access key IDs 

14 

15When ``action="redact"`` the matched content is replaced with 

16``[REDACTED:<ENTITY_TYPE>]``. When ``action="block"`` the request is 

17rejected entirely. When ``action="warn"`` the request is allowed but a 

18structured warning is returned. 

19""" 

20 

21from __future__ import annotations 

22 

23import re 

24from typing import TYPE_CHECKING, Any 

25 

26from lexigram.ai.guard.input.base import AbstractInputGuard 

27from lexigram.ai.guard.pipeline.result import GuardCheckResult 

28from lexigram.contracts.ai.guards import GuardResultProtocol 

29from lexigram.result import Ok, Result 

30 

31if TYPE_CHECKING: 

32 from lexigram.contracts.ai.exceptions import GuardError 

33 

34# --------------------------------------------------------------------------- 

35# PII patterns 

36# --------------------------------------------------------------------------- 

37 

38_PII_PATTERNS: dict[str, re.Pattern[str]] = { 

39 "EMAIL": re.compile(r"\b[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}\b"), 

40 "PHONE": re.compile(r"\b(?:\+?1[-.\s]?)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}\b"), 

41 "SSN": re.compile(r"\b(?!000|666|9\d{2})\d{3}-(?!00)\d{2}-(?!0000)\d{4}\b"), 

42 "CREDIT_CARD": re.compile( 

43 r"\b(?:4[0-9]{12}(?:[0-9]{3})?|[25][1-7][0-9]{14}|" 

44 r"6(?:011|5[0-9]{2})[0-9]{12}|3[47][0-9]{13}|" 

45 r"3(?:0[0-5]|[68][0-9])[0-9]{11})\b" 

46 ), 

47 "IP_ADDRESS": re.compile( 

48 r"\b(?:(?:25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}" 

49 r"(?:25[0-5]|2[0-4]\d|[01]?\d\d?)\b" 

50 ), 

51 "AWS_KEY": re.compile(r"\b(AKIA|ASIA|ABIA|ACCA)[0-9A-Z]{16}\b"), 

52} 

53 

54 

55def _scan_for_pii( 

56 content: str, 

57 entities: list[str], 

58) -> dict[str, list[str]]: 

59 """Scan content for PII matches. 

60 

61 Args: 

62 content: Text to scan. 

63 entities: List of entity type names to check. 

64 

65 Returns: 

66 Dict mapping entity type to list of matched strings. 

67 """ 

68 found: dict[str, list[str]] = {} 

69 for entity in entities: 

70 pattern = _PII_PATTERNS.get(entity.upper()) 

71 if pattern is None: 

72 continue 

73 matches = pattern.findall(content) 

74 if matches: 

75 found[entity.upper()] = matches 

76 return found 

77 

78 

79def _redact_pii(content: str, found: dict[str, list[str]]) -> str: 

80 """Replace PII matches in content with redaction tokens. 

81 

82 Args: 

83 content: Original content. 

84 found: Dict of {entity_type: [matched_strings]}. 

85 

86 Returns: 

87 Content with PII replaced by ``[REDACTED:<TYPE>]`` tokens. 

88 """ 

89 result = content 

90 for entity, matches in found.items(): 

91 for match in matches: 

92 result = result.replace(match, f"[REDACTED:{entity}]") 

93 return result 

94 

95 

96class PIIDetector(AbstractInputGuard): 

97 """Input guard that detects and optionally redacts PII. 

98 

99 Args: 

100 action: ``"redact"`` (default), ``"block"``, or ``"warn"``. 

101 entities: List of PII entity types to detect. 

102 Defaults to all supported types. 

103 

104 Example:: 

105 

106 guard = PIIDetector(action="redact", entities=["EMAIL", "SSN"]) 

107 result = await guard.check(user_message) 

108 safe_content = result.redacted_content or user_message 

109 """ 

110 

111 _DEFAULT_ENTITIES: list[str] = list(_PII_PATTERNS.keys()) 

112 

113 def __init__( 

114 self, 

115 action: str = "redact", 

116 entities: list[str] | None = None, 

117 ) -> None: 

118 """Initialise the PII detector. 

119 

120 Args: 

121 action: Action when PII is found — ``"redact"``, ``"block"``, or ``"warn"``. 

122 entities: PII entity types to scan for. Defaults to all types. 

123 """ 

124 super().__init__(action=action) 

125 self._entities: list[str] = [ 

126 e.upper() for e in (entities or self._DEFAULT_ENTITIES) 

127 ] 

128 

129 async def check( 

130 self, 

131 content: str, 

132 *, 

133 messages: list[Any] | None = None, 

134 metadata: dict[str, Any] | None = None, 

135 ) -> Result[GuardResultProtocol, GuardError]: 

136 """Scan content for PII and apply the configured action. 

137 

138 Args: 

139 content: Input text to scan. 

140 messages: Unused — present for protocol compatibility. 

141 metadata: Optional metadata. 

142 

143 Returns: 

144 PASS if no PII detected; otherwise the configured action result. 

145 """ 

146 found = _scan_for_pii(content, self._entities) 

147 

148 if not found: 

149 return Ok(GuardCheckResult.allow(self.name)) 

150 

151 entity_list = list(found.keys()) 

152 

153 if self._action == "block": 

154 return Ok( 

155 GuardCheckResult.block( 

156 self.name, 

157 reason=f"PII detected: {', '.join(entity_list)}", 

158 detected_entities=entity_list, 

159 ) 

160 ) 

161 

162 if self._action == "warn": 

163 return Ok( 

164 GuardCheckResult.warn( 

165 self.name, 

166 reason=f"PII present in input: {', '.join(entity_list)}", 

167 detected_entities=entity_list, 

168 ) 

169 ) 

170 

171 # Default: redact 

172 redacted = _redact_pii(content, found) 

173 return Ok( 

174 GuardCheckResult.redact( 

175 self.name, 

176 redacted_content=redacted, 

177 reason=f"PII redacted: {', '.join(entity_list)}", 

178 detected_entities=entity_list, 

179 ) 

180 ) 

181 

182 

183__all__ = ["PIIDetector"]