fix(presidio): enable custom PII entities with Union type

litellm/proxy/guardrails/guardrail_hooks/presidio.py

@@ -82,7 +82,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
         )  # mapping of PII token to original text - only used with Presidio `replace` operation
         self.mock_redacted_text = mock_redacted_text
         self.output_parse_pii = output_parse_pii or False
-        self.pii_entities_config: Dict[PiiEntityType, PiiAction] = (
+        self.pii_entities_config: Dict[Union[PiiEntityType, str], PiiAction] = (
             pii_entities_config or {}
         )
         self.presidio_language = presidio_language or "en"
@@ -302,10 +302,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
             entity_type = result.get("entity_type")
 
             if entity_type:
-                casted_entity_type: PiiEntityType = cast(PiiEntityType, entity_type)
+                # Check if entity_type is in config (supports both enum and string)
                 if (
-                    casted_entity_type in self.pii_entities_config
-                    and self.pii_entities_config[casted_entity_type] == PiiAction.BLOCK
+                    entity_type in self.pii_entities_config
+                    and self.pii_entities_config[entity_type] == PiiAction.BLOCK
                 ):
                     raise BlockedPiiEntityError(
                         entity_type=entity_type,

litellm/types/guardrails.py

@@ -253,7 +253,7 @@ class PresidioPresidioConfigModelUserInterface(BaseModel):
 class PresidioConfigModel(PresidioPresidioConfigModelUserInterface):
     """Configuration parameters for the Presidio PII masking guardrail"""
 
-    pii_entities_config: Optional[Dict[PiiEntityType, PiiAction]] = Field(
+    pii_entities_config: Optional[Dict[Union[PiiEntityType, str], PiiAction]] = Field(
         default=None, description="Configuration for PII entity types and actions"
     )
     presidio_ad_hoc_recognizers: Optional[str] = Field(

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio_union_fix.py

@@ -0,0 +1,58 @@
+"""
+Minimal test for Presidio Union[PiiEntityType, str] type fix.
+Tests only the core fix without heavy dependencies.
+"""
+from typing import Union, Dict
+from enum import Enum
+
+
+class PiiEntityType(str, Enum):
+    EMAIL_ADDRESS = "EMAIL_ADDRESS"
+
+
+class PiiAction(str, Enum):
+    BLOCK = "BLOCK"
+    MASK = "MASK"
+
+
+def test_presidio_union_type_fix():
+    """Test that Union[PiiEntityType, str] allows both enum and string entity types"""
+    
+    # Custom recognizers for EMPLOYEE_ID and CUSTOMER_ID
+    custom_recognizers = [
+        {
+            "name": "Employee ID Recognizer",
+            "supported_language": "en",
+            "patterns": [{"name": "employee id", "regex": "EMP-[0-9]{6}", "score": 0.9}],
+            "context": ["employee", "id"],
+            "supported_entity": "EMPLOYEE_ID"
+        },
+        {
+            "name": "Customer ID Recognizer",
+            "supported_language": "en",
+            "patterns": [{"name": "customer id", "regex": "CUST-[0-9]{8}", "score": 0.9}],
+            "context": ["customer", "id"],
+            "supported_entity": "CUSTOMER_ID"
+        }
+    ]
+    
+    # This is the core fix - mixed entity types in pii_entities_config
+    pii_entities_config: Dict[Union[PiiEntityType, str], PiiAction] = {
+        PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK,
+        "EMPLOYEE_ID": PiiAction.MASK,
+        "CUSTOMER_ID": PiiAction.BLOCK,
+    }
+    
+    # Verify entities can be used together (what Presidio needs)
+    entities_list = list(pii_entities_config.keys())
+    assert len(entities_list) == 3
+    assert PiiEntityType.EMAIL_ADDRESS in entities_list
+    assert "EMPLOYEE_ID" in entities_list
+    assert "CUSTOMER_ID" in entities_list
+    
+    print("✅ Union type fix verified: mixed entity types work correctly")
+    print("✅ Custom recognizers defined for EMPLOYEE_ID and CUSTOMER_ID")
+
+
+if __name__ == "__main__":
+    test_presidio_union_type_fix()