PII Remover: Restrict config to entity_types enum with pattern-based defaults

backend/app/core/enum.py

@@ -20,6 +20,24 @@ class BiasCategories(Enum):
     All = "all"
 
 
+class PIIEntityType(str, Enum):
+    CreditCard = "CREDIT_CARD"
+    EmailAddress = "EMAIL_ADDRESS"
+    IbanCode = "IBAN_CODE"
+    IpAddress = "IP_ADDRESS"
+    Location = "LOCATION"
+    MedicalLicense = "MEDICAL_LICENSE"
+    Nrp = "NRP"
+    Person = "PERSON"
+    PhoneNumber = "PHONE_NUMBER"
+    Url = "URL"
+    InAadhaar = "IN_AADHAAR"
+    InPan = "IN_PAN"
+    InPassport = "IN_PASSPORT"
+    InVehicleRegistration = "IN_VEHICLE_REGISTRATION"
+    InVoter = "IN_VOTER"
+
+
 class GuardrailOnFail(Enum):
     Exception = "exception"
     Fix = "fix"

backend/app/core/validators/README.md

@@ -175,8 +175,8 @@ Recommendation:
 
 Parameters / customization:
 
-- `entity_types: list[str] | None` (default: all supported types)
-- `threshold: float` (default: `0.5` for spaCy, `0.7` for transformers)
+- `entity_types: list[PIIEntityType]` (default: all pattern-based types; NER-based `PERSON`, `LOCATION`, `NRP` are opt-in). Unknown values are rejected.
+- `threshold: float` (default: `0.5` for spaCy, `0.7` for transformers). Constructor-only (evals), like `nlp_engine_type` and `model_name` below; not accepted via the API. Only effective in transformers mode — spaCy scores every NER hit `0.85`, and pattern recognizers are not thresholded.
 - `nlp_engine_type: str` (default: `"spacy"`) — NLP backend to use. `"spacy"` uses a spaCy model for both tokenization and NER; `"transformers"` uses a spaCy model for tokenization only and a HuggingFace token-classification model for NER.
 - `model_name: str | None` (default: `None`) — Override the NLP model. For `"spacy"`, a spaCy model name (resolved to `en_core_web_lg`); for `"transformers"`, a HuggingFace model ID (resolved to `dslim/bert-base-NER-uncased`).
 - `on_fail`
@@ -190,7 +190,7 @@ Threshold guidance:
 - Tune using real conversation samples by reviewing both missed PII and unnecessary masking.
 - If the product is privacy-critical, prefer a slightly lower threshold and tighter `entity_types`; if readability is primary, prefer a slightly higher threshold.
 
-Supported default entity types:
+Supported entity types (default excludes `LOCATION`, `NRP`, `PERSON`):
 
 - `CREDIT_CARD`, `EMAIL_ADDRESS`, `IBAN_CODE`, `IP_ADDRESS`, `LOCATION`, `MEDICAL_LICENSE`, `NRP`, `PERSON`, `PHONE_NUMBER`, `URL`, `IN_AADHAAR`, `IN_PAN`, `IN_PASSPORT`, `IN_VEHICLE_REGISTRATION`, `IN_VOTER`
 
@@ -539,31 +539,15 @@ Notes / limitations:
 
 ## Example Config Payloads
 
-Example: create validator config (stored shape) — spaCy mode (default)
+Example: create validator config (stored shape)
 
 ```json
 {
   "type": "pii_remover",
   "stage": "input",
   "on_fail_action": "fix",
   "is_enabled": true,
-  "entity_types": ["PERSON", "PHONE_NUMBER", "IN_AADHAAR"],
-  "threshold": 0.6
-}
-```
-
-Example: create validator config — transformers mode with a custom HuggingFace model
-
-```json
-{
-  "type": "pii_remover",
-  "stage": "input",
-  "on_fail_action": "fix",
-  "is_enabled": true,
-  "entity_types": ["PERSON", "PHONE_NUMBER", "LOCATION"],
-  "nlp_engine_type": "transformers",
-  "model_name": "dslim/bert-base-NER-uncased",
-  "threshold": 0.7
+  "entity_types": ["PERSON", "PHONE_NUMBER", "IN_AADHAAR"]
 }
 ```
 
@@ -573,9 +557,7 @@ Example: runtime guardrail validator object (execution shape)
 {
   "type": "pii_remover",
   "on_fail": "fix",
-  "entity_types": ["PERSON", "PHONE_NUMBER", "IN_AADHAAR"],
-  "nlp_engine_type": "spacy",
-  "threshold": 0.6
+  "entity_types": ["PERSON", "PHONE_NUMBER", "IN_AADHAAR"]
 }
 ```
 

backend/app/core/validators/config/pii_remover_safety_validator_config.py

@@ -1,22 +1,21 @@
 from __future__ import annotations
-from typing import List, Literal, Optional
+from typing import List, Literal
 
-from app.core.validators.pii_remover import PIIRemover
+from pydantic import ConfigDict
+
+from app.core.enum import PIIEntityType
+from app.core.validators.pii_remover import DEFAULT_ENTITY_TYPES, PIIRemover
 from app.core.validators.config.base_validator_config import BaseValidatorConfig
 
 
 class PIIRemoverSafetyValidatorConfig(BaseValidatorConfig):
+    model_config = ConfigDict(use_enum_values=True)
+
     type: Literal["pii_remover"]
-    entity_types: Optional[List[str]] = None
-    threshold: Optional[float] = None
-    nlp_engine_type: str = "spacy"
-    model_name: Optional[str] = None
+    entity_types: List[PIIEntityType] = DEFAULT_ENTITY_TYPES
 
     def build(self):
         return PIIRemover(
             entity_types=self.entity_types,
-            threshold=self.threshold,
-            nlp_engine_type=self.nlp_engine_type,
-            model_name=self.model_name,
             on_fail=self.resolve_on_fail(),
         )

backend/app/core/validators/pii_remover.py

@@ -30,24 +30,15 @@
     InVoterRecognizer,
 )
 
+from app.core.enum import PIIEntityType
+
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
-ALL_ENTITY_TYPES = [
-    "CREDIT_CARD",
-    "EMAIL_ADDRESS",
-    "IBAN_CODE",
-    "IP_ADDRESS",
-    "LOCATION",
-    "MEDICAL_LICENSE",
-    "NRP",
-    "PERSON",
-    "PHONE_NUMBER",
-    "URL",
-    "IN_AADHAAR",
-    "IN_PAN",
-    "IN_PASSPORT",
-    "IN_VEHICLE_REGISTRATION",
-    "IN_VOTER",
+# NER-based types (PERSON, LOCATION, NRP) are opt-in: the NER model is English-only
+# and noisy on Indic names. Everything else is pattern/regex based.
+NER_ENTITY_TYPES = {"PERSON", "LOCATION", "NRP"}
+DEFAULT_ENTITY_TYPES = [
+    e.value for e in PIIEntityType if e.value not in NER_ENTITY_TYPES
 ]
 
 SPACY_CONFIGURATION = {
@@ -190,7 +181,7 @@ def _get_cached_analyzer(
 
 @register_validator(name="pii-remover", data_type="string")
 class PIIRemover(Validator):
-    """Anonymizes detected PII entities in the text with placeholders like [REDACTED_PERSON_1]."""
+    """Anonymizes detected PII entities in the text with placeholders like <PERSON>."""
 
     def __init__(
         self,
@@ -202,7 +193,7 @@ def __init__(
     ):
         super().__init__(on_fail=on_fail)
 
-        self.entity_types = entity_types or ALL_ENTITY_TYPES
+        self.entity_types = entity_types or DEFAULT_ENTITY_TYPES
         self.nlp_engine_type = nlp_engine_type
         if nlp_engine_type == "transformers":
             self.model_name = model_name or DEFAULT_TRANSFORMERS_MODEL
@@ -211,6 +202,8 @@ def __init__(
             )
         else:
             self.model_name = model_name or "en_core_web_lg"
+            # ponytail: no-op in spaCy mode — SpacyRecognizer scores every NER hit 0.85
+            # and regex recognizers aren't thresholded. Only matters for transformers.
             self.threshold = threshold if threshold is not None else 0.5
         self.on_fail = on_fail
         self.analyzer = _get_cached_analyzer(

backend/app/schemas/validator_config.py

@@ -2,6 +2,9 @@
 from sqlmodel import Field, SQLModel
 
 from app.core.enum import GuardrailOnFail, Stage, ValidatorType
+from app.core.validators.config.pii_remover_safety_validator_config import (
+    PIIRemoverSafetyValidatorConfig,
+)
 
 
 class ValidatorBase(SQLModel):
@@ -30,6 +33,16 @@ def reject_body_tenant(cls, data: object) -> object:
                 )
         return data
 
+    @model_validator(mode="after")
+    def validate_pii_config(self):
+        # Reject bad entity_types at save time, not first guardrail run.
+        # ponytail: PII only; route other types through their configs when needed.
+        if self.type == ValidatorType.PIIRemover:
+            PIIRemoverSafetyValidatorConfig.model_validate(
+                {"type": "pii_remover", **(self.model_extra or {})}
+            )
+        return self
+
 
 class ValidatorUpdate(SQLModel):
     model_config = ConfigDict(extra="forbid")

backend/app/tests/validators/test_pii_remover.py

@@ -1,14 +1,17 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
+from pydantic import ValidationError
 
 from app.core.validators import pii_remover
 from app.core.validators.pii_remover import (
-    ALL_ENTITY_TYPES,
+    DEFAULT_ENTITY_TYPES,
     DEFAULT_TRANSFORMERS_MODEL,
     DEFAULT_TRANSFORMERS_THRESHOLD,
+    NER_ENTITY_TYPES,
     PIIRemover,
 )
+from app.schemas.validator_config import ValidatorCreate
 from app.core.validators.config.pii_remover_safety_validator_config import (
     PIIRemoverSafetyValidatorConfig,
 )
@@ -78,7 +81,30 @@ def test_analyzer_called_with_correct_arguments(validator):
 
 
 def test_default_entity_types_applied(validator):
-    assert validator.entity_types == ALL_ENTITY_TYPES
+    assert validator.entity_types == DEFAULT_ENTITY_TYPES
+    assert not NER_ENTITY_TYPES & set(DEFAULT_ENTITY_TYPES)
+
+
+def test_config_default_entity_types_passed_as_strings(mock_presidio):
+    v = PIIRemoverSafetyValidatorConfig(type="pii_remover").build()
+    assert v.entity_types == DEFAULT_ENTITY_TYPES
+
+
+def test_config_rejects_unknown_entity_type():
+    with pytest.raises(ValidationError):
+        PIIRemoverSafetyValidatorConfig(type="pii_remover", entity_types=["AADHAR"])
+
+
+def test_config_rejects_backend_params():
+    with pytest.raises(ValidationError):
+        PIIRemoverSafetyValidatorConfig(type="pii_remover", threshold=0.5)
+
+
+def test_create_payload_rejects_unknown_entity_type():
+    base = {"name": "pii-check", "type": "pii_remover", "stage": "input"}
+    ValidatorCreate(**base, entity_types=["PERSON", "IN_AADHAAR"])
+    with pytest.raises(ValidationError):
+        ValidatorCreate(**base, entity_types=["AADHAR"])
 
 
 def test_custom_entity_types_override(mock_presidio):
@@ -106,17 +132,6 @@ def test_transformers_engine_accepts_custom_model(mock_presidio):
     assert v.model_name == "dslim/bert-base-NER-uncased"
 
 
-def test_config_builds_validator_with_nlp_engine_params(mock_presidio):
-    config = PIIRemoverSafetyValidatorConfig(
-        type="pii_remover",
-        nlp_engine_type="transformers",
-        model_name="dslim/bert-base-NER-uncased",
-    )
-    v = config.build()
-    assert v.nlp_engine_type == "transformers"
-    assert v.model_name == "dslim/bert-base-NER-uncased"
-
-
 def test_config_defaults_to_spacy_engine(mock_presidio):
     config = PIIRemoverSafetyValidatorConfig(type="pii_remover")
     v = config.build()