From 82ec42993f5cf0cb3a06f63b61c517a9c6e8fa4e Mon Sep 17 00:00:00 2001
From: Arshdeep Singh <asingh1@marqeta.com>
Date: Thu, 25 Sep 2025 11:27:40 -0400
Subject: [PATCH] fix(presidio): enable custom PII entities with Union type

---
 .../guardrails/guardrail_hooks/presidio.py    |  8 +--
 litellm/types/guardrails.py                   |  2 +-
 .../test_presidio_union_fix.py                | 58 +++++++++++++++++++
 3 files changed, 63 insertions(+), 5 deletions(-)
 create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio_union_fix.py

diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py
index 8dbcf77a84..c3c6c3d2e0 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py
+++ b/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,
diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py
index e5ea89325e..0c3dc662ec 100644
--- a/litellm/types/guardrails.py
+++ b/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(
diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio_union_fix.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio_union_fix.py
new file mode 100644
index 0000000000..549f715ed5
--- /dev/null
+++ b/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()
\ No newline at end of file