Coverage for little_loops / pii.py: 42%

1"""PII detection and redaction utilities for SFT corpus filtering. 

2 

3Provides regex-based detection and redaction of email, phone, and SSN patterns. 

4Primary consumer is the ``sft-corpus`` FSM loop's ``filter`` state via 

5``apply_pii_action()``. 

6""" 

7 

8from __future__ import annotations 

9 

10import re 

11 

12# Compiled PII patterns — module-level to avoid recompilation on each call 

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

14_PHONE = re.compile(r"\b(\+?1[-.\s]?)?\(?\d{3}\)?[-.\s]\d{3}[-.\s]\d{4}\b") 

15_SSN = re.compile(r"\b\d{3}-\d{2}-\d{4}\b") 

16 

17PII_PATTERNS: dict[str, re.Pattern[str]] = { 

18 "email": _EMAIL, 

19 "phone": _PHONE, 

20 "ssn": _SSN, 

21} 

22 

23_VALID_ACTIONS = frozenset({"flag", "redact", "discard"}) 

24 

25 

26def detect_pii(text: str) -> list[str]: 

27 """Return list of PII type names found in text. 

28 

29 Args: 

30 text: Input text to scan for PII. 

31 

32 Returns: 

33 List of PII type names (e.g. ``["email", "phone"]``) present in the 

34 text. Returns an empty list when no PII is detected. 

35 """ 

36 return [name for name, pattern in PII_PATTERNS.items() if pattern.search(text)] 

37 

38 

39def redact_pii(text: str) -> str: 

40 """Replace PII spans with ``[TYPE]`` placeholders. 

41 

42 Args: 

43 text: Input text to redact. 

44 

45 Returns: 

46 Text with all PII spans replaced by their uppercased type placeholder 

47 (e.g. ``[EMAIL]``, ``[PHONE]``, ``[SSN]``). 

48 """ 

49 for name, pattern in PII_PATTERNS.items(): 

50 text = pattern.sub(f"[{name.upper()}]", text) 

51 return text 

52 

53 

54def apply_pii_action(example: dict, action: str) -> dict | None: 

55 """Apply flag/redact/discard to a formatted SFT example dict. 

56 

57 Scans all top-level string values for PII and applies the requested action. 

58 

59 Args: 

60 example: SFT example dict (e.g. Alpaca ``{"instruction": ..., "output": ...}``). 

61 action: One of ``"flag"``, ``"redact"``, or ``"discard"``. 

62 

63 Returns: 

64 - ``"flag"``: original example with ``pii_detected: True`` added if PII 

65 is present; original example unchanged if no PII found. 

66 - ``"redact"``: copy of example with PII in string values replaced by 

67 ``[TYPE]`` placeholders. 

68 - ``"discard"``: ``None`` when PII is detected; original example when 

69 no PII is found. 

70 

71 Raises: 

72 ValueError: If *action* is not ``"flag"``, ``"redact"``, or ``"discard"``. 

73 """ 

74 if action not in _VALID_ACTIONS: 

75 raise ValueError(f"Invalid pii_action {action!r}. Must be one of: {sorted(_VALID_ACTIONS)}") 

76 

77 combined = " ".join(v for v in example.values() if isinstance(v, str)) 

78 

79 if action == "discard": 

80 return None if detect_pii(combined) else example 

81 

82 if action == "flag": 

83 if detect_pii(combined): 

84 return {**example, "pii_detected": True} 

85 return example 

86 

87 # redact: replace PII in all string values 

88 return {k: redact_pii(v) if isinstance(v, str) else v for k, v in example.items()}