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..1caa314 100644 --- a/backend/app/core/validators/README.md +++ b/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,7 +539,7 @@ Notes / limitations: ## Example Config Payloads -Example: create validator config (stored shape) — spaCy mode (default) +Example: create validator config (stored shape) ```json { @@ -547,23 +547,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 +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"] } ``` 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..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,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(), ) diff --git a/backend/app/core/validators/pii_remover.py b/backend/app/core/validators/pii_remover.py index cd25d89..9b71877 100644 --- a/backend/app/core/validators/pii_remover.py +++ b/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 .""" 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/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 0518c03..4dd83a9 100644 --- a/backend/app/tests/validators/test_pii_remover.py +++ b/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()