scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-30 06:43:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 250

feat(guardrails): add PrivAiTe guardrail provider

Browse source 
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.
This commit is contained in:
crp4222 2026-06-27 23:31:19 +01:00
parent a8113b0dd7
commit 483d57ee75
6 changed files with 1060 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

32
litellm/proxy/guardrails/guardrail_hooks/privaite/__init__.py Normal file
View file

@ -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,
}

362
litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py Normal file
View file

@ -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

2
litellm/proxy/spend_tracking/spend_tracking_utils.py
View file

@ -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"})

5
litellm/types/guardrails.py
View file

@ -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(

24
litellm/types/proxy/guardrails/guardrail_hooks/privaite.py Normal file
View file

@ -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"

636
tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py Normal file
View file

@ -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 = {"<PERSON_1>": "Marie Dupont", "<EMAIL_ADDRESS_1>": "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 "<PERSON_1>" in content and "<EMAIL_ADDRESS_1>" 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 "<PERSON_1>" 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": {"<PERSON_1>": "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": "<EMAIL_ADDRESS_1>"}'))
# a tool_call with no function must be skipped, not crash.
empty_tool_call = types.SimpleNamespace(function=None)
message = types.SimpleNamespace(
content="Sending to <PERSON_1>",
tool_calls=[tool_call, empty_tool_call],
function_call=types.SimpleNamespace(arguments='{"name": "<PERSON_1>"}'),
)
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 <PERSON_1>", 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 <PERSON_1>"
@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 <PERSON_1>"),
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 <PER") # placeholder starts, held in choice-0 buffer
yield _chunk(1, "Bye") # interleaved choice-1 content
yield _chunk(0, "SON_1>", 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": "<EMAIL_ADDRESS'),
),
types.SimpleNamespace(index=1, function=None),
],
function_call=types.SimpleNamespace(arguments='{"n": "<PER'),
),
finish_reason=None,
)
]
)
yield types.SimpleNamespace(
choices=[
types.SimpleNamespace(
index=0,
delta=types.SimpleNamespace(
content=None,
tool_calls=[
types.SimpleNamespace(
index=0,
function=types.SimpleNamespace(arguments='_1>"}'),
)
],
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 "<PERSON_1>" 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 "<PERSON_1>" 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 <PERSON_1>"
@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 "<PERSON_1>" in messages[0]["content"]
assert "marie@acme.com" not in out["input"]
assert "<EMAIL_ADDRESS_1>" 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 <PERSON_1>"}],
}
# output item as an object (function_call with arguments)
obj_item = types.SimpleNamespace(
type="function_call",
content=None,
arguments='{"to": "<EMAIL_ADDRESS_1>"}',
)
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