From aaf469165259a3f1bc47715a6f95177759146032 Mon Sep 17 00:00:00 2001 From: AkhileshNegi Date: Wed, 14 Oct 2026 15:46:59 +0530 Subject: [PATCH 1/4] feat(pii): restrict PII config to entity_types enum with pattern-based default - Add PIIEntityType enum; unknown entity types are now rejected. - Default entity_types to the 12 pattern-based types; NER-based PERSON, LOCATION and NRP are opt-in. - Drop threshold, nlp_engine_type and model_name from the API config; they remain PIIRemover constructor params for evals. - Fix placeholder format in docstring and update README. Co-Authored-By: Claude Opus 5.5 --- backend/app/core/enum.py | 18 +++++++++++ backend/app/core/validators/README.md | 31 ++++-------------- .../pii_remover_safety_validator_config.py | 15 +++------ backend/app/core/validators/pii_remover.py | 31 +++++++----------- .../app/tests/validators/test_pii_remover.py | 32 +++++++++++-------- 5 files changed, 61 insertions(+), 66 deletions(-) diff --git a/backend/app/core/enum.py b/backend/app/core/enum.py index fee1a64..55b568b 100644 --- a/backend/app/core/enum.py +++ b/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" diff --git a/backend/app/core/validators/README.md b/backend/app/core/validators/README.md index 85e7a52..938922f 100644 --- a/backend/app/core/validators/README.md +++ b/backend/app/core/validators/README.md @@ -175,8 +175,9 @@ 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`, `nlp_engine_type`, `model_name` are `PIIRemover` constructor params only (used by evals); not accepted via the API. +- `threshold: float` (default: `0.5` for spaCy, `0.7` for transformers). 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 +191,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,7 +540,7 @@ Notes / limitations: ## Example Config Payloads -Example: create validator config (stored shape) — spaCy mode (default) +Example: create validator config (stored shape) ```json { @@ -547,23 +548,7 @@ Example: create validator config (stored shape) — spaCy mode (default) "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 +558,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"] } ``` diff --git a/backend/app/core/validators/config/pii_remover_safety_validator_config.py b/backend/app/core/validators/config/pii_remover_safety_validator_config.py index bd7fcd6..4523920 100644 --- a/backend/app/core/validators/config/pii_remover_safety_validator_config.py +++ b/backend/app/core/validators/config/pii_remover_safety_validator_config.py @@ -1,22 +1,17 @@ from __future__ import annotations -from typing import List, Literal, Optional +from typing import List, Literal -from app.core.validators.pii_remover import PIIRemover +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): 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] = [PIIEntityType(t) for t in 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, + entity_types=[e.value for e in self.entity_types], on_fail=self.resolve_on_fail(), ) diff --git a/backend/app/core/validators/pii_remover.py b/backend/app/core/validators/pii_remover.py index cd25d89..dd27b96 100644 --- a/backend/app/core/validators/pii_remover.py +++ b/backend/app/core/validators/pii_remover.py @@ -30,25 +30,16 @@ 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", -] +ALL_ENTITY_TYPES = [e.value for e in PIIEntityType] + +# 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 = [t for t in ALL_ENTITY_TYPES if t not in NER_ENTITY_TYPES] SPACY_CONFIGURATION = { "nlp_engine_name": "spacy", @@ -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 .""" 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( diff --git a/backend/app/tests/validators/test_pii_remover.py b/backend/app/tests/validators/test_pii_remover.py index 0518c03..3dfb898 100644 --- a/backend/app/tests/validators/test_pii_remover.py +++ b/backend/app/tests/validators/test_pii_remover.py @@ -1,10 +1,11 @@ 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, PIIRemover, @@ -78,7 +79,23 @@ 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 {"PERSON", "LOCATION", "NRP"} & 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_custom_entity_types_override(mock_presidio): @@ -106,17 +123,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() From 8c8bcb7cb5ae028a632279d4c818be8d347fd39a Mon Sep 17 00:00:00 2001 From: AkhileshNegi Date: Wed, 14 Oct 2026 15:57:00 +0530 Subject: [PATCH 2/4] refactor(pii): use enum values directly, drop ALL_ENTITY_TYPES - use_enum_values on PII config so entity_types are plain strings; removes enum<->str round trip in default and build() - fold ALL_ENTITY_TYPES into DEFAULT_ENTITY_TYPES - merge duplicate threshold bullets in README - test uses NER_ENTITY_TYPES instead of hardcoded set Co-Authored-By: Claude Opus 5.5 --- backend/app/core/validators/README.md | 3 +-- .../config/pii_remover_safety_validator_config.py | 8 ++++++-- backend/app/core/validators/pii_remover.py | 4 +--- backend/app/tests/validators/test_pii_remover.py | 3 ++- 4 files changed, 10 insertions(+), 8 deletions(-) diff --git a/backend/app/core/validators/README.md b/backend/app/core/validators/README.md index 938922f..1caa314 100644 --- a/backend/app/core/validators/README.md +++ b/backend/app/core/validators/README.md @@ -176,8 +176,7 @@ Recommendation: Parameters / customization: - `entity_types: list[PIIEntityType]` (default: all pattern-based types; NER-based `PERSON`, `LOCATION`, `NRP` are opt-in). Unknown values are rejected. -- `threshold`, `nlp_engine_type`, `model_name` are `PIIRemover` constructor params only (used by evals); not accepted via the API. -- `threshold: float` (default: `0.5` for spaCy, `0.7` for transformers). Only effective in transformers mode — spaCy scores every NER hit `0.85`, and pattern recognizers are not thresholded. +- `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` diff --git a/backend/app/core/validators/config/pii_remover_safety_validator_config.py b/backend/app/core/validators/config/pii_remover_safety_validator_config.py index 4523920..fbc7e7a 100644 --- a/backend/app/core/validators/config/pii_remover_safety_validator_config.py +++ b/backend/app/core/validators/config/pii_remover_safety_validator_config.py @@ -1,17 +1,21 @@ from __future__ import annotations from typing import List, Literal +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: List[PIIEntityType] = [PIIEntityType(t) for t in DEFAULT_ENTITY_TYPES] + entity_types: List[PIIEntityType] = DEFAULT_ENTITY_TYPES def build(self): return PIIRemover( - entity_types=[e.value for e in self.entity_types], + entity_types=self.entity_types, on_fail=self.resolve_on_fail(), ) diff --git a/backend/app/core/validators/pii_remover.py b/backend/app/core/validators/pii_remover.py index dd27b96..233897f 100644 --- a/backend/app/core/validators/pii_remover.py +++ b/backend/app/core/validators/pii_remover.py @@ -34,12 +34,10 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false" -ALL_ENTITY_TYPES = [e.value for e in PIIEntityType] - # 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 = [t for t in ALL_ENTITY_TYPES if t not in NER_ENTITY_TYPES] +DEFAULT_ENTITY_TYPES = [e.value for e in PIIEntityType if e.value not in NER_ENTITY_TYPES] SPACY_CONFIGURATION = { "nlp_engine_name": "spacy", diff --git a/backend/app/tests/validators/test_pii_remover.py b/backend/app/tests/validators/test_pii_remover.py index 3dfb898..a80fd9f 100644 --- a/backend/app/tests/validators/test_pii_remover.py +++ b/backend/app/tests/validators/test_pii_remover.py @@ -8,6 +8,7 @@ DEFAULT_ENTITY_TYPES, DEFAULT_TRANSFORMERS_MODEL, DEFAULT_TRANSFORMERS_THRESHOLD, + NER_ENTITY_TYPES, PIIRemover, ) from app.core.validators.config.pii_remover_safety_validator_config import ( @@ -80,7 +81,7 @@ def test_analyzer_called_with_correct_arguments(validator): def test_default_entity_types_applied(validator): assert validator.entity_types == DEFAULT_ENTITY_TYPES - assert not {"PERSON", "LOCATION", "NRP"} & set(DEFAULT_ENTITY_TYPES) + assert not NER_ENTITY_TYPES & set(DEFAULT_ENTITY_TYPES) def test_config_default_entity_types_passed_as_strings(mock_presidio): From 6b69d9d97fd949d6004a9fd78d91b49a49814907 Mon Sep 17 00:00:00 2001 From: AkhileshNegi Date: Wed, 14 Oct 2026 16:13:16 +0530 Subject: [PATCH 3/4] fix(pii): validate pii_remover config at create time ValidatorCreate allows extra fields, so typos like entity_types=["AADHAR"] were saved and only failed on the first guardrail run. Validate PII configs against PIIRemoverSafetyValidatorConfig before saving. Co-Authored-By: Claude Opus 5.5 --- backend/app/schemas/validator_config.py | 13 +++++++++++++ backend/app/tests/validators/test_pii_remover.py | 8 ++++++++ 2 files changed, 21 insertions(+) diff --git a/backend/app/schemas/validator_config.py b/backend/app/schemas/validator_config.py index 9fc8920..0a2ab30 100644 --- a/backend/app/schemas/validator_config.py +++ b/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") diff --git a/backend/app/tests/validators/test_pii_remover.py b/backend/app/tests/validators/test_pii_remover.py index a80fd9f..4dd83a9 100644 --- a/backend/app/tests/validators/test_pii_remover.py +++ b/backend/app/tests/validators/test_pii_remover.py @@ -11,6 +11,7 @@ NER_ENTITY_TYPES, PIIRemover, ) +from app.schemas.validator_config import ValidatorCreate from app.core.validators.config.pii_remover_safety_validator_config import ( PIIRemoverSafetyValidatorConfig, ) @@ -99,6 +100,13 @@ def test_config_rejects_backend_params(): 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): v = PIIRemover(entity_types=["EMAIL_ADDRESS"], threshold=0.5) From 8e24a1c7fed9b5e9b827d32c7b0d3cf09985e81f Mon Sep 17 00:00:00 2001 From: AkhileshNegi Date: Wed, 14 Oct 2026 16:50:41 +0530 Subject: [PATCH 4/4] style(pii): apply black formatting Co-Authored-By: Claude Opus 5.5 --- backend/app/core/validators/pii_remover.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/backend/app/core/validators/pii_remover.py b/backend/app/core/validators/pii_remover.py index 233897f..9b71877 100644 --- a/backend/app/core/validators/pii_remover.py +++ b/backend/app/core/validators/pii_remover.py @@ -37,7 +37,9 @@ # 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] +DEFAULT_ENTITY_TYPES = [ + e.value for e in PIIEntityType if e.value not in NER_ENTITY_TYPES +] SPACY_CONFIGURATION = { "nlp_engine_name": "spacy",