Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-guard/src/lexigram/ai/guard/output/pii_redactor.py: 64%

1"""PII redactor output guard. 

2 

3Scans LLM responses for PII and redacts it before returning content to 

4or storing it. Uses the same regex-based detection as the input 

5:class:`~lexigram.ai.guard.input.pii.PIIDetector`. 

6 

7This guard always uses ``action="redact"`` — if you want to *block* 

8responses containing PII use a ``PIIDetector`` with ``action="block"`` 

9as the output guard instead. 

10""" 

11 

12from __future__ import annotations 

13 

14from typing import TYPE_CHECKING, Any 

15 

16from lexigram.ai.guard.input.pii import _PII_PATTERNS, _redact_pii, _scan_for_pii 

17from lexigram.ai.guard.output.base import AbstractOutputGuard 

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

19from lexigram.contracts.ai.guards import GuardResultProtocol 

20from lexigram.result import Ok, Result 

21 

22if TYPE_CHECKING: 

23 from lexigram.contracts.ai.exceptions import GuardError 

24 

25 

26class PIIRedactor(AbstractOutputGuard): 

27 """Output guard that redacts PII from LLM responses. 

28 

29 Args: 

30 entities: List of PII entity types to redact. 

31 Defaults to all supported types. 

32 action: ``"redact"`` (default) or ``"block"`` (rejects any 

33 response containing PII rather than sanitising it). 

34 

35 Example:: 

36 

37 guard = PIIRedactor(entities=["SSN", "CREDIT_CARD", "EMAIL"]) 

38 result = await guard.check(llm_response) 

39 safe_response = result.final_content or llm_response 

40 """ 

41 

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

43 

44 def __init__( 

45 self, 

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

47 action: str = "redact", 

48 ) -> None: 

49 """Initialise the PII redactor. 

50 

51 Args: 

52 entities: PII types to redact. Defaults to all types. 

53 action: ``"redact"`` or ``"block"``. 

54 """ 

55 super().__init__(action=action) 

56 self._entities: list[str] = [ 

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

58 ] 

59 

60 async def check( 

61 self, 

62 content: str, 

63 *, 

64 original_input: str | None = None, 

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

66 ) -> Result[GuardResultProtocol, GuardError]: 

67 """Scan LLM response for PII and apply the configured action. 

68 

69 Args: 

70 content: LLM response text to scan. 

71 original_input: Unused — present for protocol compatibility. 

72 metadata: Optional metadata. 

73 

74 Returns: 

75 PASS if no PII; REDACT with sanitised content or BLOCK otherwise. 

76 """ 

77 found = _scan_for_pii(content, self._entities) 

78 

79 if not found: 

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

81 

82 entity_list = list(found.keys()) 

83 

84 if self._action == "block": 

85 return Ok( 

86 GuardCheckResult.block( 

87 self.name, 

88 reason=f"LLM response contains PII: {', '.join(entity_list)}", 

89 detected_entities=entity_list, 

90 ) 

91 ) 

92 

93 redacted = _redact_pii(content, found) 

94 return Ok( 

95 GuardCheckResult.redact( 

96 self.name, 

97 redacted_content=redacted, 

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

99 detected_entities=entity_list, 

100 ) 

101 ) 

102 

103 

104__all__ = ["PIIRedactor"]