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