ash0ts commited on
Commit
0f0578b
·
1 Parent(s): 7e16d4f

add presidio model and anonymization options

Browse files
guardrails_genie/guardrails/pii/presidio_pii_guardrail.py ADDED
@@ -0,0 +1,112 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import List, Dict, Optional, ClassVar
2
+ import weave
3
+ from pydantic import BaseModel
4
+
5
+ from presidio_analyzer import AnalyzerEngine
6
+ from presidio_anonymizer import AnonymizerEngine
7
+
8
+ from ..base import Guardrail
9
+
10
+ class PresidioPIIGuardrailResponse(BaseModel):
11
+ contains_pii: bool
12
+ detected_pii_types: Dict[str, List[str]]
13
+ safe_to_process: bool
14
+ explanation: str
15
+ anonymized_text: Optional[str] = None
16
+
17
+ class PresidioPIIGuardrail(Guardrail):
18
+ AVAILABLE_ENTITIES: ClassVar[List[str]] = [
19
+ "PERSON", "EMAIL_ADDRESS", "PHONE_NUMBER", "LOCATION",
20
+ "CREDIT_CARD", "CRYPTO", "DATE_TIME", "NRP", "MEDICAL_LICENSE",
21
+ "URL", "US_BANK_NUMBER", "US_DRIVER_LICENSE", "US_ITIN",
22
+ "US_PASSPORT", "US_SSN", "UK_NHS", "IP_ADDRESS"
23
+ ]
24
+
25
+ analyzer: AnalyzerEngine
26
+ anonymizer: AnonymizerEngine
27
+ selected_entities: List[str]
28
+ should_anonymize: bool
29
+ language: str
30
+
31
+ def __init__(
32
+ self,
33
+ selected_entities: Optional[List[str]] = None,
34
+ should_anonymize: bool = False,
35
+ language: str = "en"
36
+ ):
37
+ # Initialize default values
38
+ if selected_entities is None:
39
+ selected_entities = [
40
+ "PERSON", "EMAIL_ADDRESS", "PHONE_NUMBER",
41
+ "LOCATION", "CREDIT_CARD", "US_SSN"
42
+ ]
43
+
44
+ # Validate selected entities
45
+ invalid_entities = set(selected_entities) - set(self.AVAILABLE_ENTITIES)
46
+ if invalid_entities:
47
+ raise ValueError(f"Invalid entities: {invalid_entities}")
48
+
49
+ # Initialize Presidio engines
50
+ analyzer = AnalyzerEngine()
51
+ anonymizer = AnonymizerEngine()
52
+
53
+ # Call parent class constructor with all fields
54
+ super().__init__(
55
+ analyzer=analyzer,
56
+ anonymizer=anonymizer,
57
+ selected_entities=selected_entities,
58
+ should_anonymize=should_anonymize,
59
+ language=language
60
+ )
61
+
62
+ @weave.op()
63
+ def guard(self, prompt: str, **kwargs) -> PresidioPIIGuardrailResponse:
64
+ """
65
+ Check if the input prompt contains any PII using Presidio.
66
+ """
67
+ # Analyze text for PII
68
+ analyzer_results = self.analyzer.analyze(
69
+ text=prompt,
70
+ entities=self.selected_entities,
71
+ language=self.language
72
+ )
73
+
74
+ # Group results by entity type
75
+ detected_pii = {}
76
+ for result in analyzer_results:
77
+ entity_type = result.entity_type
78
+ text_slice = prompt[result.start:result.end]
79
+ if entity_type not in detected_pii:
80
+ detected_pii[entity_type] = []
81
+ detected_pii[entity_type].append(text_slice)
82
+
83
+ # Create explanation
84
+ explanation_parts = []
85
+ if detected_pii:
86
+ explanation_parts.append("Found the following PII in the text:")
87
+ for pii_type, instances in detected_pii.items():
88
+ explanation_parts.append(f"- {pii_type}: {len(instances)} instance(s)")
89
+ else:
90
+ explanation_parts.append("No PII detected in the text.")
91
+
92
+ # Add information about what was checked
93
+ explanation_parts.append("\nChecked for these PII types:")
94
+ for entity in self.selected_entities:
95
+ explanation_parts.append(f"- {entity}")
96
+
97
+ # Anonymize if requested
98
+ anonymized_text = None
99
+ if self.should_anonymize and detected_pii:
100
+ anonymized_result = self.anonymizer.anonymize(
101
+ text=prompt,
102
+ analyzer_results=analyzer_results
103
+ )
104
+ anonymized_text = anonymized_result.text
105
+
106
+ return PresidioPIIGuardrailResponse(
107
+ contains_pii=bool(detected_pii),
108
+ detected_pii_types=detected_pii,
109
+ safe_to_process=not bool(detected_pii),
110
+ explanation="\n".join(explanation_parts),
111
+ anonymized_text=anonymized_text
112
+ )
guardrails_genie/guardrails/pii/regex_pii_guardrail.py CHANGED
@@ -12,11 +12,13 @@ class RegexPIIGuardrailResponse(BaseModel):
12
  detected_pii_types: Dict[str, list[str]]
13
  safe_to_process: bool
14
  explanation: str
 
15
 
16
 
17
  class RegexPIIGuardrail(Guardrail):
18
  regex_model: RegexModel
19
  patterns: Dict[str, str] = {}
 
20
 
21
  DEFAULT_PII_PATTERNS: ClassVar[Dict[str, str]] = {
22
  "email": r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}",
@@ -31,7 +33,7 @@ class RegexPIIGuardrail(Guardrail):
31
  "zip_code": r"\b\d{5}(?:[-]\d{4})?\b"
32
  }
33
 
34
- def __init__(self, use_defaults: bool = True, **kwargs):
35
  patterns = {}
36
  if use_defaults:
37
  patterns = self.DEFAULT_PII_PATTERNS.copy()
@@ -42,7 +44,11 @@ class RegexPIIGuardrail(Guardrail):
42
  regex_model = RegexModel(patterns=patterns)
43
 
44
  # Initialize the base class with both the regex_model and patterns
45
- super().__init__(regex_model=regex_model, patterns=patterns)
 
 
 
 
46
 
47
  @weave.op()
48
  def guard(self, prompt: str, **kwargs) -> RegexPIIGuardrailResponse:
@@ -71,9 +77,19 @@ class RegexPIIGuardrail(Guardrail):
71
  for pattern in result.failed_patterns:
72
  explanation_parts.append(f"- {pattern}")
73
 
 
 
 
 
 
 
 
 
 
74
  return RegexPIIGuardrailResponse(
75
  contains_pii=not result.passed,
76
  detected_pii_types=result.matched_patterns,
77
  safe_to_process=result.passed,
78
- explanation="\n".join(explanation_parts)
 
79
  )
 
12
  detected_pii_types: Dict[str, list[str]]
13
  safe_to_process: bool
14
  explanation: str
15
+ anonymized_text: Optional[str] = None
16
 
17
 
18
  class RegexPIIGuardrail(Guardrail):
19
  regex_model: RegexModel
20
  patterns: Dict[str, str] = {}
21
+ should_anonymize: bool = False
22
 
23
  DEFAULT_PII_PATTERNS: ClassVar[Dict[str, str]] = {
24
  "email": r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}",
 
33
  "zip_code": r"\b\d{5}(?:[-]\d{4})?\b"
34
  }
35
 
36
+ def __init__(self, use_defaults: bool = True, should_anonymize: bool = False, **kwargs):
37
  patterns = {}
38
  if use_defaults:
39
  patterns = self.DEFAULT_PII_PATTERNS.copy()
 
44
  regex_model = RegexModel(patterns=patterns)
45
 
46
  # Initialize the base class with both the regex_model and patterns
47
+ super().__init__(
48
+ regex_model=regex_model,
49
+ patterns=patterns,
50
+ should_anonymize=should_anonymize
51
+ )
52
 
53
  @weave.op()
54
  def guard(self, prompt: str, **kwargs) -> RegexPIIGuardrailResponse:
 
77
  for pattern in result.failed_patterns:
78
  explanation_parts.append(f"- {pattern}")
79
 
80
+ # Add anonymization logic
81
+ anonymized_text = None
82
+ if getattr(self, 'should_anonymize', False) and result.matched_patterns:
83
+ anonymized_text = prompt
84
+ for pii_type, matches in result.matched_patterns.items():
85
+ for match in matches:
86
+ replacement = f"[{pii_type.upper()}]"
87
+ anonymized_text = anonymized_text.replace(match, replacement)
88
+
89
  return RegexPIIGuardrailResponse(
90
  contains_pii=not result.passed,
91
  detected_pii_types=result.matched_patterns,
92
  safe_to_process=result.passed,
93
+ explanation="\n".join(explanation_parts),
94
+ anonymized_text=anonymized_text
95
  )
guardrails_genie/guardrails/pii/run_presidio_model.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from guardrails_genie.guardrails.pii.presidio_pii_guardrail import PresidioPIIGuardrail
2
+ import weave
3
+
4
+ def run_presidio_model():
5
+ weave.init("guardrails-genie-pii-presidio-model")
6
+
7
+ # Create the guardrail with default entities and anonymization enabled
8
+ pii_guardrail = PresidioPIIGuardrail(
9
+ selected_entities=["PERSON", "EMAIL_ADDRESS", "PHONE_NUMBER"],
10
+ should_anonymize=True
11
+ )
12
+
13
+ # Check a prompt
14
+ prompt = "Please contact [email protected] or call 123-456-7890. My SSN is 123-45-6789"
15
+ result = pii_guardrail.guard(prompt)
16
+ print(result)
17
+
18
+ # Result will contain:
19
+ # - contains_pii: True
20
+ # - detected_pii_types: {
21
+ # "EMAIL_ADDRESS": ["[email protected]"],
22
+ # "PHONE_NUMBER": ["123-456-7890"],
23
+ # "US_SSN": ["123-45-6789"]
24
+ # }
25
+ # - safe_to_process: False
26
+ # - explanation: Detailed explanation of findings
27
+ # - anonymized_text: "Please contact <EMAIL_ADDRESS> or call <PHONE_NUMBER>. My SSN is <US_SSN>"
28
+
29
+ # Example with no PII
30
+ safe_prompt = "The weather is nice today"
31
+ safe_result = pii_guardrail.guard(safe_prompt)
32
+ print("\nSafe prompt result:")
33
+ print(safe_result)
34
+
35
+ if __name__ == "__main__":
36
+ run_presidio_model()
guardrails_genie/guardrails/pii/run_regex_model.py CHANGED
@@ -4,7 +4,7 @@ import weave
4
  def run_regex_model():
5
  weave.init("guardrails-genie-pii-regex-model")
6
  # Create the guardrail
7
- pii_guardrail = RegexPIIGuardrail(use_defaults=True)
8
 
9
  # Check a prompt
10
  prompt = "Please contact [email protected] or call 123-456-7890"
 
4
  def run_regex_model():
5
  weave.init("guardrails-genie-pii-regex-model")
6
  # Create the guardrail
7
+ pii_guardrail = RegexPIIGuardrail(use_defaults=True, should_anonymize=True)
8
 
9
  # Check a prompt
10
  prompt = "Please contact [email protected] or call 123-456-7890"
guardrails_genie/spacy_model.py ADDED
File without changes