add presidio model and anonymization options

guardrails_genie/guardrails/pii/presidio_pii_guardrail.py

@@ -0,0 +1,112 @@
+from typing import List, Dict, Optional, ClassVar
+import weave
+from pydantic import BaseModel
+
+from presidio_analyzer import AnalyzerEngine
+from presidio_anonymizer import AnonymizerEngine
+
+from ..base import Guardrail
+
+class PresidioPIIGuardrailResponse(BaseModel):
+    contains_pii: bool
+    detected_pii_types: Dict[str, List[str]]
+    safe_to_process: bool
+    explanation: str
+    anonymized_text: Optional[str] = None
+
+class PresidioPIIGuardrail(Guardrail):
+    AVAILABLE_ENTITIES: ClassVar[List[str]] = [
+        "PERSON", "EMAIL_ADDRESS", "PHONE_NUMBER", "LOCATION", 
+        "CREDIT_CARD", "CRYPTO", "DATE_TIME", "NRP", "MEDICAL_LICENSE",
+        "URL", "US_BANK_NUMBER", "US_DRIVER_LICENSE", "US_ITIN", 
+        "US_PASSPORT", "US_SSN", "UK_NHS", "IP_ADDRESS"
+    ]
+    
+    analyzer: AnalyzerEngine
+    anonymizer: AnonymizerEngine
+    selected_entities: List[str]
+    should_anonymize: bool
+    language: str
+    
+    def __init__(
+        self,
+        selected_entities: Optional[List[str]] = None,
+        should_anonymize: bool = False,
+        language: str = "en"
+    ):
+        # Initialize default values
+        if selected_entities is None:
+            selected_entities = [
+                "PERSON", "EMAIL_ADDRESS", "PHONE_NUMBER", 
+                "LOCATION", "CREDIT_CARD", "US_SSN"
+            ]
+        
+        # Validate selected entities
+        invalid_entities = set(selected_entities) - set(self.AVAILABLE_ENTITIES)
+        if invalid_entities:
+            raise ValueError(f"Invalid entities: {invalid_entities}")
+            
+        # Initialize Presidio engines
+        analyzer = AnalyzerEngine()
+        anonymizer = AnonymizerEngine()
+        
+        # Call parent class constructor with all fields
+        super().__init__(
+            analyzer=analyzer,
+            anonymizer=anonymizer,
+            selected_entities=selected_entities,
+            should_anonymize=should_anonymize,
+            language=language
+        )
+
+    @weave.op()
+    def guard(self, prompt: str, **kwargs) -> PresidioPIIGuardrailResponse:
+        """
+        Check if the input prompt contains any PII using Presidio.
+        """
+        # Analyze text for PII
+        analyzer_results = self.analyzer.analyze(
+            text=prompt,
+            entities=self.selected_entities,
+            language=self.language
+        )
+        
+        # Group results by entity type
+        detected_pii = {}
+        for result in analyzer_results:
+            entity_type = result.entity_type
+            text_slice = prompt[result.start:result.end]
+            if entity_type not in detected_pii:
+                detected_pii[entity_type] = []
+            detected_pii[entity_type].append(text_slice)
+        
+        # Create explanation
+        explanation_parts = []
+        if detected_pii:
+            explanation_parts.append("Found the following PII in the text:")
+            for pii_type, instances in detected_pii.items():
+                explanation_parts.append(f"- {pii_type}: {len(instances)} instance(s)")
+        else:
+            explanation_parts.append("No PII detected in the text.")
+            
+        # Add information about what was checked
+        explanation_parts.append("\nChecked for these PII types:")
+        for entity in self.selected_entities:
+            explanation_parts.append(f"- {entity}")
+        
+        # Anonymize if requested
+        anonymized_text = None
+        if self.should_anonymize and detected_pii:
+            anonymized_result = self.anonymizer.anonymize(
+                text=prompt,
+                analyzer_results=analyzer_results
+            )
+            anonymized_text = anonymized_result.text
+            
+        return PresidioPIIGuardrailResponse(
+            contains_pii=bool(detected_pii),
+            detected_pii_types=detected_pii,
+            safe_to_process=not bool(detected_pii),
+            explanation="\n".join(explanation_parts),
+            anonymized_text=anonymized_text
+        )

guardrails_genie/guardrails/pii/regex_pii_guardrail.py

@@ -12,11 +12,13 @@ class RegexPIIGuardrailResponse(BaseModel):
     detected_pii_types: Dict[str, list[str]]
     safe_to_process: bool
     explanation: str
+    anonymized_text: Optional[str] = None
 
 
 class RegexPIIGuardrail(Guardrail):
     regex_model: RegexModel
     patterns: Dict[str, str] = {}
+    should_anonymize: bool = False
     
     DEFAULT_PII_PATTERNS: ClassVar[Dict[str, str]] = {
         "email": r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}",
@@ -31,7 +33,7 @@ class RegexPIIGuardrail(Guardrail):
         "zip_code": r"\b\d{5}(?:[-]\d{4})?\b"
     }
     
-    def __init__(self, use_defaults: bool = True, **kwargs):
+    def __init__(self, use_defaults: bool = True, should_anonymize: bool = False, **kwargs):
         patterns = {}
         if use_defaults:
             patterns = self.DEFAULT_PII_PATTERNS.copy()
@@ -42,7 +44,11 @@ class RegexPIIGuardrail(Guardrail):
         regex_model = RegexModel(patterns=patterns)
         
         # Initialize the base class with both the regex_model and patterns
-        super().__init__(regex_model=regex_model, patterns=patterns)
+        super().__init__(
+            regex_model=regex_model, 
+            patterns=patterns,
+            should_anonymize=should_anonymize
+        )
 
     @weave.op()
     def guard(self, prompt: str, **kwargs) -> RegexPIIGuardrailResponse:
@@ -71,9 +77,19 @@ class RegexPIIGuardrail(Guardrail):
             for pattern in result.failed_patterns:
                 explanation_parts.append(f"- {pattern}")
                 
+        # Add anonymization logic
+        anonymized_text = None
+        if getattr(self, 'should_anonymize', False) and result.matched_patterns:
+            anonymized_text = prompt
+            for pii_type, matches in result.matched_patterns.items():
+                for match in matches:
+                    replacement = f"[{pii_type.upper()}]"
+                    anonymized_text = anonymized_text.replace(match, replacement)
+                
         return RegexPIIGuardrailResponse(
             contains_pii=not result.passed,
             detected_pii_types=result.matched_patterns,
             safe_to_process=result.passed,
-            explanation="\n".join(explanation_parts)
+            explanation="\n".join(explanation_parts),
+            anonymized_text=anonymized_text
         )

guardrails_genie/guardrails/pii/run_presidio_model.py

@@ -0,0 +1,36 @@
+from guardrails_genie.guardrails.pii.presidio_pii_guardrail import PresidioPIIGuardrail
+import weave
+
+def run_presidio_model():
+    weave.init("guardrails-genie-pii-presidio-model")
+    
+    # Create the guardrail with default entities and anonymization enabled
+    pii_guardrail = PresidioPIIGuardrail(
+        selected_entities=["PERSON", "EMAIL_ADDRESS", "PHONE_NUMBER"],
+        should_anonymize=True
+    )
+
+    # Check a prompt
+    prompt = "Please contact [email protected] or call 123-456-7890. My SSN is 123-45-6789"
+    result = pii_guardrail.guard(prompt)
+    print(result)
+
+    # Result will contain:
+    # - contains_pii: True
+    # - detected_pii_types: {
+    #     "EMAIL_ADDRESS": ["[email protected]"],
+    #     "PHONE_NUMBER": ["123-456-7890"],
+    #     "US_SSN": ["123-45-6789"]
+    # }
+    # - safe_to_process: False
+    # - explanation: Detailed explanation of findings
+    # - anonymized_text: "Please contact <EMAIL_ADDRESS> or call <PHONE_NUMBER>. My SSN is <US_SSN>"
+
+    # Example with no PII
+    safe_prompt = "The weather is nice today"
+    safe_result = pii_guardrail.guard(safe_prompt)
+    print("\nSafe prompt result:")
+    print(safe_result)
+
+if __name__ == "__main__":
+    run_presidio_model()

guardrails_genie/guardrails/pii/run_regex_model.py

@@ -4,7 +4,7 @@ import weave
 def run_regex_model():
     weave.init("guardrails-genie-pii-regex-model")
     # Create the guardrail
-    pii_guardrail = RegexPIIGuardrail(use_defaults=True)
+    pii_guardrail = RegexPIIGuardrail(use_defaults=True, should_anonymize=True)
 
     # Check a prompt
     prompt = "Please contact [email protected] or call 123-456-7890"