diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/__init__.py index d88ff434654..76c7989c5cf 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/__init__.py @@ -18,6 +18,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" preset=getattr(litellm_params, "preset", None), languages=getattr(litellm_params, "languages", None), deanonymize=getattr(litellm_params, "deanonymize", True), + block_entities=getattr(litellm_params, "block_entities", None), ) litellm.logging_callback_manager.add_litellm_callback(_privaite_callback) return _privaite_callback diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py index 635b042ca12..b0c7de9b584 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py @@ -11,6 +11,10 @@ streamed function_call. (Responses API streaming restore is not yet implemented. On the failure path the map is dropped from metadata as well, so it cannot reach a failure spend-log. +If `block_entities` is configured, a request containing any of those PII types is +rejected with an HTTP 400 in the pre-call hook, before anything is forwarded to the +model. The error names the offending type(s) only, never the underlying value. + It reuses PrivAiTe's engine, so there is no detection or masking logic here. The 'privaite' package is imported lazily so this module stays importable (for the guardrail registry) even when PrivAiTe is not installed. @@ -60,6 +64,9 @@ class PrivaiteGuardrail(CustomGuardrail): self.deanonymize = ( deanon if isinstance(deanon, bool) else str(deanon).strip().lower() not in ("false", "0", "no", "") ) + # PII TYPES to reject outright (empty = mask everything, the default). + # Accepts a list or a comma-separated string. + self.block_entities = self._parse_block_entities(kwargs.pop("block_entities", None)) if "supported_event_hooks" not in kwargs: kwargs["supported_event_hooks"] = [ GuardrailEventHooks.pre_call, @@ -92,8 +99,16 @@ class PrivaiteGuardrail(CustomGuardrail): langs = [lang.strip() for lang in self.languages.split(",") if lang.strip()] return langs or ["en"] + @staticmethod + def _parse_block_entities(raw: object) -> list: + if isinstance(raw, str): + return [e.strip() for e in raw.split(",") if e.strip()] + if isinstance(raw, (list, tuple)): + return [str(e).strip() for e in raw if str(e).strip()] + return [] + async def _engine_for(self, languages: list) -> Any: - key = (self.preset, tuple(languages), self.deanonymize) + key = (self.preset, tuple(languages), self.deanonymize, tuple(self.block_entities)) if self._engine is not None and self._engine_key == key: return self._engine @@ -116,12 +131,23 @@ class PrivaiteGuardrail(CustomGuardrail): "Install it with: pip install 'privaite>=0.2.4'" ) from exc + if self.block_entities and "block_entities" not in PIIConfig.model_fields: + # PIIConfig uses extra="allow", so an older privaite would silently + # swallow block_entities and forward the PII anyway. Fail closed + # rather than give a false sense of a policy gate that is not there. + raise RuntimeError( + "block_entities is set but the installed privaite does not " + "support it; upgrade privaite to a version that enforces " + "pii.block_entities." + ) + config = PIIConfig( enabled=True, preset=self.preset, detectors=DetectorsConfig(presidio=PresidioDetectorConfig(enabled=True, languages=languages)), anonymization=AnonymizationConfig(method="placeholder"), deanonymization=DeanonymizationConfig(enabled=self.deanonymize), + block_entities=self.block_entities, ) engine = PIIEngine(config) try: @@ -214,7 +240,16 @@ class PrivaiteGuardrail(CustomGuardrail): return data engine = await self._engine_for(self._languages()) - mapping = await self._anonymize_request(data, engine) + from privaite.pii.engine import PIIBlockedError + + try: + mapping = await self._anonymize_request(data, engine) + except PIIBlockedError as exc: + # A blocked PII type was found: reject the request outright with a 400, + # forwarding nothing. The message names TYPES only, never the values. + from fastapi import HTTPException + + raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc if mapping is not None and self.deanonymize and not mapping.is_empty: # Carry a plain fake->original dict to the post-call hook on the same diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/privaite.py b/litellm/types/proxy/guardrails/guardrail_hooks/privaite.py index 1a4db7a34a6..065cf99c37b 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/privaite.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/privaite.py @@ -18,6 +18,15 @@ class PrivaiteGuardrailConfigModel(GuardrailConfigModel): default=True, description="Restore the original PII values in the response.", ) + block_entities: str | list[str] | None = Field( + default=None, + description=( + "PII types to reject outright instead of masking: a list, or a " + "comma-separated string, e.g. ['US_SSN', 'CREDIT_CARD']. A request " + "containing any listed type is refused with a 400 and nothing is " + "forwarded. Omitted (the default) masks everything and forwards it." + ), + ) @staticmethod def ui_friendly_name() -> str: diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py index 894eea618e1..88f6476643d 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py @@ -37,12 +37,22 @@ class _FakeMapping: self._fakes[fake] = original +class _FakePIIBlockedError(Exception): + """Stand-in for privaite.pii.engine.PIIBlockedError. Names TYPES, not values.""" + + def __init__(self, entity_types): + self.entity_types = sorted(entity_types) + super().__init__("request blocked: contains disallowed PII type(s): " + ", ".join(self.entity_types)) + + class _FakeEngine: """Stand-in for privaite.pii.engine.PIIEngine. Scrubs the known reals on the way out, restores the known fakes on the way back. ``raise_oserror_until`` lets a test force the spaCy-download retry path - in the guardrail's _engine_for by failing the first N initialize() calls. + in the guardrail's _engine_for by failing the first N initialize() calls. If + ``config.block_entities`` names a type present in the text, process_request + raises _FakePIIBlockedError, exactly as the real engine would. """ raise_oserror_until = 0 @@ -82,6 +92,13 @@ class _FakeEngine: new_message = dict(message) new_message["content"] = self._scrub(message.get("content", ""), mapping) anonymized.append(new_message) + blocked = set(getattr(self.config, "block_entities", None) or []) + if blocked: + # placeholder "" -> type "EMAIL_ADDRESS" + detected = {fake.strip("<>").rsplit("_", 1)[0] for fake in mapping.get_all_fakes()} + hit = detected & blocked + if hit: + raise _FakePIIBlockedError(hit) return anonymized, mapping async def process_response(self, text, _mapping): @@ -124,6 +141,14 @@ class _Config: self.__dict__.update(kwargs) +class _PIIConfig(_Config): + """PIIConfig stand-in. ``model_fields`` mirrors a modern privaite that declares + block_entities, so the guardrail's fail-closed guard sees the field as present. + A test can override model_fields to simulate an older privaite.""" + + model_fields = {"preset": None, "languages": None, "block_entities": None} + + @pytest.fixture(autouse=True) def _fake_privaite_package(monkeypatch): """Install a fake 'privaite' (and 'spacy') package tree so the guardrail's @@ -142,11 +167,11 @@ def _fake_privaite_package(monkeypatch): AnonymizationConfig=_Config, DeanonymizationConfig=_Config, DetectorsConfig=_Config, - PIIConfig=_Config, + PIIConfig=_PIIConfig, PresidioDetectorConfig=_Config, ) mapping = _module("privaite.pii.mapping", PIIMapping=_FakeMapping) - engine = _module("privaite.pii.engine", PIIEngine=_FakeEngine) + engine = _module("privaite.pii.engine", PIIEngine=_FakeEngine, PIIBlockedError=_FakePIIBlockedError) buffer = _module("privaite.streaming.buffer", StreamingDeAnonymizer=_FakeStreamingDeAnonymizer) config_pkg = _module("privaite.config", schema=schema) pii_pkg = _module("privaite.pii", mapping=mapping, engine=engine) @@ -634,3 +659,79 @@ async def test_post_call_failure_hook_pops_map(): # missing metadata and a non-dict request_data are safe no-ops assert await gr.async_post_call_failure_hook({}, RuntimeError("x"), None) is None assert await gr.async_post_call_failure_hook(None, RuntimeError("x"), None) is None + + +def test_block_entities_config_model_field(): + from litellm.types.proxy.guardrails.guardrail_hooks.privaite import ( + PrivaiteGuardrailConfigModel, + ) + + assert PrivaiteGuardrailConfigModel.model_fields["block_entities"].default is None + + +def test_block_entities_parsed_from_string_and_list(): + assert _make_guardrail(block_entities="US_SSN, CREDIT_CARD").block_entities == [ + "US_SSN", + "CREDIT_CARD", + ] + assert _make_guardrail(block_entities=["EMAIL_ADDRESS"]).block_entities == ["EMAIL_ADDRESS"] + assert _make_guardrail().block_entities == [] + + +def test_initialize_guardrail_passes_block_entities(): + from litellm.proxy.guardrails.guardrail_hooks.privaite import initialize_guardrail + + litellm_params = types.SimpleNamespace( + mode="pre_call", + default_on=False, + preset="light", + languages="en", + deanonymize=True, + block_entities=["US_SSN"], + ) + callback = initialize_guardrail(litellm_params, {"guardrail_name": "privaite"}) + assert callback.block_entities == ["US_SSN"] + + +@pytest.mark.asyncio +async def test_block_entities_rejects_with_400(): + from fastapi import HTTPException + + gr = _make_guardrail(block_entities=["EMAIL_ADDRESS"]) + data = {"messages": [{"role": "user", "content": "reach marie@acme.com"}]} + with pytest.raises(HTTPException) as ei: + await gr.async_pre_call_hook(None, None, data, "completion") + + assert ei.value.status_code == 400 + detail = ei.value.detail + msg = detail["error"] if isinstance(detail, dict) else str(detail) + assert "EMAIL_ADDRESS" in msg + assert "marie@acme.com" not in msg # the value never leaks into the error + assert "privaite_map" not in (data.get("metadata") or {}) # nothing forwarded + + +@pytest.mark.asyncio +async def test_block_entities_ignores_types_not_present(): + # a blocked type that is absent must not disturb a normal request: the other + # PII is still masked and the request goes through with a restore map. + gr = _make_guardrail(block_entities=["US_SSN"]) + data = {"messages": [{"role": "user", "content": "I am Marie Dupont, marie@acme.com"}]} + out = await gr.async_pre_call_hook(None, None, data, "completion") + + serialized = str(out["messages"]) + assert "Marie Dupont" not in serialized + assert "marie@acme.com" not in serialized + assert out.get("metadata", {}).get("privaite_map") + + +@pytest.mark.asyncio +async def test_block_entities_fails_closed_when_privaite_too_old(monkeypatch): + # simulate an older privaite whose PIIConfig has no block_entities field: + # extra="allow" would swallow it silently, so the guardrail must refuse. + schema = sys.modules["privaite.config.schema"] + monkeypatch.setattr(schema.PIIConfig, "model_fields", {"preset": None, "languages": None}) + + gr = _make_guardrail(block_entities=["EMAIL_ADDRESS"]) + data = {"messages": [{"role": "user", "content": "marie@acme.com"}]} + with pytest.raises(RuntimeError, match="block_entities"): + await gr.async_pre_call_hook(None, None, data, "completion")