From 483d57ee75056c839de09dacb2be614983875477 Mon Sep 17 00:00:00 2001 From: crp4222 <121295348+crp4222@users.noreply.github.com> Date: Sat, 27 Jun 2026 23:31:19 +0100 Subject: [PATCH] feat(guardrails): add PrivAiTe guardrail provider PrivAiTe (https://github.com/crp4222/PrivAiTe) is a local PII anonymization engine. This guardrail runs it in-process: the pre-call hook anonymizes the request and stashes a reversible map in request metadata, and the post-call hook restores the original values in the response, including inside tool-call arguments and the legacy function_call. Streaming responses are restored too. The 'privaite' package is imported lazily, so the module stays importable for the guardrail registry even when PrivAiTe is not installed. Adds the privaite ConfigModel (preset/languages/deanonymize), the PRIVAITE enum value, registry wiring, and unit tests that fake the engine so they run without the external package. --- .../guardrail_hooks/privaite/__init__.py | 32 + .../guardrail_hooks/privaite/privaite.py | 362 ++++++++++ .../spend_tracking/spend_tracking_utils.py | 2 +- litellm/types/guardrails.py | 5 + .../guardrails/guardrail_hooks/privaite.py | 24 + .../guardrail_hooks/test_privaite.py | 636 ++++++++++++++++++ 6 files changed, 1060 insertions(+), 1 deletion(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/privaite/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/privaite.py create mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/__init__.py new file mode 100644 index 00000000000..d88ff434654 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/__init__.py @@ -0,0 +1,32 @@ +from typing import TYPE_CHECKING + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .privaite import PrivaiteGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): + import litellm + + _privaite_callback = PrivaiteGuardrail( + guardrail_name=guardrail.get("guardrail_name", ""), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + preset=getattr(litellm_params, "preset", None), + languages=getattr(litellm_params, "languages", None), + deanonymize=getattr(litellm_params, "deanonymize", True), + ) + litellm.logging_callback_manager.add_litellm_callback(_privaite_callback) + return _privaite_callback + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.PRIVAITE.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.PRIVAITE.value: PrivaiteGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py new file mode 100644 index 00000000000..635b042ca12 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py @@ -0,0 +1,362 @@ +""" +PrivAiTe guardrail for the LiteLLM proxy. + +Runs PrivAiTe's engine in-process inside LiteLLM. The pre-call hook anonymizes the +request (chat `messages` and Responses API `input`) and stashes the reversible map +in the request metadata (consumed and popped by the post-call hook); the post-call +hook restores the real values in the response, including chat tool-call arguments, +the legacy function_call, and Responses API output_text/function_call output. Chat +streaming is restored too: text content, streamed tool-call arguments, and the +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. + +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. + +See https://github.com/crp4222/PrivAiTe for the package. +""" + +from __future__ import annotations + +import asyncio +from typing import Any, cast + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.types.guardrails import GuardrailEventHooks + +_LANG_MODELS = { + "en": "en_core_web_lg", + "fr": "fr_core_news_md", + "de": "de_core_news_md", + "es": "es_core_news_md", + "it": "it_core_news_md", + "pt": "pt_core_news_md", + "nl": "nl_core_news_md", +} + + +def _obj_get(obj: Any, key: str) -> Any: + """Read a field from a dict-or-object response item.""" + return obj.get(key) if isinstance(obj, dict) else getattr(obj, key, None) + + +def _obj_set(obj: Any, key: str, value: Any) -> None: + """Write a field on a dict-or-object response item.""" + if isinstance(obj, dict): + obj[key] = value + else: + setattr(obj, key, value) + + +class PrivaiteGuardrail(CustomGuardrail): + def __init__(self, **kwargs: Any) -> None: + self.preset = kwargs.pop("preset", None) or "onnx" + if self.preset not in ("light", "onnx"): + self.preset = "onnx" + self.languages = kwargs.pop("languages", None) or "en,fr" + deanon = kwargs.pop("deanonymize", True) + self.deanonymize = ( + deanon if isinstance(deanon, bool) else str(deanon).strip().lower() not in ("false", "0", "no", "") + ) + if "supported_event_hooks" not in kwargs: + kwargs["supported_event_hooks"] = [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ] + super().__init__(**kwargs) + # Both hooks are mandatory: pre_call anonymizes the request and post_call + # restores the response. A config such as `mode: post_call` would skip the + # pre_call pass and silently forward raw PII to the model, so ensure both + # hooks run for plain string/list modes. Tag-based Mode configs are left + # untouched. + hook = self.event_hook + if hook is None or isinstance(hook, str): + base = [hook] if isinstance(hook, str) and hook else [] + self.event_hook = cast( + list[GuardrailEventHooks], + list(dict.fromkeys(base + ["pre_call", "post_call"])), + ) + elif isinstance(hook, list): + normalized = [getattr(h, "value", h) for h in hook] + self.event_hook = cast( + list[GuardrailEventHooks], + list(dict.fromkeys(normalized + ["pre_call", "post_call"])), + ) + self._engine: Any = None + self._engine_key: Any = None + self._lock = asyncio.Lock() + + def _languages(self) -> list: + langs = [lang.strip() for lang in self.languages.split(",") if lang.strip()] + return langs or ["en"] + + async def _engine_for(self, languages: list) -> Any: + key = (self.preset, tuple(languages), self.deanonymize) + if self._engine is not None and self._engine_key == key: + return self._engine + + async with self._lock: + if self._engine is not None and self._engine_key == key: + return self._engine # pragma: no cover + + try: + from privaite.config.schema import ( + AnonymizationConfig, + DeanonymizationConfig, + DetectorsConfig, + PIIConfig, + PresidioDetectorConfig, + ) + from privaite.pii.engine import PIIEngine + except ImportError as exc: # pragma: no cover + raise RuntimeError( + "The PrivAiTe guardrail requires the 'privaite' package. " + "Install it with: pip install 'privaite>=0.2.4'" + ) from exc + + config = PIIConfig( + enabled=True, + preset=self.preset, + detectors=DetectorsConfig(presidio=PresidioDetectorConfig(enabled=True, languages=languages)), + anonymization=AnonymizationConfig(method="placeholder"), + deanonymization=DeanonymizationConfig(enabled=self.deanonymize), + ) + engine = PIIEngine(config) + try: + await engine.initialize() + except OSError: + # spaCy models not present yet: download them once, then retry. + from spacy.cli import download + + for lang in languages: + model = _LANG_MODELS.get(lang) + if model: + download(model) + engine = PIIEngine(config) + await engine.initialize() + + self._engine = engine + self._engine_key = key + return engine + + @staticmethod + def _is_role_message_list(value: list) -> bool: + return all(isinstance(item, dict) and "role" in item for item in value) + + def _overwrite_snapshot_input(self, data: dict, new_input: Any) -> None: + # A plain `data["input"] = ...` rebind leaks for string input: the proxy + # snapshots the request body by shallow-copying data before this hook, + # so the original string stays in proxy_server_request.body["input"]. + # Overwrite the snapshot copy too. (List inputs are mutated in place, so + # the aliased snapshot already reflects the anonymized values.) + psr = data.get("proxy_server_request") + body = psr.get("body") if isinstance(psr, dict) else None + if isinstance(body, dict) and "input" in body: + body["input"] = new_input + + def _input_as_messages(self, input_value: Any) -> tuple: + """Represent the Responses API `input` as message dicts for the engine. + Returns (kind, [messages]) where kind drives write-back.""" + if isinstance(input_value, str) and input_value: + return "str", [{"role": "user", "content": input_value}] + if isinstance(input_value, list) and input_value: + if self._is_role_message_list(input_value): + return "role_list", list(input_value) + return "content_list", [{"role": "user", "content": input_value}] + return "none", [] + + def _write_back_input(self, data: dict, kind: str, input_value: Any, anonymized: list) -> None: + if kind == "str": + new_text = anonymized[0].get("content", input_value) + data["input"] = new_text + self._overwrite_snapshot_input(data, new_text) + elif kind == "role_list": + input_value[:] = anonymized + elif kind == "content_list": + new_content = anonymized[0].get("content") + if isinstance(new_content, list): + input_value[:] = new_content + + async def _anonymize_request(self, data: dict, engine: Any) -> Any: + """Anonymize chat `messages` AND Responses `input` in place using the + engine (span-precise), sharing ONE mapping so neither is left untouched + when a crafted request carries both. Returns the mapping, or None if + there was nothing to anonymize.""" + messages = data.get("messages") + msg_list = messages if isinstance(messages, list) else [] + input_kind, input_msgs = self._input_as_messages(data.get("input")) + + if not msg_list and not input_msgs: + return None + + batch = list(msg_list) + input_msgs + anonymized, mapping = await engine.process_request(batch) + + n = len(msg_list) + if msg_list: + # msg_list is data["messages"] (same object) -> mutate it in place. + msg_list[:] = anonymized[:n] + self._write_back_input(data, input_kind, data.get("input"), anonymized[n:]) + return mapping + + async def async_pre_call_hook(self, user_api_key_dict: Any, cache: Any, data: dict, call_type: str) -> dict: + # metadata is caller-controlled at the proxy boundary, so never trust an + # incoming privaite_map: this hook is the only authority that may set it. + # Clearing it here means the post-call hook can only ever restore from a + # map produced by this guardrail's pre-call pass on the same request. + metadata = data.get("metadata") + if isinstance(metadata, dict): + metadata.pop("privaite_map", None) + + if not data.get("messages") and not data.get("input"): + return data + + engine = await self._engine_for(self._languages()) + mapping = await self._anonymize_request(data, engine) + + 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 + # request, the same channel LiteLLM's Presidio guardrail uses. The + # post-call hook pops it again so it does not linger in metadata. + data.setdefault("metadata", {})["privaite_map"] = dict(mapping.get_all_fakes()) + return data + + async def _restore_message(self, message: Any, engine: Any, mapping: Any) -> None: + """Restore originals in one response message: content, tool-call args and + the legacy function_call args.""" + content = getattr(message, "content", None) + if isinstance(content, str) and content: + message.content = await engine.process_response(content, mapping) + for tool_call in getattr(message, "tool_calls", None) or []: + fn = getattr(tool_call, "function", None) + if fn is None: + continue + args = getattr(fn, "arguments", None) + if args: + fn.arguments = await engine.process_response(args, mapping) + function_call = getattr(message, "function_call", None) + if function_call is not None: + fc_args = getattr(function_call, "arguments", None) + if fc_args: + function_call.arguments = await engine.process_response(fc_args, mapping) + + async def _restore_responses_output(self, response: Any, engine: Any, mapping: Any) -> None: + """Restore originals in a Responses API result: output_text content + blocks and any function_call output-item arguments (dict or object).""" + for item in getattr(response, "output", None) or []: + for block in _obj_get(item, "content") or []: + text = _obj_get(block, "text") + if isinstance(text, str) and text: + _obj_set(block, "text", await engine.process_response(text, mapping)) + args = _obj_get(item, "arguments") + if isinstance(args, str) and args: + _obj_set(item, "arguments", await engine.process_response(args, mapping)) + + async def async_post_call_success_hook(self, data: dict, user_api_key_dict: Any, response: Any) -> Any: + if not self.deanonymize: + return response + # pop (not get): consume the reversible map so the originals do not linger + # in request metadata, which the proxy may persist to spend logs. + fakes = (data.get("metadata") or {}).pop("privaite_map", None) + if not fakes: + return response + + from privaite.pii.mapping import PIIMapping + + mapping = PIIMapping() + for fake, original in fakes.items(): + mapping.add(original, fake, "PII") + + engine = await self._engine_for(self._languages()) + for choice in getattr(response, "choices", None) or []: + message = getattr(choice, "message", None) + if message is not None: + await self._restore_message(message, engine, mapping) + # Responses API results carry text under `output`, not `choices`. + await self._restore_responses_output(response, engine, mapping) + return response + + def _restore_delta(self, delta: Any, index: int, finished: bool, restore) -> None: + """Restore one streamed delta in place: text content, streamed tool-call + argument fragments (per tool_call index) and the legacy function_call.""" + content = getattr(delta, "content", None) or "" + restored = restore(("content", index), content, finished) + if content or finished: + delta.content = restored + for tool_call in getattr(delta, "tool_calls", None) or []: + fn = getattr(tool_call, "function", None) + if fn is None: + continue + args = getattr(fn, "arguments", None) + if args: + tc_index = getattr(tool_call, "index", 0) or 0 + fn.arguments = restore(("tool", index, tc_index), args, False) + function_call = getattr(delta, "function_call", None) + if function_call is not None: + fc_args = getattr(function_call, "arguments", None) + if fc_args: + function_call.arguments = restore(("fc", index), fc_args, False) + + async def async_post_call_streaming_iterator_hook( + self, user_api_key_dict: Any, response: Any, request_data: dict + ) -> Any: + # pop (not get): consume the map so the originals do not linger in metadata. + fakes = (request_data.get("metadata") or {}).pop("privaite_map", None) + if not self.deanonymize or not fakes: + async for chunk in response: + yield chunk + return + + from privaite.pii.mapping import PIIMapping + from privaite.streaming.buffer import StreamingDeAnonymizer + + mapping = PIIMapping() + for fake, original in fakes.items(): + mapping.add(original, fake, "PII") + + # One de-anonymizer buffer per streamed segment, keyed by (kind, choice + # index, ...). With n>1 the provider interleaves chunks for different + # choices, and tool-call arguments stream as fragments per tool_call + # index; each segment keeps its own boundary buffer so a placeholder split + # across chunks reassembles without mixing segments. + buffers: dict = {} + + def _restore(key: tuple, text: str, finished: bool) -> str: + deanon = buffers.get(key) + if deanon is None: + deanon = buffers[key] = StreamingDeAnonymizer(mapping) + out = deanon.feed(text) if text else "" + if finished: + out += deanon.flush() + return out + + # The buffer holds back partial placeholders that span chunk boundaries. + # Whatever remains is flushed onto the chunk that carries finish_reason. + async for chunk in response: + for choice in getattr(chunk, "choices", None) or []: + delta = getattr(choice, "delta", None) + if delta is None: + continue + index = getattr(choice, "index", 0) or 0 + finished = getattr(choice, "finish_reason", None) is not None + self._restore_delta(delta, index, finished, _restore) + yield chunk + + async def async_post_call_failure_hook( + self, + request_data: Any, + original_exception: Any, + user_api_key_dict: Any, + traceback_str: Any = None, + ) -> Any: + # The success/streaming hooks pop the reversible map after restoring, but + # they never run when the call fails. Drop it here too so the originals + # are not left in metadata for a failure spend-log to persist. (LiteLLM + # core also strips "privaite_map" from the spend-log body; this is the + # defense for a stock install that predates that change.) + metadata = request_data.get("metadata") if isinstance(request_data, dict) else None + if isinstance(metadata, dict): + metadata.pop("privaite_map", None) + return None diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 94b079163e9..2e2beb4e5ab 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -1210,7 +1210,7 @@ def _get_messages_for_spend_logs_payload( return "{}" -_SENSITIVE_REQUEST_BODY_KEYS: Final = frozenset({"secret_fields"}) +_SENSITIVE_REQUEST_BODY_KEYS: Final = frozenset({"secret_fields", "privaite_map"}) _REQUEST_BODY_CREDENTIAL_MASKER: Final = SensitiveDataMasker(extra_sensitive_patterns=frozenset({"apikey"})) _TOOL_INPUT_BLOCK_TYPES: Final = frozenset({"tool_use", "server_tool_use", "mcp_tool_use"}) _TOOL_OUTPUT_BLOCK_TYPES: Final = frozenset({"tool_result", "mcp_tool_result", "function_call_output"}) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index ff063d2e233..0780023a56f 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -77,6 +77,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.vigil_guard import ( from litellm.types.proxy.guardrails.guardrail_hooks.wingback import ( WingbackGuardrailConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.privaite import ( + PrivaiteGuardrailConfigModel, +) from litellm.types.proxy.guardrails.guardrail_hooks.xecguard import ( XecGuardConfigModel, ) @@ -159,6 +162,7 @@ class SupportedGuardrailIntegrations(Enum): AIRIA = "airia" WINGBACK = "wingback" DECISION_MODEL = "decision_model" + PRIVAITE = "privaite" class Role(Enum): @@ -1400,6 +1404,7 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o SingulrGuardrailConfigModel, Agent365GuardrailConfigModel, WingbackGuardrailConfigModel, + PrivaiteGuardrailConfigModel, ): guardrail: str = Field(description="The type of guardrail integration to use") decision_model: str | None = Field( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/privaite.py b/litellm/types/proxy/guardrails/guardrail_hooks/privaite.py new file mode 100644 index 00000000000..1a4db7a34a6 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/privaite.py @@ -0,0 +1,24 @@ +from typing import Optional + +from pydantic import Field + +from .base import GuardrailConfigModel + + +class PrivaiteGuardrailConfigModel(GuardrailConfigModel): + preset: Optional[str] = Field( + default="onnx", + description=("Detection preset: 'onnx' (full suite, also detects secrets) or 'light' (faster, Presidio only)."), + ) + languages: Optional[str] = Field( + default="en,fr", + description="Comma-separated spaCy languages used for detection.", + ) + deanonymize: Optional[bool] = Field( + default=True, + description="Restore the original PII values in the response.", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "PrivAiTe" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py new file mode 100644 index 00000000000..894eea618e1 --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py @@ -0,0 +1,636 @@ +""" +Unit tests for the PrivAiTe guardrail provider. + +The PrivAiTe engine lives in the external 'privaite' package, which is NOT a +dependency of LiteLLM. These tests therefore install a fake 'privaite' package +into sys.modules (an autouse fixture) so the guardrail's lazy imports resolve to +fakes. That means the suite runs exactly the way LiteLLM CI runs it: with the +real 'privaite' package absent. The fakes let us exercise the guardrail's own +wiring end to end (registry, engine caching, pre/post/streaming hooks) without +re-testing PrivAiTe's detection internals. +""" + +import sys +import types + +import pytest + +# fake placeholder <-> original, the way PrivAiTe maps values. +_FAKES = {"": "Marie Dupont", "": "marie@acme.com"} +_REALS = {original: fake for fake, original in _FAKES.items()} + + +class _FakeMapping: + """Stand-in for privaite.pii.mapping.PIIMapping.""" + + def __init__(self, fakes=None): + self._fakes = dict(fakes or {}) + + @property + def is_empty(self): + return not self._fakes + + def get_all_fakes(self): + return dict(self._fakes) + + def add(self, original, fake, _entity_type): + self._fakes[fake] = original + + +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. + """ + + raise_oserror_until = 0 + init_calls = 0 + + def __init__(self, config): + self.config = config + + async def initialize(self): + type(self).init_calls += 1 + if self.init_calls <= type(self).raise_oserror_until: + raise OSError("spaCy model missing") + + def _scrub(self, content, mapping): + if isinstance(content, str): + for real, fake in _REALS.items(): + if real in content: + content = content.replace(real, fake) + mapping.add(real, fake, "PII") + return content + if isinstance(content, list): + out = [] + for part in content: + if isinstance(part, str): + out.append(self._scrub(part, mapping)) + elif isinstance(part, dict) and isinstance(part.get("text"), str): + out.append({**part, "text": self._scrub(part["text"], mapping)}) + else: + out.append(part) + return out + return content + + async def process_request(self, messages): + mapping = _FakeMapping() + anonymized = [] + for message in messages: + new_message = dict(message) + new_message["content"] = self._scrub(message.get("content", ""), mapping) + anonymized.append(new_message) + return anonymized, mapping + + async def process_response(self, text, _mapping): + for fake, original in _FAKES.items(): + text = text.replace(fake, original) + return text + + +class _FakeStreamingDeAnonymizer: + """Stand-in for privaite.streaming.buffer.StreamingDeAnonymizer. + + Stateful like the real one: it buffers a trailing "<...."-without-">" fragment + so a placeholder split across chunks reassembles, which is what makes per-choice + buffer isolation observable in tests. + """ + + def __init__(self, _mapping): + self._buf = "" + + def feed(self, text): + self._buf += text + for fake, original in _FAKES.items(): + self._buf = self._buf.replace(fake, original) + lt = self._buf.rfind("<") + if lt != -1 and ">" not in self._buf[lt:]: + emit, self._buf = self._buf[:lt], self._buf[lt:] + else: + emit, self._buf = self._buf, "" + return emit + + def flush(self): + out, self._buf = self._buf, "" + return out + + +class _Config: + """Generic stand-in for the privaite.config.schema dataclasses.""" + + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + +@pytest.fixture(autouse=True) +def _fake_privaite_package(monkeypatch): + """Install a fake 'privaite' (and 'spacy') package tree so the guardrail's + lazy imports resolve without the real packages installed, mirroring CI.""" + _FakeEngine.raise_oserror_until = 0 + _FakeEngine.init_calls = 0 + + def _module(name, **attrs): + mod = types.ModuleType(name) + for key, value in attrs.items(): + setattr(mod, key, value) + return mod + + schema = _module( + "privaite.config.schema", + AnonymizationConfig=_Config, + DeanonymizationConfig=_Config, + DetectorsConfig=_Config, + PIIConfig=_Config, + PresidioDetectorConfig=_Config, + ) + mapping = _module("privaite.pii.mapping", PIIMapping=_FakeMapping) + engine = _module("privaite.pii.engine", PIIEngine=_FakeEngine) + buffer = _module("privaite.streaming.buffer", StreamingDeAnonymizer=_FakeStreamingDeAnonymizer) + config_pkg = _module("privaite.config", schema=schema) + pii_pkg = _module("privaite.pii", mapping=mapping, engine=engine) + streaming_pkg = _module("privaite.streaming", buffer=buffer) + privaite_pkg = _module("privaite", config=config_pkg, pii=pii_pkg, streaming=streaming_pkg) + spacy_cli = _module("spacy.cli", download=lambda _model: None) + spacy_pkg = _module("spacy", cli=spacy_cli) + + for name, mod in { + "privaite": privaite_pkg, + "privaite.config": config_pkg, + "privaite.config.schema": schema, + "privaite.pii": pii_pkg, + "privaite.pii.mapping": mapping, + "privaite.pii.engine": engine, + "privaite.streaming": streaming_pkg, + "privaite.streaming.buffer": buffer, + "spacy": spacy_pkg, + "spacy.cli": spacy_cli, + }.items(): + monkeypatch.setitem(sys.modules, name, mod) + yield + + +def _make_guardrail(**kwargs): + from litellm.proxy.guardrails.guardrail_hooks.privaite import PrivaiteGuardrail + + return PrivaiteGuardrail(guardrail_name="privaite-test", **kwargs) + + +async def _collect(aiterator): + return [chunk async for chunk in aiterator] + + +def test_registry_and_enum_wiring(): + from litellm.proxy.guardrails.guardrail_hooks.privaite import ( + PrivaiteGuardrail, + guardrail_class_registry, + guardrail_initializer_registry, + ) + from litellm.types.guardrails import SupportedGuardrailIntegrations + + assert SupportedGuardrailIntegrations.PRIVAITE.value == "privaite" + assert "privaite" in guardrail_initializer_registry + assert guardrail_class_registry["privaite"] is PrivaiteGuardrail + + +def test_config_model_ui_name(): + from litellm.types.proxy.guardrails.guardrail_hooks.privaite import ( + PrivaiteGuardrailConfigModel, + ) + + assert PrivaiteGuardrailConfigModel.ui_friendly_name() == "PrivAiTe" + fields = PrivaiteGuardrailConfigModel.model_fields + assert fields["preset"].default == "onnx" + assert fields["languages"].default == "en,fr" + assert fields["deanonymize"].default is True + + +def test_initialize_guardrail_registers_callback(monkeypatch): + import litellm + from litellm.proxy.guardrails.guardrail_hooks.privaite import ( + PrivaiteGuardrail, + initialize_guardrail, + ) + + added = [] + monkeypatch.setattr( + litellm.logging_callback_manager, + "add_litellm_callback", + lambda cb: added.append(cb), + ) + + litellm_params = types.SimpleNamespace( + mode="pre_call", + default_on=False, + preset="light", + languages="en", + deanonymize=True, + ) + callback = initialize_guardrail(litellm_params, {"guardrail_name": "privaite"}) + + assert isinstance(callback, PrivaiteGuardrail) + assert callback.preset == "light" + assert added == [callback] + + +def test_init_normalizes_params(): + # invalid preset falls back to onnx; "false"-ish strings disable deanonymize. + gr = _make_guardrail(preset="bogus", languages="en, fr ,", deanonymize="false") + assert gr.preset == "onnx" + assert gr.deanonymize is False + assert gr._languages() == ["en", "fr"] + + # an explicit valid preset and bool are kept; empty languages fall back to + # the __init__ default of "en,fr". + gr2 = _make_guardrail(preset="light", languages="", deanonymize=True) + assert gr2.preset == "light" + assert gr2.deanonymize is True + assert gr2._languages() == ["en", "fr"] + + # a non-empty but content-free languages string falls back to ["en"]. + assert _make_guardrail(languages=", ,")._languages() == ["en"] + + +def test_event_hook_always_includes_pre_and_post(): + # A `mode: post_call` config must NOT disable pre_call anonymization. + gr = _make_guardrail(event_hook="post_call") + assert "pre_call" in gr.event_hook and "post_call" in gr.event_hook + # A list missing post_call gets it added (and vice-versa). + gr2 = _make_guardrail(event_hook=["pre_call"]) + assert "pre_call" in gr2.event_hook and "post_call" in gr2.event_hook + # The default (no event_hook configured) still runs both hooks. + gr3 = _make_guardrail() + assert "pre_call" in gr3.event_hook and "post_call" in gr3.event_hook + + +@pytest.mark.asyncio +async def test_engine_is_cached(): + gr = _make_guardrail() + engine_a = await gr._engine_for(["en"]) + engine_b = await gr._engine_for(["en"]) + assert engine_a is engine_b + + +@pytest.mark.asyncio +async def test_engine_downloads_spacy_models_on_oserror(): + _FakeEngine.raise_oserror_until = 1 # first initialize() raises, retry succeeds + gr = _make_guardrail(languages="en,fr") + engine = await gr._engine_for(["en", "fr"]) + assert isinstance(engine, _FakeEngine) + assert _FakeEngine.init_calls == 2 # failed once, then succeeded after download + + +@pytest.mark.asyncio +async def test_pre_call_no_messages_is_passthrough(): + gr = _make_guardrail() + data = {"messages": []} + assert await gr.async_pre_call_hook(None, None, data, "completion") is data + + +@pytest.mark.asyncio +async def test_pre_call_anonymizes_text_and_stashes_map(): + gr = _make_guardrail() + messages = [{"role": "user", "content": "Email Marie Dupont at marie@acme.com"}] + data = {"messages": messages} + out = await gr.async_pre_call_hook(None, None, data, "completion") + + # Mutated IN PLACE (same list object), so the proxy's shallow body snapshot + # ends up pointing at the anonymized messages, not the original raw-PII ones. + assert out["messages"] is messages + content = messages[0]["content"] + assert "Marie Dupont" not in content + assert "marie@acme.com" not in content + assert "" in content and "" in content + assert out["metadata"]["privaite_map"] == _FAKES + + +@pytest.mark.asyncio +async def test_pre_call_without_deanonymize_does_not_stash(): + gr = _make_guardrail(deanonymize=False) + data = {"messages": [{"role": "user", "content": "Hi Marie Dupont"}]} + out = await gr.async_pre_call_hook(None, None, data, "completion") + assert "" in out["messages"][0]["content"] + assert "privaite_map" not in out.get("metadata", {}) + + +@pytest.mark.asyncio +async def test_pre_call_without_pii_does_not_stash(): + gr = _make_guardrail() + data = {"messages": [{"role": "user", "content": "no personal data here"}]} + out = await gr.async_pre_call_hook(None, None, data, "completion") + assert "privaite_map" not in out.get("metadata", {}) + + +@pytest.mark.asyncio +async def test_pre_call_clears_client_supplied_map(): + # metadata is caller-controlled: a client-supplied privaite_map must be + # dropped, so it can never drive post-call restoration of model output. + gr = _make_guardrail() + data = { + "messages": [{"role": "user", "content": "no personal data here"}], + "metadata": {"privaite_map": {"": "attacker-chosen text"}}, + } + out = await gr.async_pre_call_hook(None, None, data, "completion") + assert "privaite_map" not in out["metadata"] + + +@pytest.mark.asyncio +async def test_post_call_restores_text_and_tool_call_args(): + gr = _make_guardrail() + data = {"metadata": {"privaite_map": _FAKES}} + + tool_call = types.SimpleNamespace(function=types.SimpleNamespace(arguments='{"to": ""}')) + # a tool_call with no function must be skipped, not crash. + empty_tool_call = types.SimpleNamespace(function=None) + message = types.SimpleNamespace( + content="Sending to ", + tool_calls=[tool_call, empty_tool_call], + function_call=types.SimpleNamespace(arguments='{"name": ""}'), + ) + response = types.SimpleNamespace( + choices=[ + types.SimpleNamespace(message=None), # choice with no message is skipped + types.SimpleNamespace(message=message), + ] + ) + + out = await gr.async_post_call_success_hook(data, None, response) + + msg = out.choices[1].message + assert msg.content == "Sending to Marie Dupont" + assert msg.tool_calls[0].function.arguments == '{"to": "marie@acme.com"}' + assert msg.function_call.arguments == '{"name": "Marie Dupont"}' + # the map is consumed (popped) so it cannot be persisted to spend logs. + assert "privaite_map" not in data["metadata"] + + +@pytest.mark.asyncio +async def test_post_call_deanonymize_false_skips_restore(): + gr = _make_guardrail(deanonymize=False) + data = {"metadata": {"privaite_map": _FAKES}} + message = types.SimpleNamespace(content="Sending to ", tool_calls=None, function_call=None) + response = types.SimpleNamespace(choices=[types.SimpleNamespace(message=message)]) + + out = await gr.async_post_call_success_hook(data, None, response) + assert out.choices[0].message.content == "Sending to " + + +@pytest.mark.asyncio +async def test_post_call_without_map_is_passthrough(): + gr = _make_guardrail() + response = types.SimpleNamespace(choices=[]) + assert await gr.async_post_call_success_hook({}, None, response) is response + + +@pytest.mark.asyncio +async def test_streaming_restores_content(): + gr = _make_guardrail() + request_data = {"metadata": {"privaite_map": _FAKES}} + + async def _source(): + # a choice with no delta must be skipped, not crash. + yield types.SimpleNamespace(choices=[types.SimpleNamespace(delta=None, finish_reason=None)]) + yield types.SimpleNamespace( + choices=[ + types.SimpleNamespace( + delta=types.SimpleNamespace(content="Hi "), + finish_reason=None, + ) + ] + ) + yield types.SimpleNamespace( + choices=[ + types.SimpleNamespace( + delta=types.SimpleNamespace(content=""), + finish_reason="stop", + ) + ] + ) + + chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _source(), request_data)) + assert chunks[1].choices[0].delta.content == "Hi Marie Dupont" + + +@pytest.mark.asyncio +async def test_streaming_passthrough_when_disabled(): + gr = _make_guardrail(deanonymize=False) + + async def _source(): + yield types.SimpleNamespace(choices=[]) + + chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _source(), {})) + assert len(chunks) == 1 + + +@pytest.mark.asyncio +async def test_streaming_isolates_buffers_per_choice(): + # With n>1 the provider interleaves single-choice chunks for different + # indices. A placeholder split across choice 0's chunks must not be + # corrupted by choice 1's interleaved bytes (one buffer per choice index). + gr = _make_guardrail() + request_data = {"metadata": {"privaite_map": _FAKES}} + + def _chunk(index, content, finish=None): + return types.SimpleNamespace( + choices=[ + types.SimpleNamespace( + index=index, + delta=types.SimpleNamespace(content=content), + finish_reason=finish, + ) + ] + ) + + async def _source(): + yield _chunk(0, "Hi ", finish="stop") # completes choice-0 placeholder + yield _chunk(1, "", finish="stop") + + chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _source(), request_data)) + + restored = {} + for chunk in chunks: + for choice in chunk.choices: + restored.setdefault(choice.index, "") + if choice.delta.content: + restored[choice.index] += choice.delta.content + + assert restored[0] == "Hi Marie Dupont" + assert restored[1] == "Bye" + + +@pytest.mark.asyncio +async def test_streaming_restores_tool_and_function_call_arguments(): + # Streamed tool-call and legacy function_call argument fragments must be + # de-anonymized too, with a placeholder split across chunks reassembled per + # tool_call index. A tool_call with no function is skipped, not crashed. + gr = _make_guardrail() + request_data = {"metadata": {"privaite_map": _FAKES}} + + async def _source(): + yield types.SimpleNamespace( + choices=[ + types.SimpleNamespace( + index=0, + delta=types.SimpleNamespace( + content=None, + tool_calls=[ + types.SimpleNamespace( + index=0, + function=types.SimpleNamespace(arguments='{"to": ""}'), + ) + ], + function_call=types.SimpleNamespace(arguments='SON_1>"}'), + ), + finish_reason="stop", + ) + ] + ) + + chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _source(), request_data)) + + tool_args = "".join( + tc.function.arguments + for chunk in chunks + for choice in chunk.choices + for tc in (choice.delta.tool_calls or []) + if tc.function is not None and tc.function.arguments + ) + fc_args = "".join( + choice.delta.function_call.arguments + for chunk in chunks + for choice in chunk.choices + if choice.delta.function_call is not None and choice.delta.function_call.arguments + ) + assert tool_args == '{"to": "marie@acme.com"}' + assert fc_args == '{"n": "Marie Dupont"}' + + +# --- Responses API (/v1/responses) --- + + +@pytest.mark.asyncio +async def test_pre_call_anonymizes_responses_string_input_and_fixes_snapshot(): + gr = _make_guardrail() + # the proxy takes a shallow snapshot of the body before this hook runs + body = {"input": "Email Marie Dupont at marie@acme.com"} + data = { + "input": "Email Marie Dupont at marie@acme.com", + "proxy_server_request": {"body": body}, + } + out = await gr.async_pre_call_hook(None, None, data, "aresponses") + + assert "Marie Dupont" not in out["input"] + assert "" in out["input"] + # the string-rebind would otherwise leave raw PII in the snapshot body + assert body["input"] == out["input"] + assert "Marie Dupont" not in body["input"] + assert out["metadata"]["privaite_map"] == _FAKES + + +@pytest.mark.asyncio +async def test_pre_call_anonymizes_responses_role_message_list_in_place(): + gr = _make_guardrail() + input_value = [{"role": "user", "content": "Hi Marie Dupont"}] + data = {"input": input_value} + out = await gr.async_pre_call_hook(None, None, data, "aresponses") + # mutated in place (same list object) + assert out["input"] is input_value + assert "" in input_value[0]["content"] + assert "Marie Dupont" not in input_value[0]["content"] + + +@pytest.mark.asyncio +async def test_pre_call_anonymizes_responses_content_parts_in_place(): + gr = _make_guardrail() + input_value = [{"type": "input_text", "text": "Hi Marie Dupont"}] + data = {"input": input_value} + out = await gr.async_pre_call_hook(None, None, data, "aresponses") + assert out["input"] is input_value + assert input_value[0]["text"] == "Hi " + + +@pytest.mark.asyncio +async def test_pre_call_anonymizes_both_messages_and_input(): + # a crafted /v1/responses body with decoy messages + PII in input: BOTH must + # be anonymized under one shared mapping (input must not be skipped). + gr = _make_guardrail() + messages = [{"role": "user", "content": "Email Marie Dupont"}] + data = {"messages": messages, "input": "also reach marie@acme.com"} + out = await gr.async_pre_call_hook(None, None, data, "aresponses") + + assert "Marie Dupont" not in messages[0]["content"] + assert "" in messages[0]["content"] + assert "marie@acme.com" not in out["input"] + assert "" in out["input"] + # one consistent map covering both sources + assert out["metadata"]["privaite_map"] == _FAKES + + +@pytest.mark.asyncio +async def test_pre_call_unhandled_input_shape_is_passthrough(): + # an unexpected input shape (neither str nor list) is left untouched. + gr = _make_guardrail() + data = {"input": {"unexpected": "shape"}} + out = await gr.async_pre_call_hook(None, None, data, "aresponses") + assert "privaite_map" not in out.get("metadata", {}) + + +@pytest.mark.asyncio +async def test_post_call_restores_responses_output_dict_and_object(): + gr = _make_guardrail() + data = {"metadata": {"privaite_map": _FAKES}} + + # output item as a dict (message with output_text content) + dict_item = { + "type": "message", + "content": [{"type": "output_text", "text": "Hello "}], + } + # output item as an object (function_call with arguments) + obj_item = types.SimpleNamespace( + type="function_call", + content=None, + arguments='{"to": ""}', + ) + response = types.SimpleNamespace(output=[dict_item, obj_item]) + + out = await gr.async_post_call_success_hook(data, None, response) + + assert out.output[0]["content"][0]["text"] == "Hello Marie Dupont" + assert out.output[1].arguments == '{"to": "marie@acme.com"}' + assert "privaite_map" not in data["metadata"] + + +@pytest.mark.asyncio +async def test_post_call_failure_hook_pops_map(): + # the success/streaming hooks never run on failure; the failure hook must + # still drop the reversible map so it cannot reach a failure spend-log. + gr = _make_guardrail() + rd = {"metadata": {"privaite_map": dict(_FAKES), "other": 1}} + assert await gr.async_post_call_failure_hook(rd, RuntimeError("x"), None) is None + assert "privaite_map" not in rd["metadata"] + assert rd["metadata"]["other"] == 1 + # 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