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