workign PII guardrails in chat_app

application_pages/chat_app.py

@@ -61,6 +61,27 @@ def initialize_guardrails():
                         guardrail_name,
                     )(model_name=classifier_model_name)
                 )
+        elif guardrail_name == "PresidioEntityRecognitionGuardrail":
+            st.session_state.guardrails.append(
+                getattr(
+                    importlib.import_module("guardrails_genie.guardrails"),
+                    guardrail_name,
+                )()
+            )
+        elif guardrail_name == "RegexEntityRecognitionGuardrail":
+            st.session_state.guardrails.append(
+                getattr(
+                    importlib.import_module("guardrails_genie.guardrails"),
+                    guardrail_name,
+                )()
+            )
+        elif guardrail_name == "TransformersEntityRecognitionGuardrail":
+            st.session_state.guardrails.append(
+                getattr(
+                    importlib.import_module("guardrails_genie.guardrails"),
+                    guardrail_name,
+                )()
+            )
     st.session_state.guardrails_manager = GuardrailManager(
         guardrails=st.session_state.guardrails
     )

guardrails_genie/guardrails/__init__.py

@@ -2,10 +2,18 @@ from .injection import (
     PromptInjectionClassifierGuardrail,
     PromptInjectionSurveyGuardrail,
 )
+from .entity_recognition import (
+    PresidioEntityRecognitionGuardrail,
+    RegexEntityRecognitionGuardrail,
+    TransformersEntityRecognitionGuardrail,
+)
 from .manager import GuardrailManager
 
 __all__ = [
     "PromptInjectionSurveyGuardrail",
     "PromptInjectionClassifierGuardrail",
+    "PresidioEntityRecognitionGuardrail",
+    "RegexEntityRecognitionGuardrail",
+    "TransformersEntityRecognitionGuardrail",
     "GuardrailManager",
 ]

guardrails_genie/guardrails/entity_recognition/__init__.py

@@ -0,0 +1,9 @@
+from .presidio_entity_recognition_guardrail import PresidioEntityRecognitionGuardrail 
+from .regex_entity_recognition_guardrail import RegexEntityRecognitionGuardrail
+from .transformers_entity_recognition_guardrail import TransformersEntityRecognitionGuardrail
+
+__all__ = [
+    "PresidioEntityRecognitionGuardrail",
+    "RegexEntityRecognitionGuardrail",
+    "TransformersEntityRecognitionGuardrail",
+]

guardrails_genie/guardrails/entity_recognition/pii_examples/run_transformers.py

@@ -8,7 +8,7 @@ def test_pii_detection():
     
     # Create the guardrail with default entities and anonymization enabled
     pii_guardrail = TransformersEntityRecognitionGuardrail(
-        selected_entities=["GIVENNAME", "SURNAME", "EMAIL", "PHONE_NUMBER", "SOCIALNUM"],
+        selected_entities=["GIVENNAME", "SURNAME", "EMAIL", "TELEPHONENUM", "SOCIALNUM"],
         should_anonymize=True,
         show_available_entities=True
     )

guardrails_genie/guardrails/entity_recognition/presidio_entity_recognition_guardrail.py

@@ -13,11 +13,19 @@ class PresidioEntityRecognitionResponse(BaseModel):
     explanation: str
     anonymized_text: Optional[str] = None
 
+    @property
+    def safe(self) -> bool:
+        return not self.contains_entities
+
 class PresidioEntityRecognitionSimpleResponse(BaseModel):
     contains_entities: bool
     explanation: str
     anonymized_text: Optional[str] = None
 
+    @property
+    def safe(self) -> bool:
+        return not self.contains_entities
+
 #TODO: Add support for transformers workflow and not just Spacy
 class PresidioEntityRecognitionGuardrail(Guardrail):
     @staticmethod
@@ -40,23 +48,37 @@ class PresidioEntityRecognitionGuardrail(Guardrail):
         language: str = "en",
         deny_lists: Optional[Dict[str, List[str]]] = None,
         regex_patterns: Optional[Dict[str, List[Dict[str, str]]]] = None,
-        custom_recognizers: Optional[List[Any]] = None
+        custom_recognizers: Optional[List[Any]] = None,
+        show_available_entities: bool = False
     ):
+        # If show_available_entities is True, print available entities
+        if show_available_entities:
+            available_entities = self.get_available_entities()
+            print("\nAvailable entities:")
+            print("=" * 25)
+            for entity in available_entities:
+                print(f"- {entity}")
+            print("=" * 25 + "\n")
+
         # Initialize default values
         if selected_entities is None:
             selected_entities = [
-                "PERSON", "EMAIL_ADDRESS", "PHONE_NUMBER", 
-                "LOCATION", "CREDIT_CARD", "US_SSN"
+                "CREDIT_CARD", "US_SSN", "EMAIL_ADDRESS", "PHONE_NUMBER",
+                "IP_ADDRESS", "URL", "DATE_TIME"
             ]
         
         # Get available entities dynamically
         available_entities = self.get_available_entities()
         
-        # Validate selected entities
-        invalid_entities = set(selected_entities) - set(available_entities)
+        # Filter out invalid entities and warn user
+        invalid_entities = [e for e in selected_entities if e not in available_entities]
+        valid_entities = [e for e in selected_entities if e in available_entities]
+        
         if invalid_entities:
-            raise ValueError(f"Invalid entities: {invalid_entities}")
-            
+            print(f"\nWarning: The following entities are not available and will be ignored: {invalid_entities}")
+            print(f"Continuing with valid entities: {valid_entities}")
+            selected_entities = valid_entities
+        
         # Initialize analyzer with default recognizers
         analyzer = AnalyzerEngine()
         

guardrails_genie/guardrails/entity_recognition/regex_entity_recognition_guardrail.py

@@ -13,12 +13,20 @@ class RegexEntityRecognitionResponse(BaseModel):
     explanation: str
     anonymized_text: Optional[str] = None
 
+    @property
+    def safe(self) -> bool:
+        return not self.contains_entities
+
 
 class RegexEntityRecognitionSimpleResponse(BaseModel):
     contains_entities: bool
     explanation: str
     anonymized_text: Optional[str] = None
 
+    @property
+    def safe(self) -> bool:
+        return not self.contains_entities
+
 
 class RegexEntityRecognitionGuardrail(Guardrail):
     regex_model: RegexModel

guardrails_genie/guardrails/entity_recognition/transformers_entity_recognition_guardrail.py

@@ -11,11 +11,19 @@ class TransformersEntityRecognitionResponse(BaseModel):
     explanation: str
     anonymized_text: Optional[str] = None
 
+    @property
+    def safe(self) -> bool:
+        return not self.contains_entities
+
 class TransformersEntityRecognitionSimpleResponse(BaseModel):
     contains_entities: bool
     explanation: str
     anonymized_text: Optional[str] = None
 
+    @property
+    def safe(self) -> bool:
+        return not self.contains_entities
+
 class TransformersEntityRecognitionGuardrail(Guardrail):
     """Generic guardrail for detecting entities using any token classification model."""
     
@@ -126,9 +134,12 @@ class TransformersEntityRecognitionGuardrail(Guardrail):
                 # Replace the entity with the redaction marker
                 chars[start:end] = replacement
         
-        # Join and clean up multiple spaces
+        # Join characters and clean up only consecutive spaces (preserving newlines)
         result = ''.join(chars)
-        return ' '.join(result.split())
+        # Replace multiple spaces with single space, but preserve newlines
+        lines = result.split('\n')
+        cleaned_lines = [' '.join(line.split()) for line in lines]
+        return '\n'.join(cleaned_lines)
 
     @weave.op()
     def guard(self, prompt: str, return_detected_types: bool = True, aggregate_redaction: bool = True) -> TransformersEntityRecognitionResponse | TransformersEntityRecognitionSimpleResponse:

guardrails_genie/guardrails/manager.py

@@ -1,5 +1,6 @@
 import weave
 from rich.progress import track
+from pydantic import BaseModel
 
 from .base import Guardrail
 
@@ -20,10 +21,12 @@ class GuardrailManager(weave.Model):
             alerts.append(
                 {"guardrail_name": guardrail.__class__.__name__, "response": response}
             )
-            safe = safe and response["safe"]
-            summaries += (
-                f"**{guardrail.__class__.__name__}**: {response['summary']}\n\n---\n\n"
-            )
+            if isinstance(response, BaseModel):
+                safe = safe and response.safe
+                summaries += f"**{guardrail.__class__.__name__}**: {response.explanation}\n\n---\n\n"
+            else:
+                safe = safe and response["safe"]
+                summaries += f"**{guardrail.__class__.__name__}**: {response['summary']}\n\n---\n\n"
         return {"safe": safe, "alerts": alerts, "summary": summaries}
 
     @weave.op()

pyproject.toml

@@ -20,6 +20,8 @@ dependencies = [
     "pymupdf4llm>=0.0.17",
     "transformers>=4.46.3",
     "torch>=2.5.1",
+    "presidio-analyzer>=2.2.355",
+    "presidio-anonymizer>=2.2.355",
 ]
 
 [tool.setuptools]