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

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

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

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

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 "

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