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 01/16] 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 From ed5dc6fd989d342cd7bdaa01d33ec9912936b18b Mon Sep 17 00:00:00 2001 From: crp4222 <121295348+crp4222@users.noreply.github.com> Date: Wed, 1 Jul 2026 12:38:52 +0100 Subject: [PATCH 02/16] feat(guardrails): add block_entities to the PrivAiTe guardrail A request containing any PII type listed in block_entities is now rejected with an HTTP 400 in the pre-call hook, before anything is forwarded to the model, instead of being pseudonymized. Empty by default, so existing behavior is unchanged. Accepts a list or a comma-separated string, and the error names the blocked type(s) only, never the value. Because privaite's PIIConfig uses extra="allow", an older privaite would silently swallow block_entities and forward the PII anyway; the guardrail fails closed when block rules are set against a privaite that lacks the field. Wires the option through the ConfigModel and initializer, with tests (the fake engine now raises PIIBlockedError for a blocked type). --- .../guardrail_hooks/privaite/__init__.py | 1 + .../guardrail_hooks/privaite/privaite.py | 39 ++++++- .../guardrails/guardrail_hooks/privaite.py | 9 ++ .../guardrail_hooks/test_privaite.py | 107 +++++++++++++++++- 4 files changed, 151 insertions(+), 5 deletions(-) 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") From 66eb3d56ad08336b2f71f113b923d4bb74d8b715 Mon Sep 17 00:00:00 2001 From: crp4222 <121295348+crp4222@users.noreply.github.com> Date: Thu, 2 Jul 2026 02:17:36 +0100 Subject: [PATCH 03/16] fix(guardrails): scan every Responses input item and restore the reasoning trace Second-audit fixes for the PrivAiTe guardrail: - Responses `input` is scanned item by item instead of all-or-nothing. A mixed agentic turn (a role message plus a function_call_output/tool call, a bare input_text/output_text content part, or a bare string) previously took a homogeneity path that wrapped the whole list as one message content and left every non-message item raw, bypassing detection AND the block_entities gate. Each text-bearing item (content, output, arguments, text) now flows through the shared mapping; block fires on any of them. - The reasoning trace (reasoning_content / reasoning) is de-anonymized in both the non-streaming and streaming paths. - `finished` is passed through to the streamed tool-call and function_call restores, so a fragment arriving on the same chunk as finish_reason flushes its held-back tail instead of dropping it. - The first-use spaCy model download runs off the event loop (asyncio.to_thread) so it does not stall the whole proxy worker. Tests cover the mixed input list, the block gate on a tool-output item, the input_text content-part shape, and streamed reasoning restore. --- .../guardrail_hooks/privaite/privaite.py | 120 ++++++++++++------ .../guardrail_hooks/test_privaite.py | 68 ++++++++++ 2 files changed, 149 insertions(+), 39 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py index b0c7de9b584..51aec492e51 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py @@ -154,12 +154,15 @@ class PrivaiteGuardrail(CustomGuardrail): await engine.initialize() except OSError: # spaCy models not present yet: download them once, then retry. + # The download is synchronous pip machinery pulling hundreds of + # MB; run it off the event loop so it does not stall every other + # request in this proxy worker. from spacy.cli import download for lang in languages: model = _LANG_MODELS.get(lang) if model: - download(model) + await asyncio.to_thread(download, model) engine = PIIEngine(config) await engine.initialize() @@ -167,9 +170,10 @@ class PrivaiteGuardrail(CustomGuardrail): 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) + # Text-bearing fields on non-message Responses input items: a tool output, a + # streamed/echoed tool call, and a bare input_text/output_text content part. + # Scanned individually so a mixed input list is not left raw. + _ITEM_TEXT_FIELDS = ("output", "arguments", "text") def _overwrite_snapshot_input(self, data: dict, new_input: Any) -> None: # A plain `data["input"] = ...` rebind leaks for string input: the proxy @@ -182,49 +186,76 @@ class PrivaiteGuardrail(CustomGuardrail): 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 + def _responses_item_repr(self, item: Any) -> tuple: + """Map ONE Responses `input` list item to (engine_message, field) so it + can be scanned. field=None means the item IS a message (replace it whole + with the anonymized copy); a field name means write the scrubbed content + back into item[field]; "__str__" means the item is a bare string. Returns + (None, None) for a shape with no scannable text.""" + if isinstance(item, str): + return ({"role": "user", "content": item}, "__str__") if item else (None, None) + if isinstance(item, dict): + if "role" in item: + # A message item; the engine scans its content natively. + return item, None + for field in self._ITEM_TEXT_FIELDS: + if isinstance(item.get(field), str) and item[field]: + return {"role": "user", "content": item[field]}, field + if isinstance(item.get("content"), (str, list)) and item["content"]: + return {"role": "user", "content": item["content"]}, "content" + return None, None 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.""" + engine (span-precise), sharing ONE mapping. Every Responses input item is + scanned item by item (message, tool output, tool call, bare string), so a + mixed input list no longer slips past detection and the block gate. + 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")) + batch: list = list(msg_list) - if not msg_list and not input_msgs: + input_value = data.get("input") + targets: list = [] # (index_or_"str", field) describing each write-back + if isinstance(input_value, str) and input_value: + batch.append({"role": "user", "content": input_value}) + targets.append(("str", None)) + elif isinstance(input_value, list): + for idx, item in enumerate(input_value): + rep, field = self._responses_item_repr(item) + if rep is None: + continue # no scannable text on this item shape + batch.append(rep) + targets.append((idx, field)) + + if not batch: 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:]) + + for (target, field), anon in zip(targets, anonymized[n:]): + if target == "str": + new_text = anon.get("content", input_value) + data["input"] = new_text + self._overwrite_snapshot_input(data, new_text) + continue + # every non-"str" target came from the isinstance(input_value, list) + # branch above, so input_value is a list here. + if not isinstance(input_value, list): + continue + if field is None: + input_value[target] = anon + elif field == "__str__": + input_value[target] = anon.get("content", input_value[target]) + else: + new_item = dict(input_value[target]) + new_item[field] = anon.get("content", new_item[field]) + input_value[target] = new_item return mapping async def async_pre_call_hook(self, user_api_key_dict: Any, cache: Any, data: dict, call_type: str) -> dict: @@ -259,11 +290,15 @@ class PrivaiteGuardrail(CustomGuardrail): 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.""" + """Restore originals in one response message: content, the reasoning + trace, 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 field in ("reasoning_content", "reasoning"): + value = getattr(message, field, None) + if isinstance(value, str) and value: + setattr(message, field, await engine.process_response(value, mapping)) for tool_call in getattr(message, "tool_calls", None) or []: fn = getattr(tool_call, "function", None) if fn is None: @@ -314,12 +349,17 @@ class PrivaiteGuardrail(CustomGuardrail): 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.""" + """Restore one streamed delta in place: text content, the reasoning + trace, 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 field in ("reasoning_content", "reasoning"): + value = getattr(delta, field, None) + if isinstance(value, str) and (value or finished): + setattr(delta, field, restore((field, index), value, finished)) for tool_call in getattr(delta, "tool_calls", None) or []: fn = getattr(tool_call, "function", None) if fn is None: @@ -327,12 +367,14 @@ class PrivaiteGuardrail(CustomGuardrail): 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) + # Pass `finished` so a fragment on the same chunk as finish_reason + # flushes its held-back tail instead of dropping it. + fn.arguments = restore(("tool", index, tc_index), args, finished) 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) + function_call.arguments = restore(("fc", index), fc_args, finished) async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: Any, response: Any, request_data: dict diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py index 88f6476643d..bab8c633dc4 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py @@ -486,6 +486,38 @@ async def test_streaming_isolates_buffers_per_choice(): assert restored[1] == "Bye" +@pytest.mark.asyncio +async def test_streaming_restores_reasoning_content(): + # a reasoning model echoes a placeholder in its streamed reasoning trace; it + # must be de-anonymized like content. + gr = _make_guardrail() + request_data = {"metadata": {"privaite_map": _FAKES}} + + def _chunk(reasoning, finish=None): + return types.SimpleNamespace( + choices=[ + types.SimpleNamespace( + index=0, + delta=types.SimpleNamespace( + content=None, tool_calls=None, function_call=None, reasoning_content=reasoning + ), + finish_reason=finish, + ) + ] + ) + + async def _source(): + yield _chunk("the user is ") + yield _chunk(None, finish="stop") + + chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _source(), request_data)) + trace = "".join( + getattr(choice.delta, "reasoning_content", None) or "" for chunk in chunks for choice in chunk.choices + ) + assert "Marie Dupont" in trace + assert "" not in trace + + @pytest.mark.asyncio async def test_streaming_restores_tool_and_function_call_arguments(): # Streamed tool-call and legacy function_call argument fragments must be @@ -622,6 +654,42 @@ async def test_pre_call_unhandled_input_shape_is_passthrough(): assert "privaite_map" not in out.get("metadata", {}) +@pytest.mark.asyncio +async def test_pre_call_scans_mixed_responses_input_list(): + # an agentic turn: a role message + a function_call_output + a bare string. + # the old homogeneity check wrapped the whole list as one content and left + # the non-message items raw; every item must be scanned. + gr = _make_guardrail() + data = { + "input": [ + {"role": "user", "content": "I am Marie Dupont"}, + {"type": "function_call_output", "call_id": "c1", "output": "reach marie@acme.com"}, + "also Marie Dupont", + ] + } + out = await gr.async_pre_call_hook(None, None, data, "aresponses") + serialized = str(out["input"]) + + assert "Marie Dupont" not in serialized + assert "marie@acme.com" not in serialized + assert out["input"][1]["type"] == "function_call_output" # structure kept + assert out["input"][1]["call_id"] == "c1" + assert out["metadata"]["privaite_map"] + + +@pytest.mark.asyncio +async def test_pre_call_block_gate_fires_on_mixed_responses_input(): + # a blocked type inside a function_call_output must reject the request; the + # gate is only reachable because the item is scanned in the first place. + from fastapi import HTTPException + + gr = _make_guardrail(block_entities=["EMAIL_ADDRESS"]) + data = {"input": [{"type": "function_call_output", "output": "reach marie@acme.com"}]} + with pytest.raises(HTTPException) as ei: + await gr.async_pre_call_hook(None, None, data, "aresponses") + assert ei.value.status_code == 400 + + @pytest.mark.asyncio async def test_post_call_restores_responses_output_dict_and_object(): gr = _make_guardrail() From 6c6dfcaee14b283fbff06f9f6408a0894b5f0694 Mon Sep 17 00:00:00 2001 From: crp4222 <121295348+crp4222@users.noreply.github.com> Date: Fri, 3 Jul 2026 20:04:13 +0100 Subject: [PATCH 04/16] fix(guardrails): keep the privaite restore map out of background request state Background Responses results are fetched by a later poll that the post-call success and streaming hooks never see, so the fake-to-original map stashed in request metadata would persist in stored request state with no consumer left to pop it. The pre-call hook now skips the stash for background requests; those responses keep their placeholders, which is the fail-closed behavior. Regression test included. --- .../guardrail_hooks/privaite/privaite.py | 10 +++++++++- .../guardrails/guardrail_hooks/test_privaite.py | 15 +++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py index 51aec492e51..460d478a9c6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py @@ -282,10 +282,18 @@ class PrivaiteGuardrail(CustomGuardrail): raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc - if mapping is not None and self.deanonymize and not mapping.is_empty: + if ( + mapping is not None + and self.deanonymize + and not mapping.is_empty + and not data.get("background") + ): # 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. + # Background requests are excluded: their result is fetched by a later + # poll the post-call hooks never see, so the map would persist in + # request state with no consumer; those responses keep placeholders. data.setdefault("metadata", {})["privaite_map"] = dict(mapping.get_all_fakes()) return data diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py index bab8c633dc4..b7234a8e5fc 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py @@ -347,6 +347,21 @@ async def test_pre_call_without_pii_does_not_stash(): assert "privaite_map" not in out.get("metadata", {}) +@pytest.mark.asyncio +async def test_pre_call_background_request_anonymizes_but_does_not_stash(): + # Background results are fetched by a later poll the post-call hooks never + # see, so the fake->original map must not persist in request state. + gr = _make_guardrail() + data = { + "input": "Email Marie Dupont at marie@acme.com", + "background": True, + } + out = await gr.async_pre_call_hook(None, None, data, "aresponses") + assert "Marie Dupont" not in out["input"] + assert "" in out["input"] + 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 From 73479ac2577011d37e05e11a5af609451814d080 Mon Sep 17 00:00:00 2001 From: crp4222 <121295348+crp4222@users.noreply.github.com> Date: Fri, 3 Jul 2026 20:15:16 +0100 Subject: [PATCH 05/16] fix(guardrails): narrow the spaCy retry and keep terminal None content intact Addresses the two P2 review findings: the OSError download retry now only fires for the missing-spaCy-model error (E050) and re-raises permission, disk or network OSErrors, and a terminal streaming chunk that carries finish_reason with content=None keeps its None instead of being mutated to an empty string when nothing was held back. Regression tests for both. --- .../guardrail_hooks/privaite/privaite.py | 15 ++++-- .../guardrail_hooks/test_privaite.py | 46 ++++++++++++++++++- 2 files changed, 57 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py index 460d478a9c6..c946a88f7c3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py @@ -152,11 +152,16 @@ class PrivaiteGuardrail(CustomGuardrail): engine = PIIEngine(config) try: await engine.initialize() - except OSError: + except OSError as exc: # spaCy models not present yet: download them once, then retry. # The download is synchronous pip machinery pulling hundreds of # MB; run it off the event loop so it does not stall every other - # request in this proxy worker. + # request in this proxy worker. Only the missing-model OSError + # (spaCy E050) takes this path; permission/disk/network errors + # re-raise instead of triggering a pointless download. + message = str(exc) + if "E050" not in message and "Can't find model" not in message: + raise from spacy.cli import download for lang in languages: @@ -362,7 +367,11 @@ class PrivaiteGuardrail(CustomGuardrail): the legacy function_call.""" content = getattr(delta, "content", None) or "" restored = restore(("content", index), content, finished) - if content or finished: + # Overwrite whenever there was input text (even if the whole fragment is + # held back, the raw fragment must not stay visible) or the finish flush + # produced text; a terminal chunk with content=None and nothing held back + # keeps its None instead of becoming "". + if content or restored: delta.content = restored for field in ("reasoning_content", "reasoning"): value = getattr(delta, field, None) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py index b7234a8e5fc..f2e7cdba04e 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py @@ -56,6 +56,7 @@ class _FakeEngine: """ raise_oserror_until = 0 + oserror_message = "[E050] Can't find model 'en_core_web_lg'." init_calls = 0 def __init__(self, config): @@ -64,7 +65,7 @@ class _FakeEngine: async def initialize(self): type(self).init_calls += 1 if self.init_calls <= type(self).raise_oserror_until: - raise OSError("spaCy model missing") + raise OSError(type(self).oserror_message) def _scrub(self, content, mapping): if isinstance(content, str): @@ -154,6 +155,7 @@ 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.oserror_message = "[E050] Can't find model 'en_core_web_lg'." _FakeEngine.init_calls = 0 def _module(name, **attrs): @@ -306,6 +308,18 @@ async def test_engine_downloads_spacy_models_on_oserror(): assert _FakeEngine.init_calls == 2 # failed once, then succeeded after download +@pytest.mark.asyncio +async def test_engine_does_not_retry_non_model_oserror(): + # Only the missing-spaCy-model OSError may trigger the download retry; + # a disk/permission/network OSError must propagate untouched. + _FakeEngine.raise_oserror_until = 1 + _FakeEngine.oserror_message = "disk full" + gr = _make_guardrail() + with pytest.raises(OSError, match="disk full"): + await gr._engine_for(["en"]) + assert _FakeEngine.init_calls == 1 # no retry, no download + + @pytest.mark.asyncio async def test_pre_call_no_messages_is_passthrough(): gr = _make_guardrail() @@ -452,6 +466,36 @@ async def test_streaming_restores_content(): assert chunks[1].choices[0].delta.content == "Hi Marie Dupont" +@pytest.mark.asyncio +async def test_streaming_terminal_none_content_stays_none(): + # A provider finish chunk commonly carries finish_reason with content=None; + # with nothing held back, restoration must not mutate that None into "". + gr = _make_guardrail() + request_data = {"metadata": {"privaite_map": _FAKES}} + + async def _source(): + yield types.SimpleNamespace( + choices=[ + types.SimpleNamespace( + delta=types.SimpleNamespace(content="Hi "), + finish_reason=None, + ) + ] + ) + yield types.SimpleNamespace( + choices=[ + types.SimpleNamespace( + delta=types.SimpleNamespace(content=None), + finish_reason="stop", + ) + ] + ) + + chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _source(), request_data)) + assert chunks[0].choices[0].delta.content == "Hi Marie Dupont" + assert chunks[1].choices[0].delta.content is None + + @pytest.mark.asyncio async def test_streaming_passthrough_when_disabled(): gr = _make_guardrail(deanonymize=False) From b27457862feb4323e9d8ae8d4b7eb7e45fc6e20d Mon Sep 17 00:00:00 2001 From: crp4222 <121295348+crp4222@users.noreply.github.com> Date: Fri, 3 Jul 2026 20:43:44 +0100 Subject: [PATCH 06/16] style(guardrails): ruff format the privaite hook --- .../proxy/guardrails/guardrail_hooks/privaite/privaite.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py index c946a88f7c3..62a37f4b812 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py @@ -287,12 +287,7 @@ class PrivaiteGuardrail(CustomGuardrail): raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc - if ( - mapping is not None - and self.deanonymize - and not mapping.is_empty - and not data.get("background") - ): + if mapping is not None and self.deanonymize and not mapping.is_empty and not data.get("background"): # 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. From ee88ab7183b837d13dff4da7a613b435d03656b3 Mon Sep 17 00:00:00 2001 From: crp4222 <121295348+crp4222@users.noreply.github.com> Date: Mon, 20 Jul 2026 11:45:05 +0100 Subject: [PATCH 07/16] feat(guardrails): match privaite guardrail to core 0.4 PII coverage Scan the Responses input carriers the batching pass missed (typed action carriers, custom_tool_call_output.output as a list of parts, prompt.variables) plus the auxiliary request fields (completions prompt/suffix, chat prediction.content, web_search_options.user_location), and restore the refusal and audio transcript on both the sync and streaming paths. Coverage is reimplemented inline over the public engine surface (process_request / process_response), with no dependency on privaite internals. --- .../guardrail_hooks/privaite/privaite.py | 397 ++++++++++++++---- .../guardrail_hooks/test_privaite.py | 241 +++++++++++ 2 files changed, 556 insertions(+), 82 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py index 62a37f4b812..8cf802d1160 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py @@ -2,14 +2,23 @@ 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. +request 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. The request surface scanned matches the PrivAiTe +core: chat `messages`; the Responses API `input` (every item type: role message, +tool-call arguments, tool output including `custom_tool_call_output.output` as a +list of `{type, text}` parts, typed action carriers, bare content parts and bare +strings; opaque/binary items relayed whole) and `prompt.variables`; the +`/v1/completions` `prompt` (string or batch list) and `suffix`; the chat +`prediction.content` and `web_search_options.user_location` auxiliary fields (input +side only, nothing to restore; tokenized integer-array inputs pass through +unscanned). Restore covers `content`, the reasoning trace, the refusal and the audio +transcript besides tool/function arguments. Chat streaming is restored too: text +content, the reasoning trace, the refusal, the audio transcript, 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. 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 @@ -25,6 +34,7 @@ See https://github.com/crp4222/PrivAiTe for the package. from __future__ import annotations import asyncio +from collections.abc import Callable from typing import Any, cast from litellm.integrations.custom_guardrail import CustomGuardrail @@ -40,6 +50,58 @@ _LANG_MODELS = { "nl": "nl_core_news_md", } +# Responses `input` items relayed byte-for-byte: an opaque/binary payload +# (encrypted reasoning, a screenshot, a generated image) or a tool/pointer +# definition with no user text. Parity with the PrivAiTe gateway scrubber. +_RESPONSES_OPAQUE_TYPES = frozenset( + { + "reasoning", + "compaction", + "compaction_trigger", + "computer_call_output", + "image_generation_call", + "item_reference", + "mcp_list_tools", + "tool_search_call", + "tool_search_output", + "additional_tools", + } +) + +# Responses items whose user data sits in a known structured field (a typed +# command/action, a patch diff, search queries and results, interpreter code +# and logs). Every listed field is walked leaf by leaf. +_RESPONSES_DATA_FIELDS: dict[str, tuple[str, ...]] = { + "computer_call": ("action", "actions"), + "local_shell_call": ("action",), + "shell_call": ("action",), + "web_search_call": ("action",), + "apply_patch_call": ("operation",), + "file_search_call": ("queries", "results"), + "code_interpreter_call": ("code", "outputs"), + "program": ("code",), + "program_output": ("result",), +} + +# Content/output parts whose payload is binary (base64 or a file id), not text: +# scrubbing them would corrupt the payload without removing anything a text +# detector could find, so the whole part is relayed whole. +_BINARY_PART_TYPES = frozenset( + { + "input_image", + "input_file", + "input_audio", + "image", + "output_image", + "computer_screenshot", + } +) + +# Text-bearing fields on an unknown Responses item shape (e.g. an mcp_call +# carries both `arguments` and `output`): every one is scanned, not just the +# first, so a second field is never left raw. +_GENERIC_TEXT_FIELDS = ("output", "arguments", "input", "text", "reason") + def _obj_get(obj: Any, key: str) -> Any: """Read a field from a dict-or-object response item.""" @@ -175,63 +237,214 @@ class PrivaiteGuardrail(CustomGuardrail): self._engine_key = key return engine - # Text-bearing fields on non-message Responses input items: a tool output, a - # streamed/echoed tool call, and a bare input_text/output_text content part. - # Scanned individually so a mixed input list is not left raw. - _ITEM_TEXT_FIELDS = ("output", "arguments", "text") - - 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.) + def _overwrite_snapshot_field(self, data: dict, field: str, new_value: Any) -> None: + # A plain `data[field] = ...` rebind leaks for a top-level string field + # (`input`, `suffix`, a string `prompt`): the proxy snapshots the request + # body by shallow-copying data before this hook, so the original string + # stays in proxy_server_request.body[field]. Overwrite the snapshot copy + # too. (Lists and dicts 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 + if isinstance(body, dict) and field in body: + body[field] = new_value - def _responses_item_repr(self, item: Any) -> tuple: - """Map ONE Responses `input` list item to (engine_message, field) so it - can be scanned. field=None means the item IS a message (replace it whole - with the anonymized copy); a field name means write the scrubbed content - back into item[field]; "__str__" means the item is a bare string. Returns - (None, None) for a shape with no scannable text.""" + @staticmethod + def _make_setter(container: Any, key: Any) -> Callable[[str], None]: + """A write-back that drops the scrubbed string into container[key]. The + container (a list slot or a dict field) is aliased by the proxy's shallow + body snapshot, so the in-place write lands in the snapshot too.""" + + def _set(value: str) -> None: + container[key] = value + + return _set + + def _toplevel_setter(self, data: dict, field: str) -> Callable[[str], None]: + """A write-back for a top-level string field that also overwrites the + detached body snapshot copy (see _overwrite_snapshot_field).""" + + def _set(value: str) -> None: + data[field] = value + self._overwrite_snapshot_field(data, field, value) + + return _set + + def _add_leaf(self, batch: list, setters: list, text: Any, setter: Callable[[str], None]) -> None: + """Register one non-empty string leaf: a pseudo-message carrying it (so it + flows through the engine's single choke point exactly like chat content, + block gate included) plus the write-back that restores the scrubbed value.""" + if isinstance(text, str) and text: + batch.append({"role": "user", "content": text}) + setters.append(setter) + + def _collect_content(self, batch: list, setters: list, content: Any, setter: Callable[[str], None]) -> None: + """A role-message `content`: a bare string, or a list of parts whose + `text`/`refusal` fields carry user text (a bare string in the list is + user text too). Binary parts carry nothing a text detector can find.""" + if isinstance(content, str): + self._add_leaf(batch, setters, content, setter) + return + if not isinstance(content, list): + return + for idx, part in enumerate(content): + if isinstance(part, str): + self._add_leaf(batch, setters, part, self._make_setter(content, idx)) + elif isinstance(part, dict) and part.get("type") not in _BINARY_PART_TYPES: + for field in ("text", "refusal"): + self._add_leaf(batch, setters, part.get(field), self._make_setter(part, field)) + + def _collect_data_value(self, batch: list, setters: list, value: Any, setter: Callable[[str], None]) -> None: + """Walk a tool payload (a string, a list of parts, or arbitrary nested + JSON) registering every string leaf, but relay binary parts whole. Object + keys are never scanned, matching the documented boundary.""" + if isinstance(value, str): + self._add_leaf(batch, setters, value, setter) + return + if isinstance(value, list): + for idx, part in enumerate(value): + self._collect_data_value(batch, setters, part, self._make_setter(value, idx)) + return + if isinstance(value, dict) and value.get("type") not in _BINARY_PART_TYPES: + for key, sub in value.items(): + self._collect_data_value(batch, setters, sub, self._make_setter(value, key)) + + def _collect_responses_item(self, batch: list, setters: list, container: list, idx: int, item: Any) -> None: + """Register the scannable text on ONE Responses `input` list item, by its + shape (bare string, role message, typed tool call/output, or an unknown + shape scanned field by field). Opaque/binary items are relayed whole.""" if isinstance(item, str): - return ({"role": "user", "content": item}, "__str__") if item else (None, None) - if isinstance(item, dict): - if "role" in item: - # A message item; the engine scans its content natively. - return item, None - for field in self._ITEM_TEXT_FIELDS: - if isinstance(item.get(field), str) and item[field]: - return {"role": "user", "content": item[field]}, field - if isinstance(item.get("content"), (str, list)) and item["content"]: - return {"role": "user", "content": item["content"]}, "content" - return None, None + self._add_leaf(batch, setters, item, self._make_setter(container, idx)) + return + if not isinstance(item, dict): + return + if "role" in item and "content" in item: + self._collect_content(batch, setters, item["content"], self._make_setter(item, "content")) + return + itype = item.get("type") or "" + if itype in _RESPONSES_OPAQUE_TYPES: + return + if not self._collect_typed_item(batch, setters, item, itype): + self._collect_generic_item(batch, setters, item) - 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. Every Responses input item is - scanned item by item (message, tool output, tool call, bare string), so a - mixed input list no longer slips past detection and the block gate. - Returns the mapping, or None if there was nothing to anonymize.""" - messages = data.get("messages") - msg_list = messages if isinstance(messages, list) else [] - batch: list = list(msg_list) + def _collect_typed_item(self, batch: list, setters: list, item: dict, itype: str) -> bool: + """Register the text on an item whose user data sits in a type-specific + field; returns False to hand an unknown shape to the generic scan. Covers + function_call/custom_tool_call, the typed action carriers, and every + `*_output` tool output (whose `output` can be a list of {type, text} + parts, not just a string).""" + if itype == "function_call" and isinstance(item.get("arguments"), str): + self._add_leaf(batch, setters, item["arguments"], self._make_setter(item, "arguments")) + return True + if itype == "custom_tool_call" and isinstance(item.get("input"), str): + self._add_leaf(batch, setters, item["input"], self._make_setter(item, "input")) + return True + fields = _RESPONSES_DATA_FIELDS.get(itype) + if fields is not None: + for field in fields: + if item.get(field) is not None: + self._collect_data_value(batch, setters, item[field], self._make_setter(item, field)) + return True + if itype.endswith("_output") and "output" in item: + self._collect_data_value(batch, setters, item["output"], self._make_setter(item, "output")) + return True + return False + def _collect_generic_item(self, batch: list, setters: list, item: dict) -> None: + """Fallback scan for an unknown item shape: every known text-bearing field + plus a role-less `content`, so a second field is never left raw.""" + for field in _GENERIC_TEXT_FIELDS: + self._add_leaf(batch, setters, item.get(field), self._make_setter(item, field)) + if isinstance(item.get("content"), (str, list)): + self._collect_content(batch, setters, item["content"], self._make_setter(item, "content")) + + def _collect_responses_input(self, data: dict, batch: list, setters: list) -> None: input_value = data.get("input") - targets: list = [] # (index_or_"str", field) describing each write-back if isinstance(input_value, str) and input_value: - batch.append({"role": "user", "content": input_value}) - targets.append(("str", None)) + self._add_leaf(batch, setters, input_value, self._toplevel_setter(data, "input")) elif isinstance(input_value, list): for idx, item in enumerate(input_value): - rep, field = self._responses_item_repr(item) - if rep is None: - continue # no scannable text on this item shape - batch.append(rep) - targets.append((idx, field)) + self._collect_responses_item(batch, setters, input_value, idx, item) + + @staticmethod + def _prompt_variables(data: dict) -> dict | None: + """Responses prompt-template variables carry user data (the template + id/version do not). None when there are no scannable variables; + /v1/completions sends `prompt` as a string, which is not this surface.""" + prompt = data.get("prompt") + if isinstance(prompt, dict) and isinstance(prompt.get("variables"), dict): + return prompt["variables"] + return None + + @staticmethod + def _has_aux_fields(data: dict) -> bool: + """True when the request carries one of the auxiliary text fields scanned + by _collect_aux_fields; keeps the pre-call early return from skipping a + request whose only user text sits in those fields.""" + prediction = data.get("prediction") + if isinstance(prediction, dict) and "content" in prediction: + return True + web_search = data.get("web_search_options") + if isinstance(web_search, dict) and "user_location" in web_search: + return True + suffix = data.get("suffix") + if isinstance(suffix, str) and suffix: + return True + # /v1/completions user text: a string prompt or the batch list shape. + # (A dict prompt is the Responses template, handled by _prompt_variables.) + prompt = data.get("prompt") + return isinstance(prompt, (str, list)) and bool(prompt) + + def _collect_prompt_variables(self, data: dict, batch: list, setters: list) -> None: + variables = self._prompt_variables(data) + if variables is None: + return + for key, value in variables.items(): + self._collect_data_value(batch, setters, value, self._make_setter(variables, key)) + + def _collect_aux_fields(self, data: dict, batch: list, setters: list) -> None: + """The request-side text fields outside messages/input that LiteLLM + forwards verbatim (core parity): chat `prediction.content`, + `web_search_options.user_location`, and the completions `prompt` (string + or batch list) and `suffix`. Request inputs only, nothing to restore.""" + prediction = data.get("prediction") + if isinstance(prediction, dict) and prediction.get("content") is not None: + self._collect_data_value(batch, setters, prediction["content"], self._make_setter(prediction, "content")) + web_search = data.get("web_search_options") + if isinstance(web_search, dict) and web_search.get("user_location") is not None: + self._collect_data_value( + batch, setters, web_search["user_location"], self._make_setter(web_search, "user_location") + ) + suffix = data.get("suffix") + if isinstance(suffix, str) and suffix: + self._add_leaf(batch, setters, suffix, self._toplevel_setter(data, "suffix")) + prompt = data.get("prompt") + if isinstance(prompt, str) and prompt: + self._add_leaf(batch, setters, prompt, self._toplevel_setter(data, "prompt")) + elif isinstance(prompt, list): + # The batch shape: string leaves are scrubbed, a tokenized + # (integer-array) prompt passes through unscanned as documented. + for idx, part in enumerate(prompt): + self._add_leaf(batch, setters, part, self._make_setter(prompt, idx)) + + async def _anonymize_request(self, data: dict, engine: Any) -> Any: + """Anonymize every scanned request surface in place under ONE shared + mapping (chat `messages`, the Responses `input` and `prompt.variables`, + and the auxiliary prompt/suffix/prediction/user_location fields), so no + source is left untouched when a crafted request carries several. Every + scanned string flows through the engine's single choke point, so the block + gate and the fail-closed policy apply unchanged. Returns the mapping, or + None if there was nothing to anonymize.""" + messages = data.get("messages") + msg_list = messages if isinstance(messages, list) else [] + # Chat messages are scanned natively (multimodal content, bare-string + # lists, per-message block gate); everything else is registered as string + # leaves appended after them, sharing this one process_request pass. + batch: list = list(msg_list) + setters: list = [] + self._collect_responses_input(data, batch, setters) + self._collect_prompt_variables(data, batch, setters) + self._collect_aux_fields(data, batch, setters) if not batch: return None @@ -242,25 +455,8 @@ class PrivaiteGuardrail(CustomGuardrail): if msg_list: # msg_list is data["messages"] (same object) -> mutate it in place. msg_list[:] = anonymized[:n] - - for (target, field), anon in zip(targets, anonymized[n:]): - if target == "str": - new_text = anon.get("content", input_value) - data["input"] = new_text - self._overwrite_snapshot_input(data, new_text) - continue - # every non-"str" target came from the isinstance(input_value, list) - # branch above, so input_value is a list here. - if not isinstance(input_value, list): - continue - if field is None: - input_value[target] = anon - elif field == "__str__": - input_value[target] = anon.get("content", input_value[target]) - else: - new_item = dict(input_value[target]) - new_item[field] = anon.get("content", new_item[field]) - input_value[target] = new_item + for setter, anon in zip(setters, anonymized[n:]): + setter(anon.get("content", "")) return mapping async def async_pre_call_hook(self, user_api_key_dict: Any, cache: Any, data: dict, call_type: str) -> dict: @@ -272,7 +468,12 @@ class PrivaiteGuardrail(CustomGuardrail): if isinstance(metadata, dict): metadata.pop("privaite_map", None) - if not data.get("messages") and not data.get("input"): + if ( + not data.get("messages") + and not data.get("input") + and self._prompt_variables(data) is None + and not self._has_aux_fields(data) + ): return data engine = await self._engine_for(self._languages()) @@ -297,16 +498,30 @@ class PrivaiteGuardrail(CustomGuardrail): data.setdefault("metadata", {})["privaite_map"] = dict(mapping.get_all_fakes()) return data + async def _restore_audio_transcript(self, message: Any, engine: Any, mapping: Any) -> None: + """Restore the audio transcript on a response message (restore parity with + the core: an audio reply carries its text in audio.transcript, where the + model echoes placeholders like anywhere else).""" + audio = getattr(message, "audio", None) + if audio is None: + return + transcript = _obj_get(audio, "transcript") + if isinstance(transcript, str) and transcript: + _obj_set(audio, "transcript", await engine.process_response(transcript, mapping)) + async def _restore_message(self, message: Any, engine: Any, mapping: Any) -> None: """Restore originals in one response message: content, the reasoning - trace, tool-call args and the legacy function_call args.""" + trace, the refusal, the audio transcript, 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 field in ("reasoning_content", "reasoning"): + # A refusal can quote the request, so it carries placeholders too. + for field in ("reasoning_content", "reasoning", "refusal"): value = getattr(message, field, None) if isinstance(value, str) and value: setattr(message, field, await engine.process_response(value, mapping)) + await self._restore_audio_transcript(message, engine, mapping) for tool_call in getattr(message, "tool_calls", None) or []: fn = getattr(tool_call, "function", None) if fn is None: @@ -356,10 +571,26 @@ class PrivaiteGuardrail(CustomGuardrail): 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, the reasoning - trace, streamed tool-call argument fragments (per tool_call index) and - the legacy function_call.""" + def _restore_delta_audio(self, delta: Any, index: int, finished: bool, restore: Any) -> None: + """Feed streamed audio transcript fragments through their own restore + buffer and flush the held tail on the finish chunk, creating the audio + carrier when the finish delta has none so the tail is never dropped.""" + audio = getattr(delta, "audio", None) + fragment = _obj_get(audio, "transcript") if audio is not None else None + if not isinstance(fragment, str): + fragment = "" + if not fragment and not finished: + return + restored = restore(("audio", index), fragment, finished) + if audio is not None: + _obj_set(audio, "transcript", restored) + elif restored: + _obj_set(delta, "audio", {"transcript": restored}) + + def _restore_delta(self, delta: Any, index: int, finished: bool, restore: Any) -> None: + """Restore one streamed delta in place: text content, the reasoning trace, + the refusal, the audio transcript, 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) # Overwrite whenever there was input text (even if the whole fragment is @@ -368,10 +599,12 @@ class PrivaiteGuardrail(CustomGuardrail): # keeps its None instead of becoming "". if content or restored: delta.content = restored - for field in ("reasoning_content", "reasoning"): + # A refusal can quote the request, so it carries placeholders too. + for field in ("reasoning_content", "reasoning", "refusal"): value = getattr(delta, field, None) if isinstance(value, str) and (value or finished): setattr(delta, field, restore((field, index), value, finished)) + self._restore_delta_audio(delta, index, finished, restore) for tool_call in getattr(delta, "tool_calls", None) or []: fn = getattr(tool_call, "function", None) if fn is None: diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py index f2e7cdba04e..096a8de467b 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py @@ -862,3 +862,244 @@ async def test_block_entities_fails_closed_when_privaite_too_old(monkeypatch): 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") + + +# --- Responses API: tool output list-of-parts, typed carriers, prompt variables --- + + +@pytest.mark.asyncio +async def test_pre_call_scans_custom_tool_call_output_list_of_parts(): + # a shell/custom tool that read a file returns its bytes as a list of + # {type, text} parts (not a bare string); every text leaf must be scrubbed + # while the binary image part is relayed whole (base64 carries nothing a + # text detector can find, and rewriting it would corrupt the payload). + gr = _make_guardrail() + output = [ + {"type": "output_text", "text": "the file says reach marie@acme.com"}, + {"type": "input_image", "image_url": "data:image/png;base64,QUJD"}, + ] + data = {"input": [{"type": "custom_tool_call_output", "call_id": "c1", "output": output}]} + out = await gr.async_pre_call_hook(None, None, data, "aresponses") + + scrubbed = out["input"][0]["output"] + assert scrubbed[0]["text"] == "the file says reach " + assert scrubbed[1]["image_url"] == "data:image/png;base64,QUJD" # binary relayed whole + assert out["input"][0]["call_id"] == "c1" + assert out["metadata"]["privaite_map"] + + +@pytest.mark.asyncio +async def test_pre_call_scans_typed_action_carrier(): + # a file_search_call carries user data in its `queries`/`results` fields, not + # in `content`/`output`; the old field scan never looked there. + gr = _make_guardrail() + data = { + "input": [ + { + "type": "file_search_call", + "id": "fs1", + "queries": ["records for Marie Dupont", "orders by marie@acme.com"], + "results": [{"text": "found Marie Dupont", "score": 0.9}], + } + ] + } + out = await gr.async_pre_call_hook(None, None, data, "aresponses") + item = out["input"][0] + assert item["queries"] == ["records for ", "orders by "] + assert item["results"][0]["text"] == "found " + assert item["results"][0]["score"] == 0.9 # non-text leaf untouched + assert item["id"] == "fs1" + + +@pytest.mark.asyncio +async def test_pre_call_scans_prompt_variables(): + # Responses prompt-template variables carry user data; the template id/version + # do not. A variable can be a bare string or a typed content part. + gr = _make_guardrail() + data = { + "prompt": { + "id": "pmpt_123", + "version": "2", + "variables": { + "customer": "Marie Dupont", + "note": {"type": "input_text", "text": "email marie@acme.com"}, + }, + } + } + out = await gr.async_pre_call_hook(None, None, data, "aresponses") + variables = out["prompt"]["variables"] + assert variables["customer"] == "" + assert variables["note"]["text"] == "email " + assert out["prompt"]["id"] == "pmpt_123" # template id is not user data + assert out["prompt"]["version"] == "2" + assert out["metadata"]["privaite_map"] + + +@pytest.mark.asyncio +async def test_pre_call_block_gate_fires_on_prompt_variables(): + # the block gate must reach a blocked type sitting only in prompt.variables; + # it is only reachable because the variable is scanned in the first place. + from fastapi import HTTPException + + gr = _make_guardrail(block_entities=["EMAIL_ADDRESS"]) + data = {"prompt": {"id": "pmpt_1", "variables": {"to": "marie@acme.com"}}} + with pytest.raises(HTTPException) as ei: + await gr.async_pre_call_hook(None, None, data, "aresponses") + assert ei.value.status_code == 400 + + +# --- Auxiliary request fields: completions prompt/suffix, prediction, user_location --- + + +@pytest.mark.asyncio +async def test_pre_call_scans_completions_prompt_string_and_suffix_and_fixes_snapshot(): + gr = _make_guardrail() + body = {"prompt": "Complete for Marie Dupont", "suffix": "signed marie@acme.com"} + data = { + "prompt": "Complete for Marie Dupont", + "suffix": "signed marie@acme.com", + "proxy_server_request": {"body": body}, + } + out = await gr.async_pre_call_hook(None, None, data, "text_completion") + + assert out["prompt"] == "Complete for " + assert out["suffix"] == "signed " + # the detached snapshot copy must be overwritten too, or raw PII leaks through it + assert body["prompt"] == out["prompt"] + assert body["suffix"] == out["suffix"] + assert "Marie Dupont" not in body["prompt"] + assert out["metadata"]["privaite_map"] + + +@pytest.mark.asyncio +async def test_pre_call_scans_completions_prompt_batch_list_and_skips_tokens(): + # the /v1/completions batch shape: string leaves are scrubbed; a tokenized + # (integer-array) prompt passes through unscanned as documented. + gr = _make_guardrail() + prompt = ["reach marie@acme.com", [1, 15, 27]] + data = {"prompt": prompt} + out = await gr.async_pre_call_hook(None, None, data, "text_completion") + assert out["prompt"][0] == "reach " + assert out["prompt"][1] == [1, 15, 27] # token array untouched + + +@pytest.mark.asyncio +async def test_pre_call_scans_prediction_content(): + # chat predicted outputs carry the client's current document verbatim. + gr = _make_guardrail() + prediction = {"type": "content", "content": "draft addressed to Marie Dupont"} + data = {"messages": [{"role": "user", "content": "edit this"}], "prediction": prediction} + out = await gr.async_pre_call_hook(None, None, data, "completion") + # prediction dict is aliased by the body snapshot, so it is rewritten in place + assert prediction["content"] == "draft addressed to " + assert out["metadata"]["privaite_map"] + + +@pytest.mark.asyncio +async def test_pre_call_scans_web_search_user_location(): + gr = _make_guardrail() + web_search = {"user_location": {"type": "approximate", "approximate": {"city": "Marie Dupont"}}} + data = {"messages": [{"role": "user", "content": "weather?"}], "web_search_options": web_search} + await gr.async_pre_call_hook(None, None, data, "completion") + assert web_search["user_location"]["approximate"]["city"] == "" + + +@pytest.mark.asyncio +async def test_pre_call_only_aux_field_is_not_skipped(): + # a request whose ONLY user text is an auxiliary field (no messages/input) + # must not be short-circuited by the pre-call early return. + gr = _make_guardrail() + data = {"suffix": "signed marie@acme.com"} + out = await gr.async_pre_call_hook(None, None, data, "text_completion") + assert out["suffix"] == "signed " + assert out["metadata"]["privaite_map"] + + +# --- Restore parity: refusal and audio transcript (non-streaming + streaming) --- + + +@pytest.mark.asyncio +async def test_post_call_restores_refusal_and_audio_transcript(): + gr = _make_guardrail() + data = {"metadata": {"privaite_map": _FAKES}} + + # audio as an object (transcript attribute) and refusal quoting the request. + obj_msg = types.SimpleNamespace( + content=None, + tool_calls=None, + function_call=None, + refusal="I won't email ", + audio=types.SimpleNamespace(transcript="calling now"), + ) + # audio as a dict on a second choice. + dict_msg = types.SimpleNamespace( + content=None, + tool_calls=None, + function_call=None, + audio={"transcript": "reach "}, + ) + response = types.SimpleNamespace( + choices=[types.SimpleNamespace(message=obj_msg), types.SimpleNamespace(message=dict_msg)] + ) + out = await gr.async_post_call_success_hook(data, None, response) + + assert out.choices[0].message.refusal == "I won't email Marie Dupont" + assert out.choices[0].message.audio.transcript == "calling Marie Dupont now" + assert out.choices[1].message.audio["transcript"] == "reach marie@acme.com" + + +@pytest.mark.asyncio +async def test_streaming_restores_refusal(): + gr = _make_guardrail() + request_data = {"metadata": {"privaite_map": _FAKES}} + + def _chunk(refusal, finish=None): + return types.SimpleNamespace( + choices=[ + types.SimpleNamespace( + index=0, + delta=types.SimpleNamespace(content=None, tool_calls=None, function_call=None, refusal=refusal), + finish_reason=finish, + ) + ] + ) + + async def _source(): + yield _chunk("I won't reach ", finish="stop") + + chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _source(), request_data)) + refusal = "".join(getattr(c.choices[0].delta, "refusal", None) or "" for c in chunks) + assert refusal == "I won't reach Marie Dupont" + + +@pytest.mark.asyncio +async def test_streaming_restores_audio_transcript(): + # streamed audio transcript fragments are de-anonymized with a placeholder + # split across chunks reassembled in the audio-segment buffer. + gr = _make_guardrail() + request_data = {"metadata": {"privaite_map": _FAKES}} + + def _chunk(transcript, finish=None): + return types.SimpleNamespace( + choices=[ + types.SimpleNamespace( + index=0, + delta=types.SimpleNamespace( + content=None, + tool_calls=None, + function_call=None, + audio=types.SimpleNamespace(transcript=transcript), + ), + finish_reason=finish, + ) + ] + ) + + async def _source(): + yield _chunk("Hi ", finish="stop") + + chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _source(), request_data)) + transcript = "".join(c.choices[0].delta.audio.transcript or "" for c in chunks if c.choices[0].delta.audio) + assert transcript == "Hi Marie Dupont" From bbb6bb191b6e43603f739924823e91c5256b32de Mon Sep 17 00:00:00 2001 From: crp4222 <121295348+crp4222@users.noreply.github.com> Date: Mon, 20 Jul 2026 18:50:08 +0100 Subject: [PATCH 08/16] refactor(guardrails): make the privaite payload walk iterative CI's recursive_detector gate bans self-recursive functions, and _collect_data_value tripped it. Replace the self-recursion with an explicit worklist; children are pushed reversed so string leaves register in the same depth-first order as before, keeping placeholder numbering and the setter/batch pairing byte-identical. Nesting depth is now bounded by payload size instead of the interpreter recursion limit. --- .../guardrail_hooks/privaite/privaite.py | 23 ++++++++++--------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py index 8cf802d1160..19c6aa559de 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py @@ -297,17 +297,18 @@ class PrivaiteGuardrail(CustomGuardrail): def _collect_data_value(self, batch: list, setters: list, value: Any, setter: Callable[[str], None]) -> None: """Walk a tool payload (a string, a list of parts, or arbitrary nested JSON) registering every string leaf, but relay binary parts whole. Object - keys are never scanned, matching the documented boundary.""" - if isinstance(value, str): - self._add_leaf(batch, setters, value, setter) - return - if isinstance(value, list): - for idx, part in enumerate(value): - self._collect_data_value(batch, setters, part, self._make_setter(value, idx)) - return - if isinstance(value, dict) and value.get("type") not in _BINARY_PART_TYPES: - for key, sub in value.items(): - self._collect_data_value(batch, setters, sub, self._make_setter(value, key)) + keys are never scanned, matching the documented boundary. Iterative + worklist, no recursion; children are pushed reversed so leaves register + in the same depth-first order as the recursive walk they replace.""" + stack: list[tuple[Any, Callable[[str], None]]] = [(value, setter)] + while stack: + node, write_back = stack.pop() + if isinstance(node, str): + self._add_leaf(batch, setters, node, write_back) + elif isinstance(node, list): + stack.extend((node[idx], self._make_setter(node, idx)) for idx in range(len(node) - 1, -1, -1)) + elif isinstance(node, dict) and node.get("type") not in _BINARY_PART_TYPES: + stack.extend((sub, self._make_setter(node, key)) for key, sub in reversed(node.items())) def _collect_responses_item(self, batch: list, setters: list, container: list, idx: int, item: Any) -> None: """Register the scannable text on ONE Responses `input` list item, by its From 6d216ac0687518d53ffa5ecfc7fba6cf84992acf Mon Sep 17 00:00:00 2001 From: crp4222 <121295348+crp4222@users.noreply.github.com> Date: Sun, 23 Aug 2026 12:36:48 +0200 Subject: [PATCH 09/16] fix(guardrails): flush privaite held tails onto a bare finish chunk and drain them at end of stream --- .../guardrail_hooks/privaite/privaite.py | 170 ++++++++++++++--- .../guardrail_hooks/test_privaite.py | 176 ++++++++++++++++++ 2 files changed, 320 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py index 19c6aa559de..25fb5f718af 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py @@ -34,7 +34,8 @@ See https://github.com/crp4222/PrivAiTe for the package. from __future__ import annotations import asyncio -from collections.abc import Callable +import copy +from collections.abc import Callable, Iterator from typing import Any, cast from litellm.integrations.custom_guardrail import CustomGuardrail @@ -116,6 +117,144 @@ def _obj_set(obj: Any, key: str, value: Any) -> None: setattr(obj, key, value) +def _append_nested(delta: object, carrier: str, field: str, remaining: str) -> None: + """Append text to delta.., creating the carrier as a plain + dict when the delta has none (a bare finish delta carries no function_call + or audio object).""" + holder = _obj_get(delta, carrier) + if holder is None: + _obj_set(delta, carrier, {field: remaining}) + else: + _obj_set(holder, field, (_obj_get(holder, field) or "") + remaining) + + +def _append_tool_remainder(delta: object, tc_index: int, remaining: str) -> None: + """Append a held argument fragment under its tool_call index, creating the + slot when the delta has none: the client reassembles arguments per index.""" + calls = _obj_get(delta, "tool_calls") + if calls is None: + calls = [] + _obj_set(delta, "tool_calls", calls) + for call in calls: + if (_obj_get(call, "index") or 0) == tc_index: + _append_nested(call, "function", "arguments", remaining) + return + calls.append({"index": tc_index, "function": {"arguments": remaining}}) + + +def _append_remainder(delta: object, key: tuple, remaining: str) -> None: + """Append a restore buffer's held tail onto a streamed delta. The key is + (kind, choice index[, tool_call index]) as built by the streaming hook.""" + kind = key[0] + if kind == "tool": + _append_tool_remainder(delta, key[2], remaining) + elif kind == "fc": + _append_nested(delta, "function_call", "arguments", remaining) + elif kind == "audio": + _append_nested(delta, "audio", "transcript", remaining) + else: + # content and the text fields (reasoning_content, reasoning, refusal) + _obj_set(delta, kind, (_obj_get(delta, kind) or "") + remaining) + + +_DELTA_CHANNELS = ( + "content", + "reasoning_content", + "reasoning", + "refusal", + "audio", + "tool_calls", + "function_call", +) + + +def _drain_chunk(template: object, key: tuple, remaining: str) -> object: + """A trailing chunk carrying one held tail after the provider closed the + stream without a finish_reason: a clone of the last chunk of that choice + (same id, model and object as the rest of the stream) reduced to that one + choice, its delta cleared, no finish_reason and no usage.""" + chunk = copy.deepcopy(template) + index = key[1] + choice = next(c for c in _obj_get(chunk, "choices") or [] if (_obj_get(c, "index") or 0) == index) + _obj_set(chunk, "choices", [choice]) + if _obj_get(chunk, "usage") is not None: + _obj_set(chunk, "usage", None) + _obj_set(choice, "finish_reason", None) + delta = _obj_get(choice, "delta") + for field in _DELTA_CHANNELS: + if _obj_get(delta, field) is not None: + _obj_set(delta, field, None) + _append_remainder(delta, key, remaining) + return chunk + + +class _StreamRestorer: + """Restore state for one streamed response. + + 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. The buffer holds back partial + placeholders that span chunk boundaries; whatever remains is flushed onto + the chunk that carries finish_reason (by the channel's own restore when + the finish delta carries it, by the sweep in restore_chunk when it does + not: the common finish shape is a bare `delta: {}`), or, when the provider + closes the stream without a finish_reason, drained as trailing chunks + cloned from the last chunk seen for that choice. A held tail is never + dropped.""" + + def __init__(self, mapping: object, make_buffer: Callable[[object], object]) -> None: + self._mapping = mapping + self._make_buffer = make_buffer + self._buffers: dict = {} + self._last_chunk: dict = {} + + def restore(self, key: tuple, text: str, finished: bool) -> str: + if key not in self._buffers: + self._buffers[key] = self._make_buffer(self._mapping) + deanon = self._buffers[key] + out = deanon.feed(text) if text else "" + if finished: + # The channel is closed: drop its buffer so a sweep never flushes + # it a second time. + out += self._buffers.pop(key).flush() + return out + + def close_open(self, index: int | None) -> tuple[tuple[tuple, str], ...]: + """Flush and drop every buffer still open for this choice (every choice + when index is None); returns the non-empty tails.""" + tails = [] + for key in [k for k in self._buffers if index is None or k[1] == index]: + remaining = self._buffers.pop(key).flush() + if remaining: + tails.append((key, remaining)) + return tuple(tails) + + def restore_chunk(self, chunk: object, restore_delta: Callable[..., None]) -> None: + """Restore every choice of one chunk in place; on a finished choice, + append the tails of the channels its finish delta does not carry.""" + 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 + self._last_chunk[index] = chunk + finished = getattr(choice, "finish_reason", None) is not None + restore_delta(delta, index, finished, self.restore) + if finished: + for key, remaining in self.close_open(index): + _append_remainder(delta, key, remaining) + + def drain(self) -> Iterator[object]: + """Trailing chunks for the tails still held after the provider closed + the stream without a finish_reason for their choice.""" + for key, remaining in self.close_open(None): + template = self._last_chunk.get(key[1]) + if template is not None: + yield _drain_chunk(template, key, remaining) + + class PrivaiteGuardrail(CustomGuardrail): def __init__(self, **kwargs: Any) -> None: self.preset = kwargs.pop("preset", None) or "onnx" @@ -639,33 +778,12 @@ class PrivaiteGuardrail(CustomGuardrail): 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. + restorer = _StreamRestorer(mapping, StreamingDeAnonymizer) 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) + restorer.restore_chunk(chunk, self._restore_delta) yield chunk + for drained in restorer.drain(): + yield drained async def async_post_call_failure_hook( self, diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py index 096a8de467b..acc9877b555 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py @@ -10,6 +10,7 @@ wiring end to end (registry, engine caching, pre/post/streaming hooks) without re-testing PrivAiTe's detection internals. """ +import json import sys import types @@ -1103,3 +1104,178 @@ async def test_streaming_restores_audio_transcript(): chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _source(), request_data)) transcript = "".join(c.choices[0].delta.audio.transcript or "" for c in chunks if c.choices[0].delta.audio) assert transcript == "Hi Marie Dupont" + + +def _bare_delta(**fields): + """A streamed delta carrying only the given fields (the others None), the + shape of the finish chunk most providers send: delta {} + finish_reason.""" + base = dict(content=None, tool_calls=None, function_call=None) + base.update(fields) + return types.SimpleNamespace(**base) + + +def _choice_chunk(delta, finish=None, index=0): + return types.SimpleNamespace(choices=[types.SimpleNamespace(index=index, delta=delta, finish_reason=finish)]) + + +def _get(obj, *path): + """Walk dict-or-object attributes; the guardrail may create a carrier as a + plain dict on a finish delta that had none.""" + for key in path: + if obj is None: + return None + obj = obj.get(key) if isinstance(obj, dict) else getattr(obj, key, None) + return obj + + +# The restorer holds back a tail that could begin a placeholder (", then or Date: Sun, 23 Aug 2026 12:38:32 +0200 Subject: [PATCH 10/16] fix(guardrails): keep privaite-restored tool-call arguments valid JSON, in responses and in streams --- .../guardrail_hooks/privaite/privaite.py | 67 ++++++++++++--- .../guardrail_hooks/test_privaite.py | 83 ++++++++++++++++++- 2 files changed, 138 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py index 25fb5f718af..d2bcb3ce476 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py @@ -35,6 +35,7 @@ from __future__ import annotations import asyncio import copy +import json from collections.abc import Callable, Iterator from typing import Any, cast @@ -117,6 +118,43 @@ def _obj_set(obj: Any, key: str, value: Any) -> None: setattr(obj, key, value) +def _json_escape(value: str) -> str: + """The JSON string-literal encoding of value, without the surrounding + quotes: what a restored original must look like when spliced into streamed + argument fragments. Mirrors privaite.streaming.buffer.json_escape.""" + return json.dumps(value, ensure_ascii=False)[1:-1] + + +async def _restore_json_tree(engine: Any, value: Any, mapping: Any) -> Any: + """Restore every string leaf of a parsed JSON document (keys are not + rewritten, mirroring the core engine's walker).""" + if isinstance(value, str): + return await engine.process_response(value, mapping) + if isinstance(value, dict): + return {key: await _restore_json_tree(engine, item, mapping) for key, item in value.items()} + if isinstance(value, list): + return [await _restore_json_tree(engine, item, mapping) for item in value] + return value + + +async def _restore_arguments(engine: Any, arguments: str, mapping: Any) -> str: + """Restore a tool/function call's `arguments`: a JSON document inside a + string. Plain substitution would splice a raw quote, backslash or newline + from the original into a JSON string literal and the client's json.loads of + the arguments would fail; restore on the parsed tree and re-encode instead + (core parity: PIIEngine._deanonymize_arguments). When nothing changes, or + the string is not JSON, the previous behaviour is kept byte for byte.""" + plain = await engine.process_response(arguments, mapping) + if plain == arguments: + return arguments + try: + parsed = json.loads(arguments) + except ValueError: + # Not JSON: the plain string restore is all there is. + return plain + return json.dumps(await _restore_json_tree(engine, parsed, mapping), ensure_ascii=False) + + def _append_nested(delta: object, carrier: str, field: str, remaining: str) -> None: """Append text to delta.., creating the carrier as a plain dict when the delta has none (a bare finish delta carries no function_call @@ -204,15 +242,20 @@ class _StreamRestorer: cloned from the last chunk seen for that choice. A held tail is never dropped.""" - def __init__(self, mapping: object, make_buffer: Callable[[object], object]) -> None: + def __init__(self, mapping: object, json_mapping: object, make_buffer: Callable[..., object]) -> None: self._mapping = mapping + # For the argument channels: their fragments are JSON source text, so the + # original has to arrive JSON-escaped or the client's parse of the + # reassembled arguments fails on a quote, a backslash or a newline. + # (Placeholders need no escaping, so the fakes are the same on both sides.) + self._json_mapping = json_mapping self._make_buffer = make_buffer self._buffers: dict = {} self._last_chunk: dict = {} - def restore(self, key: tuple, text: str, finished: bool) -> str: + def restore(self, key: tuple, text: str, finished: bool, json_fragment: bool = False) -> str: if key not in self._buffers: - self._buffers[key] = self._make_buffer(self._mapping) + self._buffers[key] = self._make_buffer(self._json_mapping if json_fragment else self._mapping) deanon = self._buffers[key] out = deanon.feed(text) if text else "" if finished: @@ -668,12 +711,12 @@ class PrivaiteGuardrail(CustomGuardrail): continue args = getattr(fn, "arguments", None) if args: - fn.arguments = await engine.process_response(args, mapping) + fn.arguments = await _restore_arguments(engine, 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) + function_call.arguments = await _restore_arguments(engine, 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 @@ -685,7 +728,7 @@ class PrivaiteGuardrail(CustomGuardrail): _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)) + _obj_set(item, "arguments", await _restore_arguments(engine, args, mapping)) async def async_post_call_success_hook(self, data: dict, user_api_key_dict: Any, response: Any) -> Any: if not self.deanonymize: @@ -753,13 +796,15 @@ class PrivaiteGuardrail(CustomGuardrail): if args: tc_index = getattr(tool_call, "index", 0) or 0 # Pass `finished` so a fragment on the same chunk as finish_reason - # flushes its held-back tail instead of dropping it. - fn.arguments = restore(("tool", index, tc_index), args, finished) + # flushes its held-back tail instead of dropping it. json_fragment: + # the fragment is JSON source text, so the restored original must + # be spliced in JSON-escaped. + fn.arguments = restore(("tool", index, tc_index), args, finished, True) 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, finished) + function_call.arguments = restore(("fc", index), fc_args, finished, True) async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: Any, response: Any, request_data: dict @@ -775,10 +820,12 @@ class PrivaiteGuardrail(CustomGuardrail): from privaite.streaming.buffer import StreamingDeAnonymizer mapping = PIIMapping() + json_mapping = PIIMapping() for fake, original in fakes.items(): mapping.add(original, fake, "PII") + json_mapping.add(_json_escape(original), fake, "PII") - restorer = _StreamRestorer(mapping, StreamingDeAnonymizer) + restorer = _StreamRestorer(mapping, json_mapping, StreamingDeAnonymizer) async for chunk in response: restorer.restore_chunk(chunk, self._restore_delta) yield chunk diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py index acc9877b555..c37cec6bbe1 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py @@ -117,12 +117,13 @@ class _FakeStreamingDeAnonymizer: buffer isolation observable in tests. """ - def __init__(self, _mapping): + def __init__(self, mapping): + self._fakes = mapping.get_all_fakes() self._buf = "" def feed(self, text): self._buf += text - for fake, original in _FAKES.items(): + for fake, original in self._fakes.items(): self._buf = self._buf.replace(fake, original) lt = self._buf.rfind("<") if lt != -1 and ">" not in self._buf[lt:]: @@ -1279,3 +1280,81 @@ async def test_streaming_tails_on_real_litellm_chunks_stay_serializable(): assert [p["choices"][0]["delta"].get("reasoning_content") for p in payloads] == ["then ", "", _SPECIALS_ORIGINAL) + gr = _make_guardrail() + data = {"metadata": {"privaite_map": _FAKES}} + message = types.SimpleNamespace( + content=None, + tool_calls=[ + types.SimpleNamespace(function=types.SimpleNamespace(arguments='{"name": "", "n": 1}')), + # not JSON: the plain string restore is all there is, byte for byte + types.SimpleNamespace(function=types.SimpleNamespace(arguments="name=")), + ], + function_call=types.SimpleNamespace(arguments='["", {"k": ""}]'), + ) + response = types.SimpleNamespace( + choices=[types.SimpleNamespace(message=message)], + output=[{"type": "function_call", "arguments": '{"who": ""}'}], + ) + + out = await gr.async_post_call_success_hook(data, None, response) + + msg = out.choices[0].message + assert json.loads(msg.tool_calls[0].function.arguments) == {"name": _SPECIALS_ORIGINAL, "n": 1} + assert msg.tool_calls[1].function.arguments == f"name={_SPECIALS_ORIGINAL}" + assert json.loads(msg.function_call.arguments) == [_SPECIALS_ORIGINAL, {"k": "Marie Dupont"}] + assert json.loads(out.output[0]["arguments"]) == {"who": _SPECIALS_ORIGINAL} + + +async def test_post_call_untouched_arguments_are_kept_byte_for_byte(): + # No placeholder in the arguments: the string is returned as received, not + # re-encoded (key order, spacing and escapes stay the provider's). + gr = _make_guardrail() + data = {"metadata": {"privaite_map": _FAKES}} + raw = '{"b":1,\n "a": "x\\u00e9"}' + message = types.SimpleNamespace( + content=None, + tool_calls=[types.SimpleNamespace(function=types.SimpleNamespace(arguments=raw))], + 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.tool_calls[0].function.arguments == raw + + +async def test_streaming_tool_arguments_stay_valid_json_with_special_chars(monkeypatch): + # Streamed argument fragments are JSON source text: the original spliced in + # must be its JSON-escaped form, so the reassembled arguments still parse. + # Content keeps the raw original. The placeholder is split across chunks. + monkeypatch.setitem(_FAKES, "", _SPECIALS_ORIGINAL) + gr = _make_guardrail() + request_data = {"metadata": {"privaite_map": _FAKES}} + + def _delta(args, content=None): + call = types.SimpleNamespace(index=0, function=types.SimpleNamespace(arguments=args)) + return _bare_delta(content=content, tool_calls=[call], function_call=types.SimpleNamespace(arguments=args)) + + async def _source(): + yield _choice_chunk(_delta('{"name": ""}', content="ON_2>"), finish="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 []) + ) + fc_args = "".join(choice.delta.function_call.arguments for chunk in chunks for choice in chunk.choices) + content = "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) + assert json.loads(tool_args) == {"name": _SPECIALS_ORIGINAL} + assert json.loads(fc_args) == {"name": _SPECIALS_ORIGINAL} + assert content == f"Hi {_SPECIALS_ORIGINAL}" From c6505635ecdf91499882efa40464a253ed059390 Mon Sep 17 00:00:00 2001 From: crp4222 <121295348+crp4222@users.noreply.github.com> Date: Mon, 31 Aug 2026 11:51:48 +0200 Subject: [PATCH 11/16] fix(guardrails): advertise the privaite floor the guardrail is tested against The ImportError told users to install privaite>=0.2.4 while the guardrail relies on APIs and a block_entities config field that only later releases provide; 0.4.2 is the version its tests run against. --- litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py index d2bcb3ce476..fdf4d251757 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py @@ -372,7 +372,7 @@ class PrivaiteGuardrail(CustomGuardrail): 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'" + "Install it with: pip install 'privaite>=0.4.2'" ) from exc if self.block_entities and "block_entities" not in PIIConfig.model_fields: From 5939dda0aa66310db8df6707f9522bb2fd641960 Mon Sep 17 00:00:00 2001 From: crp4222 <121295348+crp4222@users.noreply.github.com> Date: Mon, 5 Oct 2026 16:12:36 +0200 Subject: [PATCH 12/16] fix(guardrails): require privaite>=0.7.2 0.7.2 turns off the ONNX Runtime telemetry that older releases send from the process hosting the engine, which here is the proxy itself --- .../guardrails/guardrail_hooks/privaite/privaite.py | 2 +- .../proxy/guardrails/guardrail_hooks/test_privaite.py | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py index fdf4d251757..caacde82d2d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py @@ -372,7 +372,7 @@ class PrivaiteGuardrail(CustomGuardrail): except ImportError as exc: # pragma: no cover raise RuntimeError( "The PrivAiTe guardrail requires the 'privaite' package. " - "Install it with: pip install 'privaite>=0.4.2'" + "Install it with: pip install 'privaite>=0.7.2'" ) from exc if self.block_entities and "block_entities" not in PIIConfig.model_fields: diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py index c37cec6bbe1..5b3212b2a2c 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py @@ -293,6 +293,15 @@ def test_event_hook_always_includes_pre_and_post(): assert "pre_call" in gr3.event_hook and "post_call" in gr3.event_hook +@pytest.mark.asyncio +async def test_engine_missing_privaite_recommends_safe_minimum(monkeypatch: pytest.MonkeyPatch) -> None: + guardrail = _make_guardrail() + monkeypatch.setitem(sys.modules, "privaite.config.schema", None) + + with pytest.raises(RuntimeError, match=r"pip install 'privaite>=0\.7\.2'"): + await guardrail._engine_for(["en"]) + + @pytest.mark.asyncio async def test_engine_is_cached(): gr = _make_guardrail() From e9ea4b60219d4b7a0a24eae6204b73b627b1d834 Mon Sep 17 00:00:00 2001 From: crp4222 <121295348+crp4222@users.noreply.github.com> Date: Wed, 21 Oct 2026 12:59:55 +0200 Subject: [PATCH 13/16] refactor(guardrails): split the privaite hook into typed modules The hook predates the lint gates and added around 400 violations on its own, mostly Any, missing Final and in-place list building. It now follows the llm_shield_proxy layout with payload, request_walk, engine and restore modules, and reaches the privaite package through runtime-checkable protocols since that package is absent from the type-check env Behavior is unchanged. The existing tests pass with two adjustments, tuples where the hook used to hold lists, and the initializer tests building a real LitellmParams instead of a SimpleNamespace --- .../guardrail_hooks/privaite/__init__.py | 22 +- .../guardrail_hooks/privaite/engine.py | 157 +++ .../guardrail_hooks/privaite/payload.py | 61 ++ .../guardrail_hooks/privaite/privaite.py | 956 ++++-------------- .../guardrail_hooks/privaite/request_walk.py | 237 +++++ .../guardrail_hooks/privaite/restore.py | 258 +++++ litellm/types/guardrails.py | 6 +- .../guardrails/guardrail_hooks/privaite.py | 12 +- .../guardrail_hooks/test_privaite.py | 25 +- 9 files changed, 935 insertions(+), 799 deletions(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/privaite/engine.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/privaite/payload.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/privaite/request_walk.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/privaite/restore.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/__init__.py index 76c7989c5cf..630c3c2e1d2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/__init__.py @@ -1,4 +1,6 @@ -from typing import TYPE_CHECKING +from __future__ import annotations + +from typing import TYPE_CHECKING, Final from litellm.types.guardrails import SupportedGuardrailIntegrations @@ -8,26 +10,26 @@ if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams -def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): +def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> PrivaiteGuardrail: import litellm - _privaite_callback = PrivaiteGuardrail( - guardrail_name=guardrail.get("guardrail_name", ""), + _privaite_callback: Final = PrivaiteGuardrail( + guardrail_name=guardrail["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), - block_entities=getattr(litellm_params, "block_entities", None), + preset=litellm_params.preset, + languages=litellm_params.languages, + deanonymize=litellm_params.deanonymize, + block_entities=litellm_params.block_entities, ) litellm.logging_callback_manager.add_litellm_callback(_privaite_callback) return _privaite_callback -guardrail_initializer_registry = { +guardrail_initializer_registry: Final = { SupportedGuardrailIntegrations.PRIVAITE.value: initialize_guardrail, } -guardrail_class_registry = { +guardrail_class_registry: Final = { SupportedGuardrailIntegrations.PRIVAITE.value: PrivaiteGuardrail, } diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/engine.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/engine.py new file mode 100644 index 00000000000..419077cbc61 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/engine.py @@ -0,0 +1,157 @@ +"""The typed boundary to the optional `privaite` package. + +`privaite` is not a LiteLLM dependency, so it is imported lazily and nothing it returns +has a type here. Each object it hands back is checked once against the protocol below +that names what this guardrail calls on it. +""" + +import asyncio +import importlib +from collections.abc import Callable, Iterable, Mapping, Sequence +from types import MappingProxyType, ModuleType +from typing import Final, Protocol, TypeVar, runtime_checkable + +INSTALL_HINT: Final = ( + "The PrivAiTe guardrail requires the 'privaite' package. Install it with: pip install 'privaite>=0.7.2'" +) + +SPACY_MODELS: Final = MappingProxyType( + { + "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", + } +) + +_Narrowed = TypeVar("_Narrowed") # rebind-ok: a TypeVar cannot be declared Final + + +@runtime_checkable +class RestoreMapping(Protocol): + @property + def is_empty(self) -> bool: ... + + def add(self, original: str, fake: str, entity_type: str) -> None: ... + + def get_all_fakes(self) -> Mapping[str, str]: ... + + +@runtime_checkable +class PrivaiteEngine(Protocol): + async def initialize(self) -> None: ... + + async def process_request(self, messages: Sequence[object]) -> tuple[Sequence[object], RestoreMapping]: ... + + async def process_response(self, text: str, mapping: RestoreMapping) -> str: ... + + +@runtime_checkable +class RestoreBuffer(Protocol): + def feed(self, text: str) -> str: ... + + def flush(self) -> str: ... + + +def _module(name: str) -> ModuleType: + try: + return importlib.import_module(name) + except ImportError as exc: + raise RuntimeError(INSTALL_HINT) from exc + + +def _factory(module: ModuleType, name: str) -> Callable[..., object]: + value: Final = getattr(module, name, None) + if not callable(value): + raise TypeError(f"{module.__name__}.{name} is missing. {INSTALL_HINT}") + return value + + +def _narrow(value: object, kind: type[_Narrowed]) -> _Narrowed: + if not isinstance(value, kind): + raise TypeError(f"privaite returned {type(value).__name__} where {kind.__name__} was expected") + return value + + +def _declares_field(model: object, name: str) -> bool: + fields: Final = getattr(model, "model_fields", None) + return isinstance(fields, Mapping) and name in fields + + +def _is_missing_spacy_model(exc: OSError) -> bool: + """spaCy's E050. Permission, disk and network errors must not trigger a download.""" + message: Final = str(exc) + return "E050" in message or "Can't find model" in message + + +async def _download_spacy_models(languages: Sequence[str]) -> None: + """Pulls hundreds of MB through pip's synchronous machinery, so it runs off the event loop.""" + download: Final = _factory(importlib.import_module("spacy.cli"), "download") + for language in languages: + model = SPACY_MODELS.get(language) + if model: + await asyncio.to_thread(download, model) + + +async def build_engine( + preset: str, languages: Sequence[str], deanonymize: bool, block_entities: Sequence[str] +) -> PrivaiteEngine: + schema: Final = _module("privaite.config.schema") + make_engine: Final = _factory(_module("privaite.pii.engine"), "PIIEngine") + pii_config: Final = _factory(schema, "PIIConfig") + if block_entities and not _declares_field(pii_config, "block_entities"): + # PIIConfig allows extra fields, so an older privaite would accept block_entities, + # ignore it and forward the PII. Refusing is the only safe answer. + raise RuntimeError( + "block_entities is set but the installed privaite does not support it; " + "upgrade privaite to a version that enforces pii.block_entities." + ) + presidio: Final = _factory(schema, "PresidioDetectorConfig")(enabled=True, languages=list(languages)) + config: Final = pii_config( + enabled=True, + preset=preset, + detectors=_factory(schema, "DetectorsConfig")(presidio=presidio), + anonymization=_factory(schema, "AnonymizationConfig")(method="placeholder"), + deanonymization=_factory(schema, "DeanonymizationConfig")(enabled=deanonymize), + block_entities=list(block_entities), + ) + engine: Final = _narrow(make_engine(config), PrivaiteEngine) + try: + await engine.initialize() + except OSError as exc: + if not _is_missing_spacy_model(exc): + raise + await _download_spacy_models(languages) + retried: Final = _narrow(make_engine(config), PrivaiteEngine) + await retried.initialize() + return retried + return engine + + +def blocked_error() -> type[Exception]: + """The exception privaite raises when a request carries a type listed in block_entities.""" + error: Final = getattr(_module("privaite.pii.engine"), "PIIBlockedError", None) + if not isinstance(error, type) or not issubclass(error, Exception): + raise TypeError(f"privaite.pii.engine.PIIBlockedError is missing. {INSTALL_HINT}") + return error + + +def restore_mapping(pairs: Iterable[tuple[str, str]]) -> RestoreMapping: + """A privaite mapping holding each (fake, original) pair.""" + mapping: Final = _narrow(_factory(_module("privaite.pii.mapping"), "PIIMapping")(), RestoreMapping) + for fake, original in pairs: + mapping.add(original, fake, "PII") + return mapping + + +def restore_buffer_factory() -> Callable[[RestoreMapping], RestoreBuffer]: + """privaite's streaming de-anonymizer, which holds back a placeholder split across chunks.""" + make_buffer: Final = _factory(_module("privaite.streaming.buffer"), "StreamingDeAnonymizer") + + def _buffer(mapping: RestoreMapping) -> RestoreBuffer: + return _narrow(make_buffer(mapping), RestoreBuffer) + + return _buffer diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/payload.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/payload.py new file mode 100644 index 00000000000..e3b6df66a3c --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/payload.py @@ -0,0 +1,61 @@ +import functools +from collections.abc import Callable, Sequence +from typing import Final, TypeAlias + +JsonObject: TypeAlias = dict[str, object] + +JsonArray: TypeAlias = list[object] + +Write: TypeAlias = Callable[[str], None] + +Slot: TypeAlias = tuple[str, Write] + + +def as_object(value: object) -> JsonObject | None: + """`value` as a JSON object, or None. A JSON object's keys are strings, which isinstance cannot state.""" + return value if isinstance(value, dict) else None # pyright: ignore[reportUnknownVariableType] # JSON keys are str + + +def as_array(value: object) -> JsonArray | None: + """`value` as a JSON array, or None. See `as_object`.""" + return value if isinstance(value, list) else None # pyright: ignore[reportUnknownVariableType] # see as_object + + +def read_field(holder: object, name: str) -> object: + """Reads one field from a dict or from an object. + + LiteLLM hands the hooks Pydantic models on some paths and plain dicts on others, + so every walk here reads both shapes. + """ + fields: Final = as_object(holder) + if fields is not None: + return fields.get(name) + return getattr(holder, name, None) + + +def read_list(holder: object, name: str) -> Sequence[object]: + """Reads a list field from a dict or an object. Anything else reads as empty.""" + return as_array(read_field(holder, name)) or () + + +def write_field(holder: object, name: str, value: object) -> None: + """Writes one field back into a dict or an object. Pairs with read_field.""" + fields: Final = as_object(holder) + if fields is not None: + fields[name] = value + else: + setattr(holder, name, value) + + +def field_writer(fields: JsonObject, key: str) -> Write: + return functools.partial(fields.__setitem__, key) + + +def item_writer(items: JsonArray, index: int) -> Write: + return functools.partial(items.__setitem__, index) + + +def index_of(holder: object) -> int: + """A streamed choice or tool call is matched across chunks by its index, 0 when absent.""" + index: Final = read_field(holder, "index") + return index if isinstance(index, int) else 0 diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py index caacde82d2d..d31911a9555 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py @@ -1,850 +1,272 @@ """ PrivAiTe guardrail for the LiteLLM proxy. -Runs PrivAiTe's engine in-process inside LiteLLM. The pre-call hook anonymizes the -request 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. The request surface scanned matches the PrivAiTe -core: chat `messages`; the Responses API `input` (every item type: role message, -tool-call arguments, tool output including `custom_tool_call_output.output` as a -list of `{type, text}` parts, typed action carriers, bare content parts and bare -strings; opaque/binary items relayed whole) and `prompt.variables`; the -`/v1/completions` `prompt` (string or batch list) and `suffix`; the chat -`prediction.content` and `web_search_options.user_location` auxiliary fields (input -side only, nothing to restore; tokenized integer-array inputs pass through -unscanned). Restore covers `content`, the reasoning trace, the refusal and the audio -transcript besides tool/function arguments. Chat streaming is restored too: text -content, the reasoning trace, the refusal, the audio transcript, 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. +Runs PrivAiTe's engine in-process. The pre-call hook anonymizes the request and stashes +the reversible map in the request metadata, and the post-call hooks consume it to restore +the real values in the response, streaming included. The request surface scanned matches +the PrivAiTe core: chat `messages`, every Responses `input` item shape and +`prompt.variables`, the `/v1/completions` `prompt` and `suffix`, and the chat +`prediction.content` and `web_search_options.user_location` fields. Responses API +streaming restore is not implemented yet. 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. +rejected with a 400 before anything is forwarded. The error names the types, never the +values. -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. +There is no detection or masking logic here, it all lives in the `privaite` package, +which is imported lazily so this module stays importable for the guardrail registry +without it. See https://github.com/crp4222/PrivAiTe. """ from __future__ import annotations import asyncio -import copy -import json -from collections.abc import Callable, Iterator -from typing import Any, cast +import itertools +from collections.abc import AsyncGenerator, AsyncIterable, Iterable, Iterator, Sequence +from typing import TYPE_CHECKING, Final, TypeAlias + +from typing_extensions import TypedDict, Unpack from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.guardrails import GuardrailEventHooks, Mode -_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", -} - -# Responses `input` items relayed byte-for-byte: an opaque/binary payload -# (encrypted reasoning, a screenshot, a generated image) or a tool/pointer -# definition with no user text. Parity with the PrivAiTe gateway scrubber. -_RESPONSES_OPAQUE_TYPES = frozenset( - { - "reasoning", - "compaction", - "compaction_trigger", - "computer_call_output", - "image_generation_call", - "item_reference", - "mcp_list_tools", - "tool_search_call", - "tool_search_output", - "additional_tools", - } +from .engine import ( + PrivaiteEngine, + RestoreMapping, + blocked_error, + build_engine, + restore_buffer_factory, + restore_mapping, +) +from .payload import JsonObject, as_array, as_object, read_field, read_list, write_field +from .request_walk import ( + aux_field_slots, + has_aux_fields, + prompt_variable_slots, + prompt_variables, + responses_input_slots, +) +from .restore import ( + StreamRestorer, + json_escape, + restore_message, + restore_responses_output, ) -# Responses items whose user data sits in a known structured field (a typed -# command/action, a patch diff, search queries and results, interpreter code -# and logs). Every listed field is walked leaf by leaf. -_RESPONSES_DATA_FIELDS: dict[str, tuple[str, ...]] = { - "computer_call": ("action", "actions"), - "local_shell_call": ("action",), - "shell_call": ("action",), - "web_search_call": ("action",), - "apply_patch_call": ("operation",), - "file_search_call": ("queries", "results"), - "code_interpreter_call": ("code", "outputs"), - "program": ("code",), - "program_output": ("result",), -} +if TYPE_CHECKING: + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.utils import ModelResponseStream -# Content/output parts whose payload is binary (base64 or a file id), not text: -# scrubbing them would corrupt the payload without removing anything a text -# detector could find, so the whole part is relayed whole. -_BINARY_PART_TYPES = frozenset( - { - "input_image", - "input_file", - "input_audio", - "image", - "output_image", - "computer_screenshot", - } -) +RESTORE_MAP_KEY: Final = "privaite_map" -# Text-bearing fields on an unknown Responses item shape (e.g. an mcp_call -# carries both `arguments` and `output`): every one is scanned, not just the -# first, so a second field is never left raw. -_GENERIC_TEXT_FIELDS = ("output", "arguments", "input", "text", "reason") +_PRESETS: Final = frozenset({"light", "onnx"}) + +_FALSE_STRINGS: Final = frozenset({"false", "0", "no", ""}) + +_REQUIRED_HOOKS: Final = (GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call) + +_EngineKey: TypeAlias = tuple[str, tuple[str, ...], bool, tuple[str, ...]] -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) +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + pass -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) +def _parse_block_entities(raw: str | Sequence[str] | None) -> tuple[str, ...]: + entries: Final = raw.split(",") if isinstance(raw, str) else tuple(str(entry) for entry in raw or ()) + return tuple(entry.strip() for entry in entries if entry.strip()) -def _json_escape(value: str) -> str: - """The JSON string-literal encoding of value, without the surrounding - quotes: what a restored original must look like when spliced into streamed - argument fragments. Mirrors privaite.streaming.buffer.json_escape.""" - return json.dumps(value, ensure_ascii=False)[1:-1] +def _with_required_hooks( + hook: GuardrailEventHooks | Sequence[GuardrailEventHooks] | None, +) -> list[GuardrailEventHooks]: # mutable-ok: CustomGuardrail.event_hook is a list + """Both hooks are mandatory. A `mode: post_call` config alone would skip the + anonymization and forward raw PII to the model.""" + configured: Final = ((hook,) if hook else ()) if isinstance(hook, str) else tuple(hook or ()) + return list(dict.fromkeys((*configured, *_REQUIRED_HOOKS))) -async def _restore_json_tree(engine: Any, value: Any, mapping: Any) -> Any: - """Restore every string leaf of a parsed JSON document (keys are not - rewritten, mirroring the core engine's walker).""" - if isinstance(value, str): - return await engine.process_response(value, mapping) - if isinstance(value, dict): - return {key: await _restore_json_tree(engine, item, mapping) for key, item in value.items()} - if isinstance(value, list): - return [await _restore_json_tree(engine, item, mapping) for item in value] - return value +def _drop_restore_map(request: object) -> None: + metadata: Final = as_object(read_field(request, "metadata")) + if metadata is not None: + metadata.pop(RESTORE_MAP_KEY, None) -async def _restore_arguments(engine: Any, arguments: str, mapping: Any) -> str: - """Restore a tool/function call's `arguments`: a JSON document inside a - string. Plain substitution would splice a raw quote, backslash or newline - from the original into a JSON string literal and the client's json.loads of - the arguments would fail; restore on the parsed tree and re-encode instead - (core parity: PIIEngine._deanonymize_arguments). When nothing changes, or - the string is not JSON, the previous behaviour is kept byte for byte.""" - plain = await engine.process_response(arguments, mapping) - if plain == arguments: - return arguments - try: - parsed = json.loads(arguments) - except ValueError: - # Not JSON: the plain string restore is all there is. - return plain - return json.dumps(await _restore_json_tree(engine, parsed, mapping), ensure_ascii=False) +def _pop_restore_map(request: JsonObject) -> tuple[tuple[str, str], ...]: + """Consumes the reversible map, so the originals do not linger in metadata the proxy may persist to spend logs.""" + metadata: Final = as_object(request.get("metadata")) + restore_map: Final = as_object(metadata.pop(RESTORE_MAP_KEY, None)) if metadata is not None else None + pairs: Final = (restore_map or {}).items() + return tuple((fake, original) for fake, original in pairs if isinstance(original, str)) -def _append_nested(delta: object, carrier: str, field: str, remaining: str) -> None: - """Append text to delta.., creating the carrier as a plain - dict when the delta has none (a bare finish delta carries no function_call - or audio object).""" - holder = _obj_get(delta, carrier) - if holder is None: - _obj_set(delta, carrier, {field: remaining}) - else: - _obj_set(holder, field, (_obj_get(holder, field) or "") + remaining) +def _pseudo_message(text: str) -> JsonObject: + """A string leaf wrapped as a message, so it crosses the engine's single choke point like chat content.""" + return {"role": "user", "content": text} -def _append_tool_remainder(delta: object, tc_index: int, remaining: str) -> None: - """Append a held argument fragment under its tool_call index, creating the - slot when the delta has none: the client reassembles arguments per index.""" - calls = _obj_get(delta, "tool_calls") - if calls is None: - calls = [] - _obj_set(delta, "tool_calls", calls) - for call in calls: - if (_obj_get(call, "index") or 0) == tc_index: - _append_nested(call, "function", "arguments", remaining) - return - calls.append({"index": tc_index, "function": {"arguments": remaining}}) +def _json_escaped(pairs: Iterable[tuple[str, str]]) -> Iterator[tuple[str, str]]: + """The argument channels stream JSON source text, so their originals are restored JSON-escaped. - -def _append_remainder(delta: object, key: tuple, remaining: str) -> None: - """Append a restore buffer's held tail onto a streamed delta. The key is - (kind, choice index[, tool_call index]) as built by the streaming hook.""" - kind = key[0] - if kind == "tool": - _append_tool_remainder(delta, key[2], remaining) - elif kind == "fc": - _append_nested(delta, "function_call", "arguments", remaining) - elif kind == "audio": - _append_nested(delta, "audio", "transcript", remaining) - else: - # content and the text fields (reasoning_content, reasoning, refusal) - _obj_set(delta, kind, (_obj_get(delta, kind) or "") + remaining) - - -_DELTA_CHANNELS = ( - "content", - "reasoning_content", - "reasoning", - "refusal", - "audio", - "tool_calls", - "function_call", -) - - -def _drain_chunk(template: object, key: tuple, remaining: str) -> object: - """A trailing chunk carrying one held tail after the provider closed the - stream without a finish_reason: a clone of the last chunk of that choice - (same id, model and object as the rest of the stream) reduced to that one - choice, its delta cleared, no finish_reason and no usage.""" - chunk = copy.deepcopy(template) - index = key[1] - choice = next(c for c in _obj_get(chunk, "choices") or [] if (_obj_get(c, "index") or 0) == index) - _obj_set(chunk, "choices", [choice]) - if _obj_get(chunk, "usage") is not None: - _obj_set(chunk, "usage", None) - _obj_set(choice, "finish_reason", None) - delta = _obj_get(choice, "delta") - for field in _DELTA_CHANNELS: - if _obj_get(delta, field) is not None: - _obj_set(delta, field, None) - _append_remainder(delta, key, remaining) - return chunk - - -class _StreamRestorer: - """Restore state for one streamed response. - - 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. The buffer holds back partial - placeholders that span chunk boundaries; whatever remains is flushed onto - the chunk that carries finish_reason (by the channel's own restore when - the finish delta carries it, by the sweep in restore_chunk when it does - not: the common finish shape is a bare `delta: {}`), or, when the provider - closes the stream without a finish_reason, drained as trailing chunks - cloned from the last chunk seen for that choice. A held tail is never - dropped.""" - - def __init__(self, mapping: object, json_mapping: object, make_buffer: Callable[..., object]) -> None: - self._mapping = mapping - # For the argument channels: their fragments are JSON source text, so the - # original has to arrive JSON-escaped or the client's parse of the - # reassembled arguments fails on a quote, a backslash or a newline. - # (Placeholders need no escaping, so the fakes are the same on both sides.) - self._json_mapping = json_mapping - self._make_buffer = make_buffer - self._buffers: dict = {} - self._last_chunk: dict = {} - - def restore(self, key: tuple, text: str, finished: bool, json_fragment: bool = False) -> str: - if key not in self._buffers: - self._buffers[key] = self._make_buffer(self._json_mapping if json_fragment else self._mapping) - deanon = self._buffers[key] - out = deanon.feed(text) if text else "" - if finished: - # The channel is closed: drop its buffer so a sweep never flushes - # it a second time. - out += self._buffers.pop(key).flush() - return out - - def close_open(self, index: int | None) -> tuple[tuple[tuple, str], ...]: - """Flush and drop every buffer still open for this choice (every choice - when index is None); returns the non-empty tails.""" - tails = [] - for key in [k for k in self._buffers if index is None or k[1] == index]: - remaining = self._buffers.pop(key).flush() - if remaining: - tails.append((key, remaining)) - return tuple(tails) - - def restore_chunk(self, chunk: object, restore_delta: Callable[..., None]) -> None: - """Restore every choice of one chunk in place; on a finished choice, - append the tails of the channels its finish delta does not carry.""" - 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 - self._last_chunk[index] = chunk - finished = getattr(choice, "finish_reason", None) is not None - restore_delta(delta, index, finished, self.restore) - if finished: - for key, remaining in self.close_open(index): - _append_remainder(delta, key, remaining) - - def drain(self) -> Iterator[object]: - """Trailing chunks for the tails still held after the provider closed - the stream without a finish_reason for their choice.""" - for key, remaining in self.close_open(None): - template = self._last_chunk.get(key[1]) - if template is not None: - yield _drain_chunk(template, key, remaining) + Placeholders need no escaping, so the fakes are the same on both sides. + """ + for fake, original in pairs: + yield fake, json_escape(original) 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", "") + def __init__( + self, + preset: str | None = None, + languages: str | None = None, + deanonymize: bool | str | None = True, + block_entities: str | Sequence[str] | None = None, + supported_event_hooks: Sequence[GuardrailEventHooks] = _REQUIRED_HOOKS, + **kwargs: Unpack[_CustomGuardrailOptions], # noqa: LIT008 # kwargs-ok: forwarded verbatim to CustomGuardrail.__init__ + ) -> None: + self.preset: Final = preset if preset in _PRESETS else "onnx" + self.languages: Final = languages or "en,fr" + self.deanonymize: Final = ( + deanonymize if isinstance(deanonymize, bool) else str(deanonymize).strip().lower() not in _FALSE_STRINGS ) - # 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, - 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() + self.block_entities: Final = _parse_block_entities(block_entities) + super().__init__( # pyright: ignore[reportUnknownMemberType] # CustomGuardrail.__init__ takes untyped **kwargs + supported_event_hooks=list(supported_event_hooks), **kwargs + ) + if not isinstance(self.event_hook, Mode): + self.event_hook = _with_required_hooks(self.event_hook) + self._engine: PrivaiteEngine | None = None + self._engine_key: _EngineKey | None = None + self._lock: Final = asyncio.Lock() - def _languages(self) -> list: - langs = [lang.strip() for lang in self.languages.split(",") if lang.strip()] - return langs or ["en"] + def _languages(self) -> tuple[str, ...]: + languages: Final = tuple(language.strip() for language in self.languages.split(",") if language.strip()) + return languages 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, tuple(self.block_entities)) + async def _engine_for(self, languages: Sequence[str]) -> PrivaiteEngine: + key: Final = (self.preset, tuple(languages), self.deanonymize, self.block_entities) 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.7.2'" - ) 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: - await engine.initialize() - except OSError as exc: - # spaCy models not present yet: download them once, then retry. - # The download is synchronous pip machinery pulling hundreds of - # MB; run it off the event loop so it does not stall every other - # request in this proxy worker. Only the missing-model OSError - # (spaCy E050) takes this path; permission/disk/network errors - # re-raise instead of triggering a pointless download. - message = str(exc) - if "E050" not in message and "Can't find model" not in message: - raise - from spacy.cli import download - - for lang in languages: - model = _LANG_MODELS.get(lang) - if model: - await asyncio.to_thread(download, model) - engine = PIIEngine(config) - await engine.initialize() - + engine: Final = await build_engine(self.preset, languages, self.deanonymize, self.block_entities) self._engine = engine self._engine_key = key return engine - def _overwrite_snapshot_field(self, data: dict, field: str, new_value: Any) -> None: - # A plain `data[field] = ...` rebind leaks for a top-level string field - # (`input`, `suffix`, a string `prompt`): the proxy snapshots the request - # body by shallow-copying data before this hook, so the original string - # stays in proxy_server_request.body[field]. Overwrite the snapshot copy - # too. (Lists and dicts 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 field in body: - body[field] = new_value + async def _anonymize_request(self, data: JsonObject, engine: PrivaiteEngine) -> RestoreMapping | None: + """Anonymizes every scanned request surface in place under one shared mapping. - @staticmethod - def _make_setter(container: Any, key: Any) -> Callable[[str], None]: - """A write-back that drops the scrubbed string into container[key]. The - container (a list slot or a dict field) is aliased by the proxy's shallow - body snapshot, so the in-place write lands in the snapshot too.""" - - def _set(value: str) -> None: - container[key] = value - - return _set - - def _toplevel_setter(self, data: dict, field: str) -> Callable[[str], None]: - """A write-back for a top-level string field that also overwrites the - detached body snapshot copy (see _overwrite_snapshot_field).""" - - def _set(value: str) -> None: - data[field] = value - self._overwrite_snapshot_field(data, field, value) - - return _set - - def _add_leaf(self, batch: list, setters: list, text: Any, setter: Callable[[str], None]) -> None: - """Register one non-empty string leaf: a pseudo-message carrying it (so it - flows through the engine's single choke point exactly like chat content, - block gate included) plus the write-back that restores the scrubbed value.""" - if isinstance(text, str) and text: - batch.append({"role": "user", "content": text}) - setters.append(setter) - - def _collect_content(self, batch: list, setters: list, content: Any, setter: Callable[[str], None]) -> None: - """A role-message `content`: a bare string, or a list of parts whose - `text`/`refusal` fields carry user text (a bare string in the list is - user text too). Binary parts carry nothing a text detector can find.""" - if isinstance(content, str): - self._add_leaf(batch, setters, content, setter) - return - if not isinstance(content, list): - return - for idx, part in enumerate(content): - if isinstance(part, str): - self._add_leaf(batch, setters, part, self._make_setter(content, idx)) - elif isinstance(part, dict) and part.get("type") not in _BINARY_PART_TYPES: - for field in ("text", "refusal"): - self._add_leaf(batch, setters, part.get(field), self._make_setter(part, field)) - - def _collect_data_value(self, batch: list, setters: list, value: Any, setter: Callable[[str], None]) -> None: - """Walk a tool payload (a string, a list of parts, or arbitrary nested - JSON) registering every string leaf, but relay binary parts whole. Object - keys are never scanned, matching the documented boundary. Iterative - worklist, no recursion; children are pushed reversed so leaves register - in the same depth-first order as the recursive walk they replace.""" - stack: list[tuple[Any, Callable[[str], None]]] = [(value, setter)] - while stack: - node, write_back = stack.pop() - if isinstance(node, str): - self._add_leaf(batch, setters, node, write_back) - elif isinstance(node, list): - stack.extend((node[idx], self._make_setter(node, idx)) for idx in range(len(node) - 1, -1, -1)) - elif isinstance(node, dict) and node.get("type") not in _BINARY_PART_TYPES: - stack.extend((sub, self._make_setter(node, key)) for key, sub in reversed(node.items())) - - def _collect_responses_item(self, batch: list, setters: list, container: list, idx: int, item: Any) -> None: - """Register the scannable text on ONE Responses `input` list item, by its - shape (bare string, role message, typed tool call/output, or an unknown - shape scanned field by field). Opaque/binary items are relayed whole.""" - if isinstance(item, str): - self._add_leaf(batch, setters, item, self._make_setter(container, idx)) - return - if not isinstance(item, dict): - return - if "role" in item and "content" in item: - self._collect_content(batch, setters, item["content"], self._make_setter(item, "content")) - return - itype = item.get("type") or "" - if itype in _RESPONSES_OPAQUE_TYPES: - return - if not self._collect_typed_item(batch, setters, item, itype): - self._collect_generic_item(batch, setters, item) - - def _collect_typed_item(self, batch: list, setters: list, item: dict, itype: str) -> bool: - """Register the text on an item whose user data sits in a type-specific - field; returns False to hand an unknown shape to the generic scan. Covers - function_call/custom_tool_call, the typed action carriers, and every - `*_output` tool output (whose `output` can be a list of {type, text} - parts, not just a string).""" - if itype == "function_call" and isinstance(item.get("arguments"), str): - self._add_leaf(batch, setters, item["arguments"], self._make_setter(item, "arguments")) - return True - if itype == "custom_tool_call" and isinstance(item.get("input"), str): - self._add_leaf(batch, setters, item["input"], self._make_setter(item, "input")) - return True - fields = _RESPONSES_DATA_FIELDS.get(itype) - if fields is not None: - for field in fields: - if item.get(field) is not None: - self._collect_data_value(batch, setters, item[field], self._make_setter(item, field)) - return True - if itype.endswith("_output") and "output" in item: - self._collect_data_value(batch, setters, item["output"], self._make_setter(item, "output")) - return True - return False - - def _collect_generic_item(self, batch: list, setters: list, item: dict) -> None: - """Fallback scan for an unknown item shape: every known text-bearing field - plus a role-less `content`, so a second field is never left raw.""" - for field in _GENERIC_TEXT_FIELDS: - self._add_leaf(batch, setters, item.get(field), self._make_setter(item, field)) - if isinstance(item.get("content"), (str, list)): - self._collect_content(batch, setters, item["content"], self._make_setter(item, "content")) - - def _collect_responses_input(self, data: dict, batch: list, setters: list) -> None: - input_value = data.get("input") - if isinstance(input_value, str) and input_value: - self._add_leaf(batch, setters, input_value, self._toplevel_setter(data, "input")) - elif isinstance(input_value, list): - for idx, item in enumerate(input_value): - self._collect_responses_item(batch, setters, input_value, idx, item) - - @staticmethod - def _prompt_variables(data: dict) -> dict | None: - """Responses prompt-template variables carry user data (the template - id/version do not). None when there are no scannable variables; - /v1/completions sends `prompt` as a string, which is not this surface.""" - prompt = data.get("prompt") - if isinstance(prompt, dict) and isinstance(prompt.get("variables"), dict): - return prompt["variables"] - return None - - @staticmethod - def _has_aux_fields(data: dict) -> bool: - """True when the request carries one of the auxiliary text fields scanned - by _collect_aux_fields; keeps the pre-call early return from skipping a - request whose only user text sits in those fields.""" - prediction = data.get("prediction") - if isinstance(prediction, dict) and "content" in prediction: - return True - web_search = data.get("web_search_options") - if isinstance(web_search, dict) and "user_location" in web_search: - return True - suffix = data.get("suffix") - if isinstance(suffix, str) and suffix: - return True - # /v1/completions user text: a string prompt or the batch list shape. - # (A dict prompt is the Responses template, handled by _prompt_variables.) - prompt = data.get("prompt") - return isinstance(prompt, (str, list)) and bool(prompt) - - def _collect_prompt_variables(self, data: dict, batch: list, setters: list) -> None: - variables = self._prompt_variables(data) - if variables is None: - return - for key, value in variables.items(): - self._collect_data_value(batch, setters, value, self._make_setter(variables, key)) - - def _collect_aux_fields(self, data: dict, batch: list, setters: list) -> None: - """The request-side text fields outside messages/input that LiteLLM - forwards verbatim (core parity): chat `prediction.content`, - `web_search_options.user_location`, and the completions `prompt` (string - or batch list) and `suffix`. Request inputs only, nothing to restore.""" - prediction = data.get("prediction") - if isinstance(prediction, dict) and prediction.get("content") is not None: - self._collect_data_value(batch, setters, prediction["content"], self._make_setter(prediction, "content")) - web_search = data.get("web_search_options") - if isinstance(web_search, dict) and web_search.get("user_location") is not None: - self._collect_data_value( - batch, setters, web_search["user_location"], self._make_setter(web_search, "user_location") - ) - suffix = data.get("suffix") - if isinstance(suffix, str) and suffix: - self._add_leaf(batch, setters, suffix, self._toplevel_setter(data, "suffix")) - prompt = data.get("prompt") - if isinstance(prompt, str) and prompt: - self._add_leaf(batch, setters, prompt, self._toplevel_setter(data, "prompt")) - elif isinstance(prompt, list): - # The batch shape: string leaves are scrubbed, a tokenized - # (integer-array) prompt passes through unscanned as documented. - for idx, part in enumerate(prompt): - self._add_leaf(batch, setters, part, self._make_setter(prompt, idx)) - - async def _anonymize_request(self, data: dict, engine: Any) -> Any: - """Anonymize every scanned request surface in place under ONE shared - mapping (chat `messages`, the Responses `input` and `prompt.variables`, - and the auxiliary prompt/suffix/prediction/user_location fields), so no - source is left untouched when a crafted request carries several. Every - scanned string flows through the engine's single choke point, so the block - gate and the fail-closed policy apply unchanged. Returns the mapping, or - None if there was nothing to anonymize.""" - messages = data.get("messages") - msg_list = messages if isinstance(messages, list) else [] - # Chat messages are scanned natively (multimodal content, bare-string - # lists, per-message block gate); everything else is registered as string - # leaves appended after them, sharing this one process_request pass. - batch: list = list(msg_list) - setters: list = [] - self._collect_responses_input(data, batch, setters) - self._collect_prompt_variables(data, batch, setters) - self._collect_aux_fields(data, batch, setters) - - if not batch: + Chat messages go to the engine as they are. Every other string leaf is appended to + the same batch as a pseudo-message, so one process_request pass covers a crafted + request that carries several surfaces, under the same block gate and fail-closed policy. + """ + messages: Final = as_array(data.get("messages")) or [] + slots: Final = tuple( + itertools.chain(responses_input_slots(data), prompt_variable_slots(data), aux_field_slots(data)) + ) + if not messages and not slots: return None - - 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] - for setter, anon in zip(setters, anonymized[n:]): - setter(anon.get("content", "")) + anonymized, mapping = await engine.process_request([*messages, *(_pseudo_message(text) for text, _ in slots)]) + count: Final = len(messages) + if messages: + messages[:] = anonymized[:count] + for (_, write), leaf_message in zip(slots, anonymized[count:]): + content = read_field(leaf_message, "content") + write(content if isinstance(content, str) else "") 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) - + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: JsonObject, + call_type: str, + ) -> JsonObject: + # metadata is caller controlled, so an incoming privaite_map is never trusted: + # the post-call hooks may only restore from a map this pre-call pass produced. + _drop_restore_map(data) if ( not data.get("messages") and not data.get("input") - and self._prompt_variables(data) is None - and not self._has_aux_fields(data) + and prompt_variables(data) is None + and not has_aux_fields(data) ): return data - engine = await self._engine_for(self._languages()) - from privaite.pii.engine import PIIBlockedError - + engine: Final = await self._engine_for(self._languages()) + blocked: Final = blocked_error() 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. + mapping: Final = await self._anonymize_request(data, engine) + except blocked as exc: 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 and not data.get("background"): - # 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. - # Background requests are excluded: their result is fetched by a later - # poll the post-call hooks never see, so the map would persist in - # request state with no consumer; those responses keep placeholders. - data.setdefault("metadata", {})["privaite_map"] = dict(mapping.get_all_fakes()) + # A background request's result is fetched by a later poll the post-call hooks + # never see, so its map would stay in request state with no consumer. + if mapping is None or not self.deanonymize or mapping.is_empty or data.get("background"): + return data + restore_map: Final = dict(mapping.get_all_fakes()) + metadata: Final = as_object(data.get("metadata")) + if metadata is None: + write_field(data, "metadata", {RESTORE_MAP_KEY: restore_map}) + else: + metadata[RESTORE_MAP_KEY] = restore_map return data - async def _restore_audio_transcript(self, message: Any, engine: Any, mapping: Any) -> None: - """Restore the audio transcript on a response message (restore parity with - the core: an audio reply carries its text in audio.transcript, where the - model echoes placeholders like anywhere else).""" - audio = getattr(message, "audio", None) - if audio is None: - return - transcript = _obj_get(audio, "transcript") - if isinstance(transcript, str) and transcript: - _obj_set(audio, "transcript", await engine.process_response(transcript, mapping)) - - async def _restore_message(self, message: Any, engine: Any, mapping: Any) -> None: - """Restore originals in one response message: content, the reasoning - trace, the refusal, the audio transcript, 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) - # A refusal can quote the request, so it carries placeholders too. - for field in ("reasoning_content", "reasoning", "refusal"): - value = getattr(message, field, None) - if isinstance(value, str) and value: - setattr(message, field, await engine.process_response(value, mapping)) - await self._restore_audio_transcript(message, engine, 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 _restore_arguments(engine, 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 _restore_arguments(engine, 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 _restore_arguments(engine, args, mapping)) - - async def async_post_call_success_hook(self, data: dict, user_api_key_dict: Any, response: Any) -> Any: + async def async_post_call_success_hook( + self, + data: JsonObject, + user_api_key_dict: UserAPIKeyAuth, + response: object, + ) -> object: 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: + pairs: Final = _pop_restore_map(data) + if not pairs: 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) + mapping: Final = restore_mapping(pairs) + engine: Final = await self._engine_for(self._languages()) + for choice in read_list(response, "choices"): + message = read_field(choice, "message") 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) + await restore_message(message, engine, mapping) + await restore_responses_output(response, engine, mapping) return response - def _restore_delta_audio(self, delta: Any, index: int, finished: bool, restore: Any) -> None: - """Feed streamed audio transcript fragments through their own restore - buffer and flush the held tail on the finish chunk, creating the audio - carrier when the finish delta has none so the tail is never dropped.""" - audio = getattr(delta, "audio", None) - fragment = _obj_get(audio, "transcript") if audio is not None else None - if not isinstance(fragment, str): - fragment = "" - if not fragment and not finished: - return - restored = restore(("audio", index), fragment, finished) - if audio is not None: - _obj_set(audio, "transcript", restored) - elif restored: - _obj_set(delta, "audio", {"transcript": restored}) - - def _restore_delta(self, delta: Any, index: int, finished: bool, restore: Any) -> None: - """Restore one streamed delta in place: text content, the reasoning trace, - the refusal, the audio transcript, 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) - # Overwrite whenever there was input text (even if the whole fragment is - # held back, the raw fragment must not stay visible) or the finish flush - # produced text; a terminal chunk with content=None and nothing held back - # keeps its None instead of becoming "". - if content or restored: - delta.content = restored - # A refusal can quote the request, so it carries placeholders too. - for field in ("reasoning_content", "reasoning", "refusal"): - value = getattr(delta, field, None) - if isinstance(value, str) and (value or finished): - setattr(delta, field, restore((field, index), value, finished)) - self._restore_delta_audio(delta, index, finished, restore) - 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 - # Pass `finished` so a fragment on the same chunk as finish_reason - # flushes its held-back tail instead of dropping it. json_fragment: - # the fragment is JSON source text, so the restored original must - # be spliced in JSON-escaped. - fn.arguments = restore(("tool", index, tc_index), args, finished, True) - 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, finished, True) - 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: + self, + user_api_key_dict: UserAPIKeyAuth, + response: AsyncIterable[ModelResponseStream], + request_data: JsonObject, + ) -> AsyncGenerator[ModelResponseStream, None]: + pairs: Final = _pop_restore_map(request_data) + if not self.deanonymize or not pairs: async for chunk in response: yield chunk return - from privaite.pii.mapping import PIIMapping - from privaite.streaming.buffer import StreamingDeAnonymizer - - mapping = PIIMapping() - json_mapping = PIIMapping() - for fake, original in fakes.items(): - mapping.add(original, fake, "PII") - json_mapping.add(_json_escape(original), fake, "PII") - - restorer = _StreamRestorer(mapping, json_mapping, StreamingDeAnonymizer) + restorer: Final = StreamRestorer( + restore_mapping(pairs), + restore_mapping(_json_escaped(pairs)), + restore_buffer_factory(), + ) async for chunk in response: - restorer.restore_chunk(chunk, self._restore_delta) + restorer.restore_chunk(chunk) yield chunk for drained in restorer.drain(): yield drained 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 + request_data: object, + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + # The success and streaming hooks never run when the call fails, so the map is + # dropped here too, before a failure spend-log can persist it. + _drop_restore_map(request_data) diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/request_walk.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/request_walk.py new file mode 100644 index 00000000000..c07b9aa539d --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/request_walk.py @@ -0,0 +1,237 @@ +import functools +from collections.abc import Iterator, Sequence +from types import MappingProxyType +from typing import Final + +from .payload import ( + JsonArray, + JsonObject, + Slot, + Write, + as_array, + as_object, + field_writer, + item_writer, +) + +RESPONSES_OPAQUE_TYPES: Final = frozenset( + { + "reasoning", + "compaction", + "compaction_trigger", + "computer_call_output", + "image_generation_call", + "item_reference", + "mcp_list_tools", + "tool_search_call", + "tool_search_output", + "additional_tools", + } +) +"""Responses `input` items relayed whole: an opaque or binary payload, or a tool or +pointer definition with no user text. Parity with the PrivAiTe gateway scrubber.""" + +RESPONSES_DATA_FIELDS: Final = MappingProxyType( + { + "computer_call": ("action", "actions"), + "local_shell_call": ("action",), + "shell_call": ("action",), + "web_search_call": ("action",), + "apply_patch_call": ("operation",), + "file_search_call": ("queries", "results"), + "code_interpreter_call": ("code", "outputs"), + "program": ("code",), + "program_output": ("result",), + } +) +"""Responses items whose user data sits in known structured fields, walked leaf by leaf.""" + +BINARY_PART_TYPES: Final = frozenset( + { + "input_image", + "input_file", + "input_audio", + "image", + "output_image", + "computer_screenshot", + } +) +"""Parts whose payload is base64 or a file id. Scrubbing one would corrupt it without +removing anything a text detector can find.""" + +GENERIC_TEXT_FIELDS: Final = ("output", "arguments", "input", "text", "reason") +"""Text-bearing fields of an unknown Responses item. Every one is scanned, since an +mcp_call carries both `arguments` and `output`.""" + + +def leaf(text: object, write: Write) -> Iterator[Slot]: + if isinstance(text, str) and text: + yield text, write + + +def content_slots(content: object, write: Write) -> Iterator[Slot]: + """A role message's `content`: a bare string, or parts whose `text` and `refusal` carry user text.""" + if isinstance(content, str): + yield from leaf(content, write) + return + parts: Final = as_array(content) + if parts is None: + return + for index, part in enumerate(parts): + if isinstance(part, str): + yield from leaf(part, item_writer(parts, index)) + continue + fields = as_object(part) + if fields is None or fields.get("type") in BINARY_PART_TYPES: + continue + for name in ("text", "refusal"): + yield from leaf(fields.get(name), field_writer(fields, name)) + + +def data_slots(value: object, write: Write) -> Iterator[Slot]: + """Every string leaf of a tool payload, binary parts relayed whole and object keys never scanned. + + The payload is caller controlled, so it is walked with an explicit stack rather than + by recursion. Children are pushed reversed, so leaves come out depth first. + """ + pending: Final[list[tuple[object, Write]]] = [(value, write)] # mutable-ok: local walk stack, never escapes + while pending: + node, write_back = pending.pop() + if isinstance(node, str): + yield from leaf(node, write_back) + continue + items = as_array(node) + if items is not None: + pending.extend((items[index], item_writer(items, index)) for index in range(len(items) - 1, -1, -1)) + continue + fields = as_object(node) + if fields is not None and fields.get("type") not in BINARY_PART_TYPES: + pending.extend((child, field_writer(fields, key)) for key, child in reversed(fields.items())) + + +def _named_data_slots(fields: JsonObject, names: Sequence[str]) -> Iterator[Slot]: + for name in names: + if fields.get(name) is not None: + yield from data_slots(fields[name], field_writer(fields, name)) + + +def _typed_item_slots(fields: JsonObject, item_type: str) -> Iterator[Slot] | None: + """The text of an item whose user data sits in a type-specific field, or None for an unknown shape.""" + if item_type == "function_call" and isinstance(fields.get("arguments"), str): + return leaf(fields["arguments"], field_writer(fields, "arguments")) + if item_type == "custom_tool_call" and isinstance(fields.get("input"), str): + return leaf(fields["input"], field_writer(fields, "input")) + data_fields: Final = RESPONSES_DATA_FIELDS.get(item_type) + if data_fields is not None: + return _named_data_slots(fields, data_fields) + if item_type.endswith("_output") and "output" in fields: + return data_slots(fields["output"], field_writer(fields, "output")) + return None + + +def _generic_item_slots(fields: JsonObject) -> Iterator[Slot]: + for name in GENERIC_TEXT_FIELDS: + yield from leaf(fields.get(name), field_writer(fields, name)) + content: Final = fields.get("content") + if isinstance(content, str) or as_array(content) is not None: + yield from content_slots(content, field_writer(fields, "content")) + + +def _responses_item_slots(items: JsonArray, index: int) -> Iterator[Slot]: + item: Final = items[index] + if isinstance(item, str): + yield from leaf(item, item_writer(items, index)) + return + fields: Final = as_object(item) + if fields is None: + return + if "role" in fields and "content" in fields: + yield from content_slots(fields["content"], field_writer(fields, "content")) + return + declared_type: Final = fields.get("type") + item_type: Final = declared_type if isinstance(declared_type, str) else "" + if item_type in RESPONSES_OPAQUE_TYPES: + return + typed: Final = _typed_item_slots(fields, item_type) + yield from typed if typed is not None else _generic_item_slots(fields) + + +def _write_top_level(data: JsonObject, key: str, value: str) -> None: + """Writes a top-level string field, and the copy of it in the proxy's request-body snapshot. + + The proxy snapshots the body by shallow copy before the hooks run. A nested list or + dict is shared with that snapshot, but a top-level string is not, so without the + second write the original would stay in proxy_server_request.body. + """ + data[key] = value # rebind-ok: the pre-call hook contract rewrites the request in place + request: Final = as_object(data.get("proxy_server_request")) + body: Final = as_object(request.get("body")) if request is not None else None + if body is not None and key in body: + body[key] = value + + +def top_level_writer(data: JsonObject, key: str) -> Write: + return functools.partial(_write_top_level, data, key) + + +def responses_input_slots(data: JsonObject) -> Iterator[Slot]: + request_input: Final = data.get("input") + if isinstance(request_input, str): + yield from leaf(request_input, top_level_writer(data, "input")) + return + items: Final = as_array(request_input) + if items is None: + return + for index in range(len(items)): + yield from _responses_item_slots(items, index) + + +def prompt_variables(data: JsonObject) -> JsonObject | None: + """Responses prompt-template variables, which carry user data. A string `prompt` is /v1/completions.""" + prompt: Final = as_object(data.get("prompt")) + return as_object(prompt.get("variables")) if prompt is not None else None + + +def prompt_variable_slots(data: JsonObject) -> Iterator[Slot]: + variables: Final = prompt_variables(data) + if variables is None: + return + for key, value in variables.items(): + yield from data_slots(value, field_writer(variables, key)) + + +def has_aux_fields(data: JsonObject) -> bool: + """Whether the request carries a field scanned by `aux_field_slots`.""" + prediction: Final = as_object(data.get("prediction")) + if prediction is not None and "content" in prediction: + return True + web_search: Final = as_object(data.get("web_search_options")) + if web_search is not None and "user_location" in web_search: + return True + suffix: Final = data.get("suffix") + if isinstance(suffix, str) and suffix: + return True + prompt: Final = data.get("prompt") + return (isinstance(prompt, str) and bool(prompt)) or bool(as_array(prompt)) + + +def aux_field_slots(data: JsonObject) -> Iterator[Slot]: + """Request text outside messages and input that LiteLLM forwards verbatim: chat + `prediction.content`, `web_search_options.user_location`, and the completions `suffix` + and `prompt`. A tokenized (integer array) prompt passes through unscanned.""" + prediction: Final = as_object(data.get("prediction")) + if prediction is not None and prediction.get("content") is not None: + yield from data_slots(prediction["content"], field_writer(prediction, "content")) + web_search: Final = as_object(data.get("web_search_options")) + if web_search is not None and web_search.get("user_location") is not None: + yield from data_slots(web_search["user_location"], field_writer(web_search, "user_location")) + yield from leaf(data.get("suffix"), top_level_writer(data, "suffix")) + prompt: Final = data.get("prompt") + if isinstance(prompt, str): + yield from leaf(prompt, top_level_writer(data, "prompt")) + return + parts: Final = as_array(prompt) + if parts is None: + return + for index, part in enumerate(parts): + yield from leaf(part, item_writer(parts, index)) diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/restore.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/restore.py new file mode 100644 index 00000000000..8ac3fdb54a7 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/restore.py @@ -0,0 +1,258 @@ +from __future__ import annotations + +import copy +import json +from collections.abc import Callable, Iterator +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, TypeAlias + +from .engine import PrivaiteEngine, RestoreBuffer, RestoreMapping +from .payload import as_array, as_object, index_of, read_field, read_list, write_field + +if TYPE_CHECKING: + from litellm.types.utils import ModelResponseStream + +TEXT_FIELDS: Final = ("reasoning_content", "reasoning", "refusal") +"""Text fields restored besides `content`. A refusal can quote the request, so it carries placeholders too.""" + +DELTA_CHANNELS: Final = ("content", *TEXT_FIELDS, "audio", "tool_calls", "function_call") + +JSON_CHANNELS: Final = frozenset({"tool", "function_call"}) +"""Channels that stream JSON source text, where a restored original must arrive JSON-escaped.""" + + +def json_escape(value: str) -> str: + """The JSON string-literal encoding of value without its quotes, as privaite.streaming.buffer does.""" + return json.dumps(value, ensure_ascii=False)[1:-1] + + +async def _restore_json_tree(engine: PrivaiteEngine, value: object, mapping: RestoreMapping) -> object: + """Restores every string leaf of a parsed JSON document. Keys are left as they are, as in the core engine.""" + if isinstance(value, str): + return await engine.process_response(value, mapping) + fields: Final = as_object(value) + if fields is not None: + return {key: await _restore_json_tree(engine, item, mapping) for key, item in fields.items()} + items: Final = as_array(value) + if items is not None: + return [await _restore_json_tree(engine, item, mapping) for item in items] + return value + + +async def restore_arguments(engine: PrivaiteEngine, arguments: str, mapping: RestoreMapping) -> str: + """Restores a tool call's `arguments`, a JSON document inside a string. + + Plain substitution would splice a raw quote, backslash or newline from the original + into a JSON string literal, and the client's parse of the arguments would fail. So + the restore runs on the parsed tree and re-encodes it, as PIIEngine does. When nothing + changes, or the string is not JSON, the plain restore is returned as is. + """ + plain: Final = await engine.process_response(arguments, mapping) + if plain == arguments: + return arguments + try: + parsed: Final[object] = json.loads(arguments) # pyright: ignore[reportAny] # json.loads is untyped + except ValueError: + return plain + return json.dumps(await _restore_json_tree(engine, parsed, mapping), ensure_ascii=False) + + +async def _restore_text(holder: object, name: str, engine: PrivaiteEngine, mapping: RestoreMapping) -> None: + value: Final = read_field(holder, name) + if isinstance(value, str) and value: + write_field(holder, name, await engine.process_response(value, mapping)) + + +async def _restore_call_arguments(holder: object, engine: PrivaiteEngine, mapping: RestoreMapping) -> None: + arguments: Final = read_field(holder, "arguments") + if isinstance(arguments, str) and arguments: + write_field(holder, "arguments", await restore_arguments(engine, arguments, mapping)) + + +async def restore_message(message: object, engine: PrivaiteEngine, mapping: RestoreMapping) -> None: + """Restores one response message: content, the text fields, the audio transcript, and tool-call arguments.""" + for name in ("content", *TEXT_FIELDS): + await _restore_text(message, name, engine, mapping) + await _restore_text(read_field(message, "audio"), "transcript", engine, mapping) + for tool_call in read_list(message, "tool_calls"): + await _restore_call_arguments(read_field(tool_call, "function"), engine, mapping) + await _restore_call_arguments(read_field(message, "function_call"), engine, mapping) + + +async def restore_responses_output(response: object, engine: PrivaiteEngine, mapping: RestoreMapping) -> None: + """Restores a Responses API result: output_text blocks and function_call arguments.""" + for item in read_list(response, "output"): + for block in read_list(item, "content"): + await _restore_text(block, "text", engine, mapping) + await _restore_call_arguments(item, engine, mapping) + + +@dataclass(frozen=True, slots=True) +class Channel: + """One streamed segment of one choice. Each keeps its own restore buffer.""" + + kind: str + choice: int + tool_call: int = 0 + + +Restore: TypeAlias = Callable[[Channel, str, bool], str] + + +def _append_text(holder: object, name: str, tail: str) -> None: + existing: Final = read_field(holder, name) + write_field(holder, name, (existing if isinstance(existing, str) else "") + tail) + + +def _append_nested(delta: object, carrier: str, name: str, tail: str) -> None: + """Appends to delta.., creating the carrier when a bare finish delta has none.""" + holder: Final = read_field(delta, carrier) + if holder is None: + write_field(delta, carrier, {name: tail}) + return + _append_text(holder, name, tail) + + +def _append_tool_tail(delta: object, tool_call: int, tail: str) -> None: + """Appends an argument tail under its tool_call index, since the client reassembles arguments per index.""" + calls: Final = as_array(read_field(delta, "tool_calls")) + added: Final = {"index": tool_call, "function": {"arguments": tail}} + if calls is None: + write_field(delta, "tool_calls", [added]) + return + existing: Final = next((call for call in calls if index_of(call) == tool_call), None) + if existing is None: + calls.append(added) + return + _append_nested(existing, "function", "arguments", tail) + + +def append_tail(delta: object, channel: Channel, tail: str) -> None: + """Appends a restore buffer's held tail onto a streamed delta.""" + match channel.kind: + case "tool": + _append_tool_tail(delta, channel.tool_call, tail) + case "function_call": + _append_nested(delta, "function_call", "arguments", tail) + case "audio": + _append_nested(delta, "audio", "transcript", tail) + case _: + _append_text(delta, channel.kind, tail) + + +def _restore_delta_audio(delta: object, choice: int, finished: bool, restore: Restore) -> None: + """Restores a streamed audio transcript, creating the carrier on a finish delta so a held tail is kept.""" + audio: Final = read_field(delta, "audio") + transcript: Final = read_field(audio, "transcript") + fragment: Final = transcript if isinstance(transcript, str) else "" + if not fragment and not finished: + return + restored: Final = restore(Channel("audio", choice), fragment, finished) + if audio is not None: + write_field(audio, "transcript", restored) + elif restored: + write_field(delta, "audio", {"transcript": restored}) + + +def restore_delta(delta: object, choice: int, finished: bool, restore: Restore) -> None: + """Restores one streamed delta: content, the text fields, the audio transcript and argument fragments.""" + content: Final = read_field(delta, "content") + text: Final = content if isinstance(content, str) else "" + restored: Final = restore(Channel("content", choice), text, finished) + # A fragment held back whole must not stay visible raw, but a terminal None with + # nothing held back keeps its None rather than becoming "". + if text or restored: + write_field(delta, "content", restored) + for name in TEXT_FIELDS: + value = read_field(delta, name) + if isinstance(value, str) and (value or finished): + write_field(delta, name, restore(Channel(name, choice), value, finished)) + _restore_delta_audio(delta, choice, finished, restore) + for tool_call in read_list(delta, "tool_calls"): + function = read_field(tool_call, "function") + arguments = read_field(function, "arguments") + if isinstance(arguments, str) and arguments: + channel = Channel("tool", choice, index_of(tool_call)) + write_field(function, "arguments", restore(channel, arguments, finished)) + function_call: Final = read_field(delta, "function_call") + call_arguments: Final = read_field(function_call, "arguments") + if isinstance(call_arguments, str) and call_arguments: + write_field(function_call, "arguments", restore(Channel("function_call", choice), call_arguments, finished)) + + +def drain_chunk(template: ModelResponseStream, channel: Channel, tail: str) -> ModelResponseStream: + """A trailing chunk carrying one held tail, for a stream the provider closed without a finish_reason. + + It is a copy of the last chunk of that choice, so id, model and object match the rest + of the stream, reduced to that one choice with its delta cleared and no usage. + """ + chunk: Final = copy.deepcopy(template) + choice: Final = next(item for item in read_list(chunk, "choices") if index_of(item) == channel.choice) + write_field(chunk, "choices", [choice]) + if read_field(chunk, "usage") is not None: + write_field(chunk, "usage", None) + write_field(choice, "finish_reason", None) + delta: Final = read_field(choice, "delta") + for name in DELTA_CHANNELS: + if read_field(delta, name) is not None: + write_field(delta, name, None) + append_tail(delta, channel, tail) + return chunk + + +class StreamRestorer: + """Restore state for one streamed response. + + With n>1 the provider interleaves chunks of different choices, and tool-call arguments + stream as fragments per tool_call index, so every channel keeps its own buffer and a + placeholder split across chunks reassembles without mixing channels. A held tail is + flushed onto the chunk carrying finish_reason, or, when the provider closes the stream + without one, drained as trailing chunks. It is never dropped. + """ + + def __init__( + self, + mapping: RestoreMapping, + json_mapping: RestoreMapping, + make_buffer: Callable[[RestoreMapping], RestoreBuffer], + ) -> None: + self._mapping: Final = mapping + self._json_mapping: Final = json_mapping + self._make_buffer: Final = make_buffer + self._buffers: Final[dict[Channel, RestoreBuffer]] = {} # mutable-ok: per-stream restore state + self._last_chunk: Final[dict[int, ModelResponseStream]] = {} # mutable-ok: per-stream restore state + + def restore(self, channel: Channel, text: str, finished: bool) -> str: + if channel not in self._buffers: + mapping: Final = self._json_mapping if channel.kind in JSON_CHANNELS else self._mapping + self._buffers[channel] = self._make_buffer(mapping) + restored: Final = self._buffers[channel].feed(text) if text else "" + if not finished: + return restored + return restored + self._buffers.pop(channel).flush() + + def _close(self, choice: int | None) -> tuple[tuple[Channel, str], ...]: + """Flushes and drops every buffer still open for a choice, or for all of them; returns the non-empty tails.""" + closing: Final = tuple(channel for channel in self._buffers if choice is None or channel.choice == choice) + flushed: Final = tuple((channel, self._buffers.pop(channel).flush()) for channel in closing) + return tuple((channel, tail) for channel, tail in flushed if tail) + + def restore_chunk(self, chunk: ModelResponseStream) -> None: + """Restores every choice of a chunk, and appends the tails of the channels a finish delta does not carry.""" + for choice in read_list(chunk, "choices"): + delta = read_field(choice, "delta") + if delta is None: + continue + index = index_of(choice) + self._last_chunk[index] = chunk + finished = read_field(choice, "finish_reason") is not None + restore_delta(delta, index, finished, self.restore) + if finished: + for channel, tail in self._close(index): + append_tail(delta, channel, tail) + + def drain(self) -> Iterator[ModelResponseStream]: + for channel, tail in self._close(None): + template = self._last_chunk.get(channel.choice) + if template is not None: + yield drain_chunk(template, channel, tail) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 0780023a56f..62c2f9b2857 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -50,6 +50,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter impor from litellm.types.proxy.guardrails.guardrail_hooks.ovalix import ( OvalixGuardrailConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.privaite import ( + PrivaiteGuardrailConfigModel, +) from litellm.types.proxy.guardrails.guardrail_hooks.promptguard import ( PromptGuardConfigModel, ) @@ -77,9 +80,6 @@ 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, ) diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/privaite.py b/litellm/types/proxy/guardrails/guardrail_hooks/privaite.py index 065cf99c37b..98db0c9211a 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/privaite.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/privaite.py @@ -1,24 +1,22 @@ -from typing import Optional - from pydantic import Field from .base import GuardrailConfigModel -class PrivaiteGuardrailConfigModel(GuardrailConfigModel): - preset: Optional[str] = Field( +class PrivaiteGuardrailConfigModel(GuardrailConfigModel): # pyright: ignore[reportMissingTypeArgument] # parameterizing adds an optional_params field to the UI form + preset: str | None = Field( default="onnx", description=("Detection preset: 'onnx' (full suite, also detects secrets) or 'light' (faster, Presidio only)."), ) - languages: Optional[str] = Field( + languages: str | None = Field( default="en,fr", description="Comma-separated spaCy languages used for detection.", ) - deanonymize: Optional[bool] = Field( + deanonymize: bool | None = Field( default=True, description="Restore the original PII values in the response.", ) - block_entities: str | list[str] | None = Field( + block_entities: str | tuple[str, ...] | None = Field( default=None, description=( "PII types to reject outright instead of masking: a list, or a " diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py index 5b3212b2a2c..bc53aca9b7a 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py @@ -241,6 +241,7 @@ def test_initialize_guardrail_registers_callback(monkeypatch): PrivaiteGuardrail, initialize_guardrail, ) + from litellm.types.guardrails import LitellmParams added = [] monkeypatch.setattr( @@ -249,7 +250,8 @@ def test_initialize_guardrail_registers_callback(monkeypatch): lambda cb: added.append(cb), ) - litellm_params = types.SimpleNamespace( + litellm_params = LitellmParams( + guardrail="privaite", mode="pre_call", default_on=False, preset="light", @@ -268,17 +270,17 @@ def test_init_normalizes_params(): gr = _make_guardrail(preset="bogus", languages="en, fr ,", deanonymize="false") assert gr.preset == "onnx" assert gr.deanonymize is False - assert gr._languages() == ["en", "fr"] + 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"] + assert gr2._languages() == ("en", "fr") # a non-empty but content-free languages string falls back to ["en"]. - assert _make_guardrail(languages=", ,")._languages() == ["en"] + assert _make_guardrail(languages=", ,")._languages() == ("en",) def test_event_hook_always_includes_pre_and_post(): @@ -808,18 +810,17 @@ def test_block_entities_config_model_field(): 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 == [] + 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 + from litellm.types.guardrails import LitellmParams - litellm_params = types.SimpleNamespace( + litellm_params = LitellmParams( + guardrail="privaite", mode="pre_call", default_on=False, preset="light", @@ -828,7 +829,7 @@ def test_initialize_guardrail_passes_block_entities(): block_entities=["US_SSN"], ) callback = initialize_guardrail(litellm_params, {"guardrail_name": "privaite"}) - assert callback.block_entities == ["US_SSN"] + assert callback.block_entities == ("US_SSN",) @pytest.mark.asyncio From 3486434407e9b2c4c91f6a752366b4b5bbe26ea0 Mon Sep 17 00:00:00 2001 From: crp4222 <121295348+crp4222@users.noreply.github.com> Date: Wed, 21 Oct 2026 13:00:15 +0200 Subject: [PATCH 14/16] fix(guardrails): refuse a privaite request the engine did not return in full The string leaves outside messages are written back by position from the engine's output, and zip stopped at the shorter side. An engine handing back fewer messages than it was given left the remaining leaves unscrubbed and the request went out with them. The count is now checked before anything is written and the request is refused instead --- .../guardrail_hooks/privaite/privaite.py | 7 ++++++- .../guardrail_hooks/test_privaite.py | 18 ++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py index d31911a9555..f382bd84925 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py @@ -171,9 +171,14 @@ class PrivaiteGuardrail(CustomGuardrail): return None anonymized, mapping = await engine.process_request([*messages, *(_pseudo_message(text) for text, _ in slots)]) count: Final = len(messages) + if len(anonymized) != count + len(slots): + raise RuntimeError( + f"privaite returned {len(anonymized)} messages for {count + len(slots)}, " + "refusing to forward a request it did not scrub in full" + ) if messages: messages[:] = anonymized[:count] - for (_, write), leaf_message in zip(slots, anonymized[count:]): + for (_, write), leaf_message in zip(slots, anonymized[count:], strict=True): content = read_field(leaf_message, "content") write(content if isinstance(content, str) else "") return mapping diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py index bc53aca9b7a..8e69b9d2268 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py @@ -717,6 +717,24 @@ async def test_pre_call_anonymizes_both_messages_and_input(): assert out["metadata"]["privaite_map"] == _FAKES +class _EngineReturningOneMessageFewer(_FakeEngine): + async def process_request(self, messages): + anonymized, mapping = await super().process_request(messages) + return anonymized[:-1], mapping + + +@pytest.mark.asyncio +async def test_pre_call_refuses_a_request_the_engine_did_not_fully_return(monkeypatch): + monkeypatch.setattr(sys.modules["privaite.pii.engine"], "PIIEngine", _EngineReturningOneMessageFewer) + gr = _make_guardrail() + data = {"messages": [{"role": "user", "content": "Email Marie Dupont"}], "input": "also reach marie@acme.com"} + + messages = data["messages"] + with pytest.raises(RuntimeError, match="returned 1 messages for 2"): + await gr.async_pre_call_hook(None, None, data, "aresponses") + assert data["messages"] is messages and messages[0]["content"] == "Email Marie Dupont", "nothing is half written" + + @pytest.mark.asyncio async def test_pre_call_unhandled_input_shape_is_passthrough(): # an unexpected input shape (neither str nor list) is left untouched. From 333d99e25bcac87ae167f261659166fac580063a Mon Sep 17 00:00:00 2001 From: crp4222 <121295348+crp4222@users.noreply.github.com> Date: Wed, 21 Oct 2026 13:06:04 +0200 Subject: [PATCH 15/16] fix(guardrails): pin how the privaite hook handles malformed input The typed rewrite changed three paths that used to crash or skip. A Responses item whose type is not a string is now scanned like any unknown item instead of raising on endswith, a request with metadata set to null now gets its restore map instead of a TypeError, and a dict-shaped response is now restored like an object one. Each one has a test now It also put a privaite install missing PIIEngine on a TypeError, where the old import raised the install hint. That path is back on the hint and tested --- .../guardrail_hooks/privaite/engine.py | 8 ++-- .../guardrail_hooks/test_privaite.py | 38 +++++++++++++++++++ 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/engine.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/engine.py index 419077cbc61..3cc873f7b91 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/engine.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/engine.py @@ -65,9 +65,11 @@ def _module(name: str) -> ModuleType: def _factory(module: ModuleType, name: str) -> Callable[..., object]: value: Final = getattr(module, name, None) - if not callable(value): - raise TypeError(f"{module.__name__}.{name} is missing. {INSTALL_HINT}") - return value + if callable(value): + return value + if hasattr(module, name): + raise TypeError(f"{module.__name__}.{name} is not callable. {INSTALL_HINT}") + raise RuntimeError(f"{module.__name__}.{name} is missing. {INSTALL_HINT}") def _narrow(value: object, kind: type[_Narrowed]) -> _Narrowed: diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py index 8e69b9d2268..6fb70e02171 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py @@ -1386,3 +1386,41 @@ async def test_streaming_tool_arguments_stay_valid_json_with_special_chars(monke assert json.loads(tool_args) == {"name": _SPECIALS_ORIGINAL} assert json.loads(fc_args) == {"name": _SPECIALS_ORIGINAL} assert content == f"Hi {_SPECIALS_ORIGINAL}" + + +@pytest.mark.asyncio +async def test_pre_call_scans_a_responses_item_whose_type_is_not_a_string(): + gr = _make_guardrail() + data = {"input": [{"type": 5, "output": "Email Marie Dupont"}]} + + out = await gr.async_pre_call_hook(None, None, data, "aresponses") + + assert out["input"][0]["output"] == "Email " + + +@pytest.mark.asyncio +async def test_pre_call_stashes_the_map_when_metadata_is_null(): + gr = _make_guardrail() + data = {"messages": [{"role": "user", "content": "Email Marie Dupont"}], "metadata": None} + + out = await gr.async_pre_call_hook(None, None, data, "acompletion") + + assert out["metadata"]["privaite_map"] == {"": "Marie Dupont"} + + +@pytest.mark.asyncio +async def test_post_call_restores_a_dict_shaped_response(): + gr = _make_guardrail() + response = {"choices": [{"message": {"content": "Sending to "}}]} + + out = await gr.async_post_call_success_hook({"metadata": {"privaite_map": _FAKES}}, None, response) + + assert out["choices"][0]["message"]["content"] == "Sending to Marie Dupont" + + +@pytest.mark.asyncio +async def test_engine_for_a_privaite_without_pii_engine_recommends_the_install(monkeypatch): + monkeypatch.delattr(sys.modules["privaite.pii.engine"], "PIIEngine") + + with pytest.raises(RuntimeError, match=r"pip install 'privaite>=0\.7\.2'"): + await _make_guardrail()._engine_for(["en"]) From f4729027fd08504a5cba30043ef4be5febe85cfe Mon Sep 17 00:00:00 2001 From: crp4222 <121295348+crp4222@users.noreply.github.com> Date: Wed, 21 Oct 2026 13:22:07 +0200 Subject: [PATCH 16/16] fix(guardrails): restore privaite responses on a copy The post-call hooks wrote the restored originals into the objects LiteLLM handed them. LiteLLM keeps those objects, the response for its cache write and success logs and every stream chunk in the wrapper's chunk list, which it reassembles into the cached and logged response at the end of the stream. A cached reply could then carry one caller's values under a key built from the anonymized request, where another caller's identical-looking request would hit them Both hooks now restore a deep copy and return or yield that, so what LiteLLM keeps holds placeholders only. The proxy adopts a guardrail's returned response, so callers still get the restored values --- .../guardrail_hooks/privaite/privaite.py | 26 +++++++++++--- .../guardrail_hooks/test_privaite.py | 34 +++++++++++++++++++ 2 files changed, 55 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py index f382bd84925..4ec7c3f0748 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py @@ -21,6 +21,7 @@ without it. See https://github.com/crp4222/PrivAiTe. from __future__ import annotations import asyncio +import copy import itertools from collections.abc import AsyncGenerator, AsyncIterable, Iterable, Iterator, Sequence from typing import TYPE_CHECKING, Final, TypeAlias @@ -228,6 +229,13 @@ class PrivaiteGuardrail(CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, response: object, ) -> object: + """Restores the original values in a copy of the response. + + LiteLLM keeps the object it hands this hook for its response cache and its + success logs, and a cache write may serialize it after this hook ran. Restoring it + in place would put this caller's values in the cache under a key built from the + anonymized request, where another caller's identical-looking request hits them. + """ if not self.deanonymize: return response pairs: Final = _pop_restore_map(data) @@ -235,12 +243,13 @@ class PrivaiteGuardrail(CustomGuardrail): return response mapping: Final = restore_mapping(pairs) engine: Final = await self._engine_for(self._languages()) - for choice in read_list(response, "choices"): + restored: Final = copy.deepcopy(response) + for choice in read_list(restored, "choices"): message = read_field(choice, "message") if message is not None: await restore_message(message, engine, mapping) - await restore_responses_output(response, engine, mapping) - return response + await restore_responses_output(restored, engine, mapping) + return restored async def async_post_call_streaming_iterator_hook( self, @@ -248,6 +257,12 @@ class PrivaiteGuardrail(CustomGuardrail): response: AsyncIterable[ModelResponseStream], request_data: JsonObject, ) -> AsyncGenerator[ModelResponseStream, None]: + """Restores the original values in copies of the chunks. + + The stream wrapper keeps every chunk it yields and rebuilds the whole response from + them for the cache and the success logs once the stream ends, so the chunks + themselves keep their placeholders, as in the non-streaming hook. + """ pairs: Final = _pop_restore_map(request_data) if not self.deanonymize or not pairs: async for chunk in response: @@ -260,8 +275,9 @@ class PrivaiteGuardrail(CustomGuardrail): restore_buffer_factory(), ) async for chunk in response: - restorer.restore_chunk(chunk) - yield chunk + restored = copy.deepcopy(chunk) + restorer.restore_chunk(restored) + yield restored for drained in restorer.drain(): yield drained diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py index 6fb70e02171..5a19ce91fe2 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py @@ -1424,3 +1424,37 @@ async def test_engine_for_a_privaite_without_pii_engine_recommends_the_install(m with pytest.raises(RuntimeError, match=r"pip install 'privaite>=0\.7\.2'"): await _make_guardrail()._engine_for(["en"]) + + +@pytest.mark.asyncio +async def test_post_call_restores_a_copy_and_leaves_the_response_litellm_caches_and_logs(): + gr = _make_guardrail() + 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({"metadata": {"privaite_map": _FAKES}}, None, response) + + assert out.choices[0].message.content == "Sending to Marie Dupont" + assert message.content == "Sending to ", "the object LiteLLM caches and logs must keep the placeholder" + + +@pytest.mark.asyncio +async def test_streaming_restores_copies_and_leaves_the_chunks_litellm_reassembles(): + gr = _make_guardrail() + source = [ + _choice_chunk(_bare_delta(content="Hi "], ( + "the chunks LiteLLM keeps to rebuild the logged and cached response must keep the placeholder" + )