Merge f4729027fd into f2118fc767

litellm/proxy/guardrails/guardrail_hooks/privaite/__init__.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Final
+
+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) -> PrivaiteGuardrail:
+    import litellm
+
+    _privaite_callback: Final = PrivaiteGuardrail(
+        guardrail_name=guardrail["guardrail_name"],
+        event_hook=litellm_params.mode,
+        default_on=litellm_params.default_on,
+        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: Final = {
+    SupportedGuardrailIntegrations.PRIVAITE.value: initialize_guardrail,
+}
+
+guardrail_class_registry: Final = {
+    SupportedGuardrailIntegrations.PRIVAITE.value: PrivaiteGuardrail,
+}

litellm/proxy/guardrails/guardrail_hooks/privaite/engine.py

@@ -0,0 +1,159 @@
+"""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 callable(value):
+        return value
+    if hasattr(module, name):
+        raise TypeError(f"{module.__name__}.{name} is not callable. {INSTALL_HINT}")
+    raise RuntimeError(f"{module.__name__}.{name} is missing. {INSTALL_HINT}")
+
+
+def _narrow(value: object, kind: type[_Narrowed]) -> _Narrowed:
+    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

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

litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py

@@ -0,0 +1,293 @@
+"""
+PrivAiTe guardrail for the LiteLLM proxy.
+
+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 a 400 before anything is forwarded. The error names the types, never the
+values.
+
+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 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, Mode
+
+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,
+)
+
+if TYPE_CHECKING:
+    from litellm.caching.caching import DualCache
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.types.utils import ModelResponseStream
+
+RESTORE_MAP_KEY: Final = "privaite_map"
+
+_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, ...]]
+
+
+class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object):
+    pass
+
+
+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 _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)))
+
+
+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)
+
+
+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 _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 _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.
+
+    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,
+        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
+        )
+        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) -> tuple[str, ...]:
+        languages: Final = tuple(language.strip() for language in self.languages.split(",") if language.strip())
+        return languages or ("en",)
+
+    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
+            engine: Final = await build_engine(self.preset, languages, self.deanonymize, self.block_entities)
+            self._engine = engine
+            self._engine_key = key
+            return engine
+
+    async def _anonymize_request(self, data: JsonObject, engine: PrivaiteEngine) -> RestoreMapping | None:
+        """Anonymizes every scanned request surface in place under one shared mapping.
+
+        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([*messages, *(_pseudo_message(text) for text, _ in slots)])
+        count: Final = len(messages)
+        if len(anonymized) != count + len(slots):
+            raise RuntimeError(
+                f"privaite returned {len(anonymized)} messages for {count + len(slots)}, "
+                "refusing to forward a request it did not scrub in full"
+            )
+        if messages:
+            messages[:] = anonymized[:count]
+        for (_, write), leaf_message in zip(slots, anonymized[count:], strict=True):
+            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: 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 prompt_variables(data) is None
+            and not has_aux_fields(data)
+        ):
+            return data
+
+        engine: Final = await self._engine_for(self._languages())
+        blocked: Final = blocked_error()
+        try:
+            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
+
+        # 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 async_post_call_success_hook(
+        self,
+        data: JsonObject,
+        user_api_key_dict: UserAPIKeyAuth,
+        response: object,
+    ) -> object:
+        """Restores the original values in a copy of the response.
+
+        LiteLLM keeps the object it hands this hook for its response cache and its
+        success logs, and a cache write may serialize it after this hook ran. Restoring it
+        in place would put this caller's values in the cache under a key built from the
+        anonymized request, where another caller's identical-looking request hits them.
+        """
+        if not self.deanonymize:
+            return response
+        pairs: Final = _pop_restore_map(data)
+        if not pairs:
+            return response
+        mapping: Final = restore_mapping(pairs)
+        engine: Final = await self._engine_for(self._languages())
+        restored: Final = copy.deepcopy(response)
+        for choice in read_list(restored, "choices"):
+            message = read_field(choice, "message")
+            if message is not None:
+                await restore_message(message, engine, mapping)
+        await restore_responses_output(restored, engine, mapping)
+        return restored
+
+    async def async_post_call_streaming_iterator_hook(
+        self,
+        user_api_key_dict: UserAPIKeyAuth,
+        response: AsyncIterable[ModelResponseStream],
+        request_data: JsonObject,
+    ) -> AsyncGenerator[ModelResponseStream, None]:
+        """Restores the original values in copies of the chunks.
+
+        The stream wrapper keeps every chunk it yields and rebuilds the whole response from
+        them for the cache and the success logs once the stream ends, so the chunks
+        themselves keep their placeholders, as in the non-streaming hook.
+        """
+        pairs: Final = _pop_restore_map(request_data)
+        if not self.deanonymize or not pairs:
+            async for chunk in response:
+                yield chunk
+            return
+
+        restorer: Final = StreamRestorer(
+            restore_mapping(pairs),
+            restore_mapping(_json_escaped(pairs)),
+            restore_buffer_factory(),
+        )
+        async for chunk in response:
+            restored = copy.deepcopy(chunk)
+            restorer.restore_chunk(restored)
+            yield restored
+        for drained in restorer.drain():
+            yield drained
+
+    async def async_post_call_failure_hook(
+        self,
+        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)

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

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.<carrier>.<name>, 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)

litellm/proxy/spend_tracking/spend_tracking_utils.py

@@ -1210,7 +1210,7 @@ def _get_messages_for_spend_logs_payload(
     return "{}"
 
 
-_SENSITIVE_REQUEST_BODY_KEYS: Final = frozenset({"secret_fields"})
+_SENSITIVE_REQUEST_BODY_KEYS: Final = frozenset({"secret_fields", "privaite_map"})
 _REQUEST_BODY_CREDENTIAL_MASKER: Final = SensitiveDataMasker(extra_sensitive_patterns=frozenset({"apikey"}))
 _TOOL_INPUT_BLOCK_TYPES: Final = frozenset({"tool_use", "server_tool_use", "mcp_tool_use"})
 _TOOL_OUTPUT_BLOCK_TYPES: Final = frozenset({"tool_result", "mcp_tool_result", "function_call_output"})

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

litellm/types/proxy/guardrails/guardrail_hooks/privaite.py

@@ -0,0 +1,31 @@
+from pydantic import Field
+
+from .base import GuardrailConfigModel
+
+
+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: str | None = Field(
+        default="en,fr",
+        description="Comma-separated spaCy languages used for detection.",
+    )
+    deanonymize: bool | None = Field(
+        default=True,
+        description="Restore the original PII values in the response.",
+    )
+    block_entities: str | tuple[str, ...] | None = Field(
+        default=None,
+        description=(
+            "PII types to reject outright instead of masking: a list, or a "
+            "comma-separated string, e.g. ['US_SSN', 'CREDIT_CARD']. A request "
+            "containing any listed type is refused with a 400 and nothing is "
+            "forwarded. Omitted (the default) masks everything and forwards it."
+        ),
+    )
+
+    @staticmethod
+    def ui_friendly_name() -> str:
+        return "PrivAiTe"