Merge pull request #14899 from arsh72/fix/presidio-custom-entities-union-type

Fix: Support custom entity types in Presidio guardrail with Union[PiiEntityType, str]

litellm/proxy/guardrails/guardrail_hooks/presidio.py

@@ -68,7 +68,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
         output_parse_pii: Optional[bool] = False,
         presidio_ad_hoc_recognizers: Optional[str] = None,
         logging_only: Optional[bool] = None,
-        pii_entities_config: Optional[Dict[PiiEntityType, PiiAction]] = None,
+        pii_entities_config: Optional[Dict[Union[PiiEntityType, str], PiiAction]] = None,
         presidio_language: Optional[str] = None,
         **kwargs,
     ):
@@ -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,
@@ -402,10 +402,14 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
             content_safety = data.get("content_safety", None)
             verbose_proxy_logger.debug("content_safety: %s", content_safety)
             presidio_config = self.get_presidio_settings_from_request_data(data)
-            if call_type in [
-                LitellmCallTypes.completion.value,
-                LitellmCallTypes.acompletion.value,
-            ] or call_type == "mcp_call":
+            if (
+                call_type
+                in [
+                    LitellmCallTypes.completion.value,
+                    LitellmCallTypes.acompletion.value,
+                ]
+                or call_type == "mcp_call"
+            ):
                 messages = data["messages"]
                 tasks = []
                 for m in messages:

litellm/types/guardrails.py

@@ -40,7 +40,6 @@ class SupportedGuardrailIntegrations(Enum):
     TOOL_PERMISSION = "tool_permission"
 
 
-
 class Role(Enum):
     SYSTEM = "system"
     ASSISTANT = "assistant"
@@ -253,7 +252,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(

litellm/types/proxy/guardrails/guardrail_hooks/presidio.py

@@ -9,7 +9,7 @@ class PresidioAnalyzeRequest(TypedDict, total=False):
     text: str
     language: Optional[str]
     ad_hoc_recognizers: Optional[List[str]]
-    entities: Optional[List[PiiEntityType]]
+    entities: Optional[List[Union[PiiEntityType, str]]]
 
 
 class PresidioAnalyzeResponseItem(TypedDict, total=False):

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

@@ -0,0 +1,39 @@
+"""
+Minimal test for Presidio Union[PiiEntityType, str] type fix.
+Tests only the core fix without heavy dependencies.
+"""
+from enum import Enum
+from typing import Dict, Union
+
+
+class PiiEntityType(str, Enum):
+    EMAIL_ADDRESS = "EMAIL_ADDRESS"
+
+
+class PiiAction(str, Enum):
+    BLOCK = "BLOCK"
+    MASK = "MASK"
+
+
+def test_presidio_union_type_fix() -> None:
+    """Test that Union[PiiEntityType, str] allows both enum and string entity types"""
+
+    # 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
+
+    # Union type fix verified: mixed entity types work correctly
+
+
+if __name__ == "__main__":
+    test_presidio_union_type_fix()