From 464516a18878bed8cf4a7ccfebb43cdbbd9dbb1b Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Mon, 5 Oct 2026 09:58:48 +0000 Subject: [PATCH] refactor(guardrails): split llm_shield_proxy into payload, request walk and stream modules The module had grown past 1,700 lines. Shared payload types and helpers move to payload.py, the request walk to request_walk.py and the stream restorers to stream_restorers.py; llm_shield_proxy.py keeps the guardrail class. No behaviour change. --- .../llm_shield_proxy/llm_shield_proxy.py | 1070 ++--------------- .../llm_shield_proxy/payload.py | 206 ++++ .../llm_shield_proxy/request_walk.py | 402 +++++++ .../llm_shield_proxy/stream_restorers.py | 368 ++++++ 4 files changed, 1063 insertions(+), 983 deletions(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/payload.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/request_walk.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py index dc0dbfda242..0291dba2ac3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py @@ -7,14 +7,9 @@ import copy import functools -import itertools -import json import os -import re import uuid -from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence -from enum import Enum -from types import MappingProxyType +from collections.abc import AsyncGenerator, Callable, Mapping, Sequence from typing import ( TYPE_CHECKING, Any, # noqa: TID251 # **kwargs forwards verbatim to CustomGuardrail.__init__ @@ -22,7 +17,6 @@ from typing import ( Final, Literal, Optional, - TypeAlias, ) import httpx @@ -48,6 +42,35 @@ if TYPE_CHECKING: LLMShieldProxyGuardrailConfigModel, ) from litellm.types.utils import CallTypes, LLMResponseTypes +from .payload import ( + JsonBody, + MutableRequest, + RequestTooDeep, + Slot, + SlotSink, + as_array, + as_object, + choice_index, + collect_json_leaves, + collect_response_item, + detached, + read_field, + read_list, + rehydrate_slots, + write_field, +) +from .request_walk import ( + locate_request_texts, +) +from .stream_restorers import ( + AnthropicSSERestorer, + CarryKey, + CarryWindows, + ResponsesStreamRestorer, + carry_sort_key, + continuation_delta, + responses_event_type, +) GUARDRAIL_NAME: Final = "llm_shield_proxy" @@ -69,10 +92,6 @@ _SESSION_METADATA_KEY: Final = "llm_shield_session_id" # a proxy request into deployment-level restoration by sending the key themselves. _DEPLOYMENT_RESTORE_KEY: Final = "llm_shield_restore_at_deployment" -# Roles whose text the application author wrote and the caller never sees. Their -# PII is still redacted outbound, but it is not restorable from the reply. -_PRIVILEGED_ROLES: Final = frozenset({"system", "developer"}) - # Vault ids are minted here and never derived from anything the caller sends. The # vault holds the plaintext behind every placeholder, so an id a caller could # supply or guess would let one user rehydrate another user's values by getting a @@ -82,877 +101,6 @@ _VAULT_PREFIX: Final = f"litellm-{uuid.uuid4().hex}" _DEFAULT_TIMEOUT_SECONDS: Final = 10.0 -# The proxy's own request dict. Mutable by design: a pre-call guardrail rewrites -# the caller's payload in place, which is the entire point of the hook. -MutableRequest: TypeAlias = dict[str, object] - -# A JSON body on its way to httpx, which requires a real dict rather than a view. -JsonBody: TypeAlias = dict[str, object] - -# One redactable span: the text as it stands, and the write that puts the -# replacement back where it came from. -# How far a tool_result chain is followed. Real payloads nest one or two deep; the -# bound is what stops a crafted one from becoming an unbounded walk. A request that -# nests deeper is refused rather than forwarded, because text past the bound would -# otherwise reach the provider unredacted. -_MAX_CONTENT_DEPTH: Final = 8 - -# How far a JSON value -- a tool input, a parameter schema -- is followed on the request -# side. Legitimate JSON nests far deeper than content blocks do, so the bound is -# generous; past it the request is refused, for the same reason as above. -_MAX_JSON_DEPTH: Final = 64 - -_Slot: TypeAlias = tuple[str, Callable[[str], None]] - -# One incremental rehydration step for a stream the caller has already bound to its -# vault: (new text, carried window, final) -> (text safe to emit, window still held). -_StreamStep: TypeAlias = Callable[[str, str, bool], Awaitable[tuple[str, str]]] - -# A batch rehydration already bound to the request's vault. -_Rehydrate: TypeAlias = Callable[[Sequence[str]], Awaitable[Sequence[str]]] - -# Anthropic /v1/messages delta types that carry restorable text, and the field holding -# it. `thinking_delta` is left out on purpose: a thinking block is signed, and one -# rewritten here fails verification when the client sends it back on the next turn. -_ANTHROPIC_DELTA_FIELDS: Final = MappingProxyType({"text_delta": "text", "input_json_delta": "partial_json"}) - -# A blank line ends an SSE event. Frames are cut there, never inside an event, so a -# `data:` line split across two network chunks is parsed only once it is whole. -_SSE_EVENT_BOUNDARY: Final = re.compile(rb"(\r?\n\r?\n)") - -# What an SSE stream can open with: one of its fields, or a `:` comment. -_SSE_OPENINGS: Final = (b"event:", b"data:", b"id:", b"retry:", b":") - -# Responses API delta events whose `delta` is not text. Audio arrives base64-encoded; -# sending it through the shield would cost a round trip per chunk to restore nothing. -_RESPONSES_BINARY_DELTAS: Final = frozenset(("response.audio.delta",)) - -# Fields on a Responses API event that identify something rather than say something. -# Every other string field on a `.done` event is model text and is restored, so an event -# type added upstream is covered by default instead of leaking a placeholder. -_RESPONSES_STRUCTURAL_FIELDS: Final = frozenset( - ("type", "id", "item_id", "call_id", "name", "server_label", "status", "obfuscation") -) - -# Terminal Responses API events that repeat the whole reply under `response`. -_RESPONSES_TERMINAL_EVENTS: Final = frozenset(("response.completed", "response.incomplete")) - -# JSON Schema keywords whose value has to reach the model or a validator verbatim, so the -# schema walk leaves them alone: types, formats, patterns, references and -# required-property lists. Everything else is scanned. -_SCHEMA_STRUCTURAL_KEYWORDS: Final = frozenset( - ( - "type", - "format", - "pattern", - "required", - "dependentRequired", - "propertyOrdering", - "discriminator", - "contentEncoding", - "contentMediaType", - "$ref", - "$id", - "$schema", - "$anchor", - "$dynamicRef", - "$dynamicAnchor", - "$recursiveRef", - "$recursiveAnchor", - "$vocabulary", - ) -) - -# Keywords holding JSON values rather than schemas: every string in them is collected, -# whatever the keys around it are called. -_SCHEMA_VALUE_KEYWORDS: Final = frozenset(("examples", "default")) - -# Keywords holding the literal values the model must reproduce. These go to the CALLER's -# vault, not the privileged one: the model emits the stand-in in its tool arguments or -# structured output, and restoring the reply turns it back into the value the schema -# allows, so the call still routes. In the non-restorable vault it would come back as a -# stand-in no validator accepts. -_SCHEMA_LITERAL_KEYWORDS: Final = frozenset(("enum", "const")) - -# Keywords whose value maps names to subschemas. Their keys are property names, not -# keywords, so a property called `type` or `enum` is walked like any other subschema. -_SCHEMA_MAP_KEYWORDS: Final = frozenset( - ("properties", "patternProperties", "$defs", "definitions", "dependentSchemas", "dependencies") -) - -# The accumulator the collectors below append into. It never escapes -# _locate_request_texts, which freezes it into a tuple before returning. -_SlotSink: TypeAlias = list[_Slot] - -# Sliding windows keyed by (choice index, tool-call index | None), threaded through one -# stream. `None` is the content channel; an int is one tool call's accumulating -# `arguments`. Content and each tool call are separate token streams, so each needs its -# own window -- one shared window would splice one stream's held-back tail onto another. -_CarryKey: TypeAlias = tuple[int, int | None] -_CarryWindows: TypeAlias = dict[_CarryKey, str] - -# A caller-owned list whose entries are rewritten in place, such as a Completions -# `prompt` sent as an array of strings. -MutableSeq: TypeAlias = list[object] - -# A Responses API delta stream: (event family, item id, output index, part index). -_ResponsesStreamKey: TypeAlias = tuple[str, object, object, object] - - -def _as_object(value: object) -> MutableRequest | None: - """`value` as a JSON object, or None. - - `isinstance(value, dict)` alone leaves the keys and values unknown to the type - checker. A JSON object's keys are strings, so the type is stated once, here. - """ - return value if isinstance(value, dict) else None - - -def _as_array(value: object) -> MutableSeq | None: - """`value` as a JSON array, or None. See `_as_object`.""" - return value if isinstance(value, list) else None - - -def _detached(value: object) -> object: - """A deep copy of a reply or chunk, for restoring without touching LiteLLM's own object. - - LiteLLM keeps the object it handed the hooks to fill its response cache and its - logs, so writing restored plaintext into that object would put it there too. - """ - return copy.deepcopy(value) - - -def _is_container(value: object) -> bool: - """Whether `value` is a JSON object or array, without narrowing it to unknown types.""" - return isinstance(value, (dict, list)) - - -def _collect(container: MutableRequest, key: str, slots: _SlotSink) -> None: - """Records the string at `key`, along with the write that replaces it.""" - value: Final = container.get(key) - if isinstance(value, str) and value: - slots.append((value, functools.partial(container.__setitem__, key))) - - -def _collect_entry(entries: MutableSeq, index: int, slots: _SlotSink) -> None: - """Records a string held directly in a list, rather than under a key.""" - value: Final = entries[index] - if isinstance(value, str) and value: - slots.append((value, functools.partial(entries.__setitem__, index))) - - -def _collect_prompt(data: MutableRequest, slots: _SlotSink) -> None: - """The Completions API sends its text in `prompt`, and its tail in `suffix`.""" - _collect(data, "suffix", slots) - prompt: Final = data.get("prompt") - if isinstance(prompt, str): - _collect(data, "prompt", slots) - return - prompt_object: Final = _as_object(prompt) - if prompt_object is not None: - # A Responses API PromptObject. `variables` are substituted into the stored - # prompt on the provider side, so they are caller text. `id` and `version` - # identify which prompt to use and must arrive unchanged. A variable is a string - # or a typed input such as `{"type": "input_text", "text": ...}`. - variables: Final = _as_object(prompt_object.get("variables")) - if variables is not None: - for name in tuple(variables): - _collect(variables, name, slots) - typed = _as_object(variables[name]) - if typed is not None: - _collect(typed, "text", slots) - return - entries: Final = _as_array(prompt) - if entries is None: - return - for index in range(len(entries)): - _collect_entry(entries, index, slots) - - -class _RequestTooDeep(Exception): - """A request nests text past a walk's bound. - - Skipping the rest would forward it unredacted while the guardrail reports as - enabled, so the pre-call hook refuses the request instead. - """ - - -def _collect_content(container: MutableRequest, slots: _SlotSink) -> None: - """Collects `content`, a string or a list of typed parts. - - An Anthropic tool_result nests its own content, so this has to descend. It walks - with an explicit stack and a depth bound rather than by recursion: the nesting is - caller controlled, and an unbounded descent is a JSON bomb. Content nested past the - bound raises `_RequestTooDeep` rather than being skipped. - """ - # Walked in document order: the shield maps its replies back by position, so the - # order spans are collected in is part of the contract. - pending: Final[list[tuple[MutableRequest, int]]] = [(container, 0)] # mutable-ok: local queue, never escapes. - cursor = 0 # rebind-ok: advances through the queue. - while cursor < len(pending): - node, depth = pending[cursor] - cursor += 1 - content = node.get("content") - if isinstance(content, str): - _collect(node, "content", slots) - continue - if depth >= _MAX_CONTENT_DEPTH and content: - raise _RequestTooDeep("content") - for item in _as_array(content) or (): - part = _as_object(item) - if part is None: - continue - # Image and audio parts have no text and fall through untouched. - _collect(part, "text", slots) - if part.get("type") == "tool_use": - # A replayed Anthropic tool call. Its `input` is a JSON object rather than - # a string, so a value can sit at any depth -- the reply side walks the - # same leaves when it restores one. Its own JSON bound applies, not the - # content one, and past it the request is refused. - _collect_json_leaves(part.get("input"), slots, strict=True) - source = _as_object(part.get("source")) if part.get("type") == "document" else None - if source is not None: - # An Anthropic document carries text inline: a `text` source holds it in - # `data`, a `content` source as a string or blocks, walked like any other - # content. Base64, URL and file sources are binary or remote, and pass - # untouched. Its `title` and `context` are caller text too. - _collect(part, "title", slots) - _collect(part, "context", slots) - if source.get("type") == "text": - _collect(source, "data", slots) - elif source.get("type") == "content": - pending.append((source, depth + 1)) - if "content" in part: - pending.append((part, depth + 1)) - - -def _collect_participant_name(message: MutableRequest, slots: _SlotSink) -> None: - """Redacts `name` where it identifies a person, never where it names a function. - - On a user or assistant turn `name` is the participant, which is personal data. - On a tool or function turn the same field carries the function's name and has - to reach the provider unchanged, or the call no longer routes. - """ - if message.get("role") in ("tool", "function"): - return - _collect(message, "name", slots) - - -def _collect_tool_arguments(message: MutableRequest, slots: _SlotSink) -> None: - """Tool arguments carry the values a user asked the model to act on.""" - for tool_call in _read_list(message, "tool_calls"): - tool_call_object = _as_object(tool_call) - function = _as_object(tool_call_object.get("function")) if tool_call_object is not None else None - if function is not None: - _collect(function, "arguments", slots) - legacy: Final = _as_object(message.get("function_call")) - if legacy is not None: - _collect(legacy, "arguments", slots) - - -def _collect_system(data: MutableRequest, slots: _SlotSink) -> None: - """Anthropic's /v1/messages carries its system prompt at the top level.""" - system: Final = data.get("system") - if isinstance(system, str): - _collect(data, "system", slots) - return - _collect_text_parts(data, "system", slots) - - -def _collect_responses_fields(data: MutableRequest, slots: _SlotSink, privileged: _SlotSink) -> None: - """The Responses API sends text outside `messages`, in `instructions` and `input`. - - `instructions` is written by the application, not by the caller, so it is - collected into the privileged sink; `input` is the caller's own text, except for - system and developer items in it, which go to the privileged sink like their Chat - counterparts. - """ - _collect(data, "instructions", privileged) - request_input: Final = data.get("input") - if isinstance(request_input, str): - _collect(data, "input", slots) - return - entries: Final = _as_array(request_input) - if entries is None: - return - for index, entry in enumerate(entries): - if isinstance(entry, str): - # The embeddings and moderations shape: `input` as an array of strings. - _collect_entry(entries, index, slots) - continue - item = _as_object(entry) - if item is None: - continue - _collect_content(item, privileged if item.get("role") in _PRIVILEGED_ROLES else slots) - # A function_call item holds `arguments`; a function_call_output holds `output`, - # as a string or as a list of input_text parts. A custom_tool_call holds `input` - # and a code_interpreter_call `code` -- the fields the reply side restores. - _collect(item, "arguments", slots) - _collect(item, "output", slots) - _collect_text_parts(item, "output", slots) - _collect(item, "input", slots) - _collect(item, "code", slots) - # A replayed reasoning item carries the model's summary of its own reasoning, - # which quotes whatever the conversation contained. - _collect_text_parts(item, "summary", slots) - - -def _collect_text_parts(container: MutableRequest, key: str, slots: _SlotSink) -> None: - """Collects the `text` of every part in the list held at `key`.""" - for entry in _as_array(container.get(key)) or (): - part = _as_object(entry) - if part is not None: - _collect(part, "text", slots) - - -def _collect_tool_definitions(data: MutableRequest, slots: _SlotSink, privileged: _SlotSink) -> None: - """Tool definitions are application-authored free text bound for the provider. - - A tool's description and the free text in its parameter schema are where callers put - examples and customer context, so they carry PII as often as a prompt does. They are - collected into the privileged sink, like a system prompt: redacted outbound, and never - restorable from the reply. `enum` and `const` values are the exception, and go to the - caller's vault -- see `_SCHEMA_LITERAL_KEYWORDS`. Names and types are left as sent. - - Covers Chat `tools[].function`, the legacy `functions[]`, and the flat tool shape the - Responses API and Anthropic share, whose schema is `parameters` or `input_schema`. - """ - for key in ("tools", "functions"): - for entry in _as_array(data.get(key)) or (): - tool = _as_object(entry) - if tool is None: - continue - function = _as_object(tool.get("function")) - for holder in (tool, function) if function is not None else (tool,): - _collect(holder, "description", privileged) - _collect_schema_text(holder.get("parameters"), slots, privileged) - _collect_schema_text(holder.get("input_schema"), slots, privileged) - - -def _collect_schema_text(schema: object, slots: _SlotSink, privileged: _SlotSink) -> None: - """Collects the text in a JSON Schema, at any depth. - - Scan by default: every string is collected except under the keywords in - `_SCHEMA_STRUCTURAL_KEYWORDS`, whose values must go out verbatim. A list of keywords - *to* collect would leak every one it forgot -- draft-07 `dependencies`, a `$comment`, - a vendor `x-` extension -- which is how this walk started out. Free text goes to the - privileged sink; `enum` / `const` literals go to the caller's, so the model's use of - them is restored. - - Structure matters in two places. Under `properties` and the other name -> subschema - maps, keys are property names rather than keywords, so a property called `type` is a - subschema to walk, not a keyword to skip. And `examples` / `default` hold JSON values, - so all their strings are collected whatever the keys around them are called. Nested - past `_MAX_JSON_DEPTH`, the request is refused. - """ - pending: Final[list[tuple[object, int]]] = [(schema, 0)] # mutable-ok: local walk stack. - while pending: - node, depth = pending.pop() - if depth > _MAX_JSON_DEPTH: - if _is_container(node) and node: - raise _RequestTooDeep("schema") - continue - entries = _as_array(node) - if entries is not None: - for index, item in enumerate(entries): - _collect_entry(entries, index, privileged) - if _is_container(item): - pending.append((item, depth + 1)) - continue - schema_object = _as_object(node) - if schema_object is None: - continue - for keyword, value in tuple(schema_object.items()): - if keyword in _SCHEMA_STRUCTURAL_KEYWORDS: - continue - subschemas = _as_object(value) if keyword in _SCHEMA_MAP_KEYWORDS else None - if keyword in _SCHEMA_LITERAL_KEYWORDS: - _collect(schema_object, keyword, slots) - _collect_json_leaves(value, slots, strict=True) - elif keyword in _SCHEMA_VALUE_KEYWORDS: - _collect(schema_object, keyword, privileged) - _collect_json_leaves(value, privileged, strict=True) - elif subschemas is not None: - pending.extend((child, depth + 1) for child in subschemas.values()) - elif isinstance(value, str): - _collect(schema_object, keyword, privileged) - elif _is_container(value): - pending.append((value, depth + 1)) - - -def _collect_output_contracts(data: MutableRequest, slots: _SlotSink, privileged: _SlotSink) -> None: - """Text the caller sends to shape the reply rather than to prompt it. - - A predicted output (`prediction.content`) is the caller's own draft of the answer, so - it goes with their text: the model largely repeats it, and it has to come back. A - structured-output schema -- Chat `response_format.json_schema`, Responses - `text.format` -- is application-authored like a tool schema, so its free text goes - to the privileged sink, and its names and types stay as sent. - """ - prediction: Final = _as_object(data.get("prediction")) - if prediction is not None: - _collect(prediction, "content", slots) - _collect_text_parts(prediction, "content", slots) - response_format: Final = _as_object(data.get("response_format")) - text_options: Final = _as_object(data.get("text")) - for declared in ( - response_format.get("json_schema") if response_format is not None else None, - text_options.get("format") if text_options is not None else None, - ): - wrapper = _as_object(declared) - if wrapper is not None: - _collect(wrapper, "description", privileged) - _collect_schema_text(wrapper.get("schema"), slots, privileged) - - -def _collect_user_locations(data: MutableRequest, privileged: _SlotSink) -> None: - """Web search forwards the user's approximate location, whose `city` and `region` - are free text and can hold a street address. - - Chat carries it in `web_search_options.user_location.approximate`; the Responses - and Anthropic web-search tools carry it flat on the tool's `user_location`. Nothing - restores it from a reply, hence the privileged sink. - """ - options: Final = data.get("web_search_options") - tools: Final = _as_array(data.get("tools")) or () - for declared in (options, *tools): - holder = _as_object(declared) - location = _as_object(holder.get("user_location")) if holder is not None else None - if location is None: - continue - approximate = _as_object(location.get("approximate")) - for container in (location, approximate) if approximate is not None else (location,): - _collect(container, "city", privileged) - _collect(container, "region", privileged) - - -def _collect_end_user_ids(data: MutableRequest, privileged: _SlotSink) -> None: - """`user` and `safety_identifier` are forwarded to the provider and often hold an email. - - Only detected PII is replaced, so an opaque id reaches the provider unchanged. LiteLLM's - own end-user spend tracking reads the id resolved at authentication, before this hook - runs, so rewriting the field here does not move spend. Nothing restores these from a - reply, hence the privileged sink. - """ - _collect(data, "user", privileged) - _collect(data, "safety_identifier", privileged) - - -def _choice_index(choice: object) -> int: - """Streaming choices are matched across chunks by their index.""" - index: Final = getattr(choice, "index", 0) - return index if isinstance(index, int) else 0 - - -def _read_field(holder: object, name: str) -> object: - """Reads one field from a dict or from an object. - - LiteLLM's replies arrive as Pydantic models on some paths and as plain dicts on - others, depending how far they have been deserialised, so every response walk here - has to handle 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. - - The entries are the reply's own objects, so writing through them edits the reply. - """ - value: Final = _read_field(holder, name) - if isinstance(value, tuple): - return value - return tuple(_as_array(value) or ()) - - -def _write_field(holder: object, name: str, value: object) -> None: - """Writes one string field back into a dict or an object. Pairs with _read_field.""" - if isinstance(holder, dict): - holder[name] = value - else: - setattr(holder, name, value) - - -def _collect_json_leaves(node: object, slots: _SlotSink, *, strict: bool = False) -> None: - """Collects every string leaf of a JSON-ish structure, with a write-back per leaf. - - An Anthropic `tool_use` block carries `input`, an arbitrary JSON object rather than a - string, so a value worth restoring can sit at any depth. Bounded by `_MAX_JSON_DEPTH`: - the shape is caller or model controlled, and the bound is what stops a crafted one from - becoming an unbounded descent. Walked with an explicit stack rather than recursively, - so a deeply nested value cannot spend stack frames proportional to attacker-chosen - depth. - - `strict` is for the request side, where a leaf left behind would reach the provider - unredacted: past the bound it raises `_RequestTooDeep`. On the reply side a leaf past - the bound just keeps its placeholder, which leaks nothing, so it is skipped. - """ - pending: Final[list[tuple[object, int]]] = [(node, 0)] # mutable-ok: local walk stack. - while pending: - current, current_depth = pending.pop() - if current_depth > _MAX_JSON_DEPTH: - if strict and _is_container(current) and current: - raise _RequestTooDeep("json") - continue - current_object = _as_object(current) - if current_object is not None: - for key in tuple(current_object): - value = current_object[key] - if isinstance(value, str) and value: - slots.append((value, functools.partial(current_object.__setitem__, key))) - else: - pending.append((value, current_depth + 1)) - continue - entries = _as_array(current) - if entries is not None: - for index, value in enumerate(entries): - if isinstance(value, str) and value: - slots.append((value, functools.partial(entries.__setitem__, index))) - else: - pending.append((value, current_depth + 1)) - - -def _carry_sort_key(key: _CarryKey) -> tuple[int, int]: - """Orders streaming windows without ever comparing None to an int. - - `sorted()` over the raw keys raises as soon as one choice holds both a content window - and a tool-call window, because `None < 0` is not orderable. Content sorts first, then - tool calls by their index. - """ - choice_index, tool_index = key - return (choice_index, -1 if tool_index is None else tool_index) - - -def _continuation_delta(tool_index: int, text: str) -> list[dict[str, object]]: - """A `tool_calls` delta carrying `text` as an index-only continuation. - - Clients concatenate tool-call fragments by index, so no id or name is needed. - """ - return [{"index": tool_index, "function": {"arguments": text}}] - - -def _collect_response_item(item: object, slots: _SlotSink) -> None: - """Restorable spans in one Responses API output item, dict or object. - - Mirrors `_collect_responses_fields` on the request side -- a function_call or - mcp_call item holds `arguments`, their outputs `output`, a reasoning item `summary` - parts -- so the two directions stay symmetric. A custom tool call carries `input` and - a code interpreter call `code`, both model-written. - """ - for block in _read_list(item, "content"): - for field in ("text", "refusal"): - text = _read_field(block, field) - if isinstance(text, str) and text: - slots.append((text, lambda new, b=block, f=field: _write_field(b, f, new))) - for part in _read_list(item, "summary"): - text = _read_field(part, "text") - if isinstance(text, str) and text: - slots.append((text, lambda new, p=part: _write_field(p, "text", new))) - for field in ("arguments", "output", "input", "code"): - value = _read_field(item, field) - if isinstance(value, str) and value: - slots.append((value, lambda new, i=item, f=field: _write_field(i, f, new))) - - -async def _rehydrate_slots(slots: Sequence[_Slot], rehydrate: _Rehydrate) -> None: - """Restores every span in `slots` in one batch and writes each result back.""" - if not slots: - return - restored: Final = await rehydrate(tuple(text for text, _ in slots)) - for (_, write), replacement in zip(slots, restored): - write(replacement) - - -def _opens_like_sse(head: bytes) -> bool | None: - """Whether a raw stream is SSE, judged by its opening bytes; None while undecidable. - - An SSE stream opens with a field name or a `:` comment. Anything else -- a JSON array - streamed in pieces, say -- has no event boundaries to wait for. A chunk that ends - partway through a field name decides nothing yet, so that case waits for more. - """ - opening: Final = head.lstrip() - if not opening: - return None - if opening.startswith(_SSE_OPENINGS): - return True - if any(field.startswith(opening) for field in _SSE_OPENINGS): - return None - return False - - -def _responses_event_type(chunk: object) -> str | None: - """The event type of a Responses API stream event, or None for any other chunk. - - The type arrives as a plain string on dicts and as a str-valued Enum on LiteLLM's - event models. The Enum is unwrapped because it does not hash like its value, so it - would miss every lookup in the event tables above. - """ - if isinstance(chunk, (bytes, str)): - return None - kind: Final = _read_field(chunk, "type") - value: Final = kind.value if isinstance(kind, Enum) else kind - return value if isinstance(value, str) and value.startswith("response.") else None - - -class _AnthropicSSERestorer: - """Restores an Anthropic `/v1/messages` stream, which reaches the hook as raw SSE. - - Each content block is its own token stream with its own window, keyed by the block's - `index`: `text_delta` carries prose and `input_json_delta` a tool call's arguments. - When a block stops, whatever its window still holds is emitted as one more delta for - that block, just ahead of the `content_block_stop` frame, so the client has the whole - block before it is told the block is complete. - - Frames are processed whole. A network chunk can end in the middle of an event, so the - unfinished tail is kept until the rest arrives; that delays one partial event, never - a completed one. A frame that is not an Anthropic event -- another endpoint's SSE, or - anything that fails to parse -- is passed through byte for byte, and a raw stream that - does not open like SSE at all is passed through chunk by chunk, never buffered. - """ - - def __init__(self, step: _StreamStep) -> None: - self._step: Final = step - self._carries: Final[dict[int, str]] = {} # mutable-ok: per-block windows advanced in place. - self._delta_types: Final[dict[int, str]] = {} # mutable-ok: each block's delta type, for its flush. - self._pending = b"" - self._as_text = False - self._is_sse: bool | None = None - - async def feed(self, chunk: bytes | str) -> tuple[bytes | str, ...]: - """Restores every event this chunk completes; holds back an unfinished tail.""" - if isinstance(chunk, str): - self._as_text = True - if self._is_sse is False: - return (chunk,) - raw: Final = chunk.encode("utf-8") if isinstance(chunk, str) else chunk - buffered: Final = self._pending + raw - if self._is_sse is None: - self._is_sse = _opens_like_sse(buffered) - if self._is_sse is None: - # Too little has arrived to tell -- `b"eve"` could still become `event:`. - self._pending = buffered - return () - if not self._is_sse: - self._pending = b"" - return self._emit(buffered) - boundaries: Final = tuple(_SSE_EVENT_BOUNDARY.finditer(buffered)) - if not boundaries: - self._pending = buffered - return () - cut: Final = boundaries[-1].end() - self._pending = buffered[cut:] - # With a capturing group, split alternates event, separator, ..., and ends in the - # empty remainder after the last separator. - parts: Final = _SSE_EVENT_BOUNDARY.split(buffered[:cut]) - restored: Final = tuple( - [await self._restore_event(parts[index]) + parts[index + 1] for index in range(0, len(parts) - 1, 2)] - ) - return self._emit(b"".join(restored)) - - async def finish(self) -> tuple[bytes | str, ...]: - """Emits an unterminated final event and any window a block never closed.""" - held: Final = self._pending - self._pending = b"" - if not self._is_sse: - # The stream ended before it could be told apart from SSE: hand it back as is. - return self._emit(held) - tail: Final = await self._restore_event(held) if held.strip() else held - flushed: Final = await self._flush_all() - # The tail had no blank line after it; one is needed before another frame follows. - separator: Final = b"\n\n" if tail.strip() and flushed else b"" - return self._emit(tail + separator + flushed) - - def _emit(self, frames: bytes) -> tuple[bytes | str, ...]: - if not frames: - return () - return (frames.decode("utf-8") if self._as_text else frames,) - - async def _restore_event(self, block: bytes) -> bytes: - """Rewrites one SSE event, or returns it untouched if it carries nothing to restore.""" - try: - lines: Final = block.decode("utf-8").split("\n") - except UnicodeDecodeError: - return block - data_lines: Final = tuple(index for index, line in enumerate(lines) if line.startswith("data:")) - if len(data_lines) != 1: - return block - line: Final = lines[data_lines[0]] - try: - parsed: Final[object] = json.loads(line[len("data:") :]) - except ValueError: - return block - event: Final = _as_object(parsed) - if event is None: - return block - kind: Final = event.get("type") - index: Final = event.get("index") - if kind == "content_block_stop" and isinstance(index, int): - return await self._flush(index) + block - if kind == "message_stop": - return await self._flush_all() + block - if kind != "content_block_delta" or not await self._restore_delta(event): - return block - ending: Final = "\r" if line.endswith("\r") else "" - rewritten: Final = ( - *lines[: data_lines[0]], - f"data: {json.dumps(event, ensure_ascii=False)}{ending}", - *lines[data_lines[0] + 1 :], - ) - return "\n".join(rewritten).encode("utf-8") - - async def _restore_delta(self, event: MutableRequest) -> bool: - """Advances one block's window through this delta. False if it holds no text.""" - index: Final = event.get("index") - delta: Final = event.get("delta") - if not isinstance(index, int) or not isinstance(delta, dict): - return False - delta_type: Final = delta.get("type") - if not isinstance(delta_type, str): - return False - field: Final = _ANTHROPIC_DELTA_FIELDS.get(delta_type) - text: Final = delta.get(field) if field is not None else None - if field is None or not isinstance(text, str) or not text: - return False - emitted, remaining = await self._step(text, self._carries.get(index, ""), False) - self._carries[index] = remaining - self._delta_types[index] = delta_type - delta[field] = emitted - return True - - async def _flush(self, index: int) -> bytes: - """One synthetic delta frame carrying whatever `index`'s window still holds.""" - carry: Final = self._carries.pop(index, "") - delta_type: Final = self._delta_types.pop(index, None) - field: Final = _ANTHROPIC_DELTA_FIELDS.get(delta_type) if isinstance(delta_type, str) else None - if not carry or field is None: - return b"" - text, _ = await self._step("", carry, True) - if not text: - return b"" - event: Final[JsonBody] = { - "type": "content_block_delta", - "index": index, - "delta": {"type": delta_type, field: text}, - } - return f"event: content_block_delta\ndata: {json.dumps(event, ensure_ascii=False)}\n\n".encode() - - async def _flush_all(self) -> bytes: - flushed: Final = tuple([await self._flush(index) for index in tuple(self._carries)]) - return b"".join(flushed) - - -class _ResponsesStreamRestorer: - """Restores a Responses API event stream. - - The event families are matched by shape rather than listed, so a text stream the - API adds later is restored by default instead of leaking a placeholder: - - - Any `*.delta` event whose `delta` is a string is a token stream (output_text, - refusal, function-call and MCP arguments, reasoning summaries, ...). Each gets its - own window, keyed by the family, the item id and the part index. - - Any `*.done` event closes the stream of the same family. Whatever its window still - holds goes out first, as a copy of that stream's last delta event -- so it carries - the stream's own ids, and repeats that event's `sequence_number`. Then every text - field on the done event is restored in full: its string fields other than - identifiers, plus any `part` or `item` it repeats. - - `response.completed` / `response.incomplete` repeat the whole reply, and are - restored the same way the non-streaming reply is. - """ - - def __init__(self, step: _StreamStep, rehydrate: _Rehydrate) -> None: - self._step: Final = step - self._rehydrate: Final = rehydrate - self._carries: Final[dict[_ResponsesStreamKey, str]] = {} # mutable-ok: per-stream windows advanced in place. - self._last_deltas: Final[dict[_ResponsesStreamKey, object]] = {} # mutable-ok: newest delta per stream. - - async def restore(self, event: object) -> tuple[object, ...]: - """The events to emit in place of `event`: any flush, then the event itself.""" - kind: Final = _responses_event_type(event) - if kind is None: - return (event,) - if kind.endswith(".delta") and kind not in _RESPONSES_BINARY_DELTAS: - await self._restore_delta(event, kind) - return (event,) - slots: Final[_SlotSink] = [] - flushed: Final = await self._flush(_responses_stream_key(event, kind)) if kind.endswith(".done") else () - if kind.endswith(".done"): - _collect_event_text(event, slots) - part: Final = _read_field(event, "part") - if part is not None: - _collect_response_item({"content": [part]}, slots) - _collect_response_item(_read_field(event, "item"), slots) - elif kind in _RESPONSES_TERMINAL_EVENTS: - for item in _read_list(_read_field(event, "response"), "output"): - _collect_response_item(item, slots) - await _rehydrate_slots(slots, self._rehydrate) - return (*flushed, event) - - async def finish(self) -> tuple[object, ...]: - """Flushes every stream the provider never closed, e.g. a truncated reply.""" - flushed: Final = tuple([await self._flush(key) for key in tuple(self._carries)]) - return tuple(itertools.chain.from_iterable(flushed)) - - async def _restore_delta(self, event: object, kind: str) -> None: - text: Final = _read_field(event, "delta") - if not isinstance(text, str) or not text: - return - key: Final = _responses_stream_key(event, kind) - emitted, remaining = await self._step(text, self._carries.get(key, ""), False) - self._carries[key] = remaining - self._last_deltas[key] = event - _write_field(event, "delta", emitted) - - async def _flush(self, key: _ResponsesStreamKey) -> tuple[object, ...]: - carry: Final = self._carries.pop(key, "") - template: Final = self._last_deltas.pop(key, None) - if not carry or template is None: - return () - text, _ = await self._step("", carry, True) - if not text: - return () - flush: Final = copy.deepcopy(template) - _write_field(flush, "delta", text) - return (flush,) - - -def _responses_stream_key(event: object, kind: str) -> _ResponsesStreamKey: - """Identifies the delta stream an event belongs to, the same for its delta and done. - - The family is the event type without its `.delta` / `.done` suffix, so an output_text - stream and a refusal stream on the same part never share a window. - """ - family: Final = kind.rsplit(".", 1)[0] - part_index: Final = _read_field(event, "content_index") - summary_index: Final = _read_field(event, "summary_index") - return ( - family, - _read_field(event, "item_id"), - _read_field(event, "output_index"), - part_index if part_index is not None else summary_index, - ) - - -def _collect_event_text(event: object, slots: _SlotSink) -> None: - """Collects every top-level text field of a Responses API event, dict or model. - - Scan by default, with identifiers excluded, rather than a list of known fields: the - `.done` event of each stream family names its text differently (`text`, `refusal`, - `arguments`, ...), and a family added upstream would otherwise leak a placeholder. - """ - # A model's fields live in its `__dict__`; an empty dict has none either way. - attributes: Final[object] = getattr(event, "__dict__", None) - fields: Final = _as_object(event) or _as_object(attributes) - if fields is None: - return - for name, value in tuple(fields.items()): - if name in _RESPONSES_STRUCTURAL_FIELDS or name.endswith("_id"): - continue - if isinstance(value, str) and value: - slots.append((value, functools.partial(_write_field, event, name))) - class LLMShieldProxyGuardrail(CustomGuardrail): """Redacts PII before it leaves the proxy and restores it in the response. @@ -1031,10 +179,10 @@ class LLMShieldProxyGuardrail(CustomGuardrail): "outside the LiteLLM proxy; send the request through the proxy or without stream=True." ), ) - metadata: Final = _as_object(kwargs.get("litellm_metadata")) + metadata: Final = as_object(kwargs.get("litellm_metadata")) if metadata is not None: metadata[_DEPLOYMENT_RESTORE_KEY] = session_id - cache_controls: Final = _as_object(kwargs.get("cache")) + cache_controls: Final = as_object(kwargs.get("cache")) kwargs["cache"] = {**(cache_controls or {}), "no-cache": True, "no-store": True} return kwargs @@ -1052,7 +200,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): a key built from the redacted request. Outside the proxy nothing restores later, and the pre-call deployment hook has already kept that request out of the cache. """ - metadata: Final = _as_object(request_data.get("litellm_metadata")) + metadata: Final = as_object(request_data.get("litellm_metadata")) marker: Final = metadata.get(_DEPLOYMENT_RESTORE_KEY) if metadata is not None else None session_id: Final = self._minted_session_id(request_data) if session_id is None or marker != session_id: @@ -1111,7 +259,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): def _same_length_or_raise(self, returned: object, sent: Sequence[str], operation: str) -> Sequence[str]: """Guards the positional mapping the callers rely on to write results back.""" - entries: Final = _as_array(returned) + entries: Final = as_array(returned) texts: Final = tuple(entry for entry in entries or () if isinstance(entry, str)) # A non-string entry would be written into the request or reply as is. if entries is None or len(entries) != len(sent) or len(texts) != len(entries): @@ -1148,7 +296,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): Read only from `litellm_metadata`, the proxy-private store `_mint_session_id` writes to. A caller can populate `metadata`; they cannot populate this. """ - metadata: Final = _as_object(data.get("litellm_metadata")) + metadata: Final = as_object(data.get("litellm_metadata")) existing: Final = metadata.get(_SESSION_METADATA_KEY) if metadata is not None else None return existing if isinstance(existing, str) and existing.startswith(_VAULT_PREFIX) else None @@ -1166,52 +314,8 @@ class LLMShieldProxyGuardrail(CustomGuardrail): # --- request traversal -------------------------------------------------------- @staticmethod - def _locate_request_texts( - data: MutableRequest, - ) -> tuple[Sequence[_Slot], Sequence[_Slot]]: - """Finds every redactable span, split by whether the caller can see it. - - Anything missed here reaches the provider in the clear while the guardrail - still reports as enabled, so the walk covers every request shape that - carries text. - - The split exists because the response is restored against one vault only. - Server-authored spans -- system and developer turns, Anthropic's top-level - `system`, the Responses API `instructions`, tool and output schemas -- go into a - vault nothing is ever restored against, so a caller who gets the model to - echo one of their placeholders back receives the placeholder, not the value - behind it. End-user identifiers go there too: nothing in a reply needs them. - - Tool *results* stay on the caller's side deliberately. The model reads them in - order to answer, so it can already repeat anything in them; restoring the - placeholder gives the caller the answer they would have had without this - guardrail, and an agent that reads a file and quotes an address from it needs - that address back. - - `extra_body` is walked the same way as the request itself. LiteLLM merges it over - the transformed request just before sending, so a field there -- `input`, - `messages`, `system` -- replaces the redacted one on the wire. - """ - slots: Final[_SlotSink] = [] - privileged: Final[_SlotSink] = [] - for payload in (data, _as_object(data.get("extra_body"))): - if payload is None: - continue - for entry in _read_list(payload, "messages"): - message = _as_object(entry) - if message is not None: - sink = privileged if message.get("role") in _PRIVILEGED_ROLES else slots - _collect_content(message, sink) - _collect_participant_name(message, sink) - _collect_tool_arguments(message, sink) - _collect_responses_fields(payload, slots, privileged) - _collect_prompt(payload, slots) - _collect_system(payload, privileged) - _collect_tool_definitions(payload, slots, privileged) - _collect_output_contracts(payload, slots, privileged) - _collect_user_locations(payload, privileged) - _collect_end_user_ids(payload, privileged) - return tuple(slots), tuple(privileged) + def _locate_request_texts(data: MutableRequest) -> tuple[Sequence[Slot], Sequence[Slot]]: + return locate_request_texts(data) # --- hooks -------------------------------------------------------------------- @@ -1229,7 +333,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): try: slots, privileged = self._locate_request_texts(data) - except _RequestTooDeep as exc: + except RequestTooDeep as exc: raise GuardrailRaisedException( guardrail_name=self.guardrail_name, message=f"Request {exc} nests deeper than LLM Shield Proxy inspects; blocking the request.", @@ -1247,7 +351,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): await self._redact_into(slots, session_id) return data - async def _redact_into(self, slots: Sequence[_Slot], session_id: str) -> None: + async def _redact_into(self, slots: Sequence[Slot], session_id: str) -> None: """Redacts every span in `slots` under one vault and writes the result back.""" redacted: Final = await self._redact(tuple(text for text, _ in slots), session_id) for (_, write), replacement in zip(slots, redacted): @@ -1280,7 +384,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): """ if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) is not True: return response - response = _detached(response) # rebind-ok: everything below restores the copy. + response = detached(response) # rebind-ok: everything below restores the copy. if self._is_anthropic_message_response(response): return await self._restore_anthropic_response(response, data) @@ -1299,14 +403,14 @@ class LLMShieldProxyGuardrail(CustomGuardrail): # characters, and `_same_length_or_raise` is what guarantees the positional # mapping -- so a reply carrying more spans than that fails closed, which is this # guardrail's posture everywhere else. - pending: Final[_SlotSink] = [] + pending: Final[SlotSink] = [] for choice in choices: message = getattr(choice, "message", None) if message is None: # A Completions reply carries its text on the choice itself. - text = _read_field(choice, "text") + text = read_field(choice, "text") if isinstance(text, str) and text: - pending.append((text, functools.partial(_write_field, choice, "text"))) + pending.append((text, functools.partial(write_field, choice, "text"))) continue content = getattr(message, "content", None) if isinstance(content, str) and content: @@ -1334,7 +438,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): @staticmethod def _is_anthropic_message_response(response: object) -> bool: """Anthropic's native /v1/messages reply arrives as a plain dict.""" - body: Final = _as_object(response) + body: Final = as_object(response) return body is not None and body.get("type") == "message" and isinstance(body.get("content"), list) async def _restore_anthropic_response(self, response: MutableRequest, data: MutableRequest) -> MutableRequest: @@ -1347,17 +451,17 @@ class LLMShieldProxyGuardrail(CustomGuardrail): string, and the request path redacts its string leaves -- so the reply's leaves have to come back or the application invokes the tool with placeholders. """ - slots: Final[_SlotSink] = [] - for entry in _read_list(response, "content"): - block = _as_object(entry) + slots: Final[SlotSink] = [] + for entry in read_list(response, "content"): + block = as_object(entry) if block is None: continue kind = block.get("type") text = block.get("text") if kind == "text" and isinstance(text, str) and text: slots.append((text, functools.partial(block.__setitem__, "text"))) - elif kind == "tool_use" and _as_object(block.get("input")) is not None: - _collect_json_leaves(block.get("input"), slots) + elif kind == "tool_use" and as_object(block.get("input")) is not None: + collect_json_leaves(block.get("input"), slots) if not slots: return response @@ -1367,7 +471,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): return response @staticmethod - def _responses_api_slots(response: object) -> Sequence[_Slot]: + def _responses_api_slots(response: object) -> Sequence[Slot]: """Restorable spans in a Responses API reply. That shape carries `output` items rather than `choices`, so it needs its own @@ -1376,18 +480,18 @@ class LLMShieldProxyGuardrail(CustomGuardrail): come through as dicts or as objects depending on how far the reply has been deserialised, so both are handled. - The item-level fields mirror `_collect_responses_fields`, which walks the same + The item-level fields mirror `collect_responses_fields`, which walks the same fields on the request side -- a function_call item holds `arguments`, a function_call_output holds `output` -- so the two directions stay symmetric. """ - slots: Final[_SlotSink] = [] + slots: Final[SlotSink] = [] for item in getattr(response, "output", None) or (): - _collect_response_item(item, slots) + collect_response_item(item, slots) return tuple(slots) - async def _restore_responses_api_response(self, response: Any, slots: Sequence[_Slot], data: MutableRequest) -> Any: + async def _restore_responses_api_response(self, response: Any, slots: Sequence[Slot], data: MutableRequest) -> Any: """Puts the original values back into a Responses API reply.""" - await _rehydrate_slots(slots, functools.partial(self._rehydrate, session_id=self._session_id(data))) + await rehydrate_slots(slots, functools.partial(self._rehydrate, session_id=self._session_id(data))) return response async def async_post_call_streaming_iterator_hook( @@ -1416,9 +520,9 @@ class LLMShieldProxyGuardrail(CustomGuardrail): session_id: Final = self._session_id(request_data) step: Final = functools.partial(self._stream_step, session_id=session_id) rehydrate: Final = functools.partial(self._rehydrate, session_id=session_id) - sse: Final = _AnthropicSSERestorer(step) - events: Final = _ResponsesStreamRestorer(step, rehydrate) - carries: Final[_CarryWindows] = {} + sse: Final = AnthropicSSERestorer(step) + events: Final = ResponsesStreamRestorer(step, rehydrate) + carries: Final[CarryWindows] = {} last_chunk = None # rebind-ok: tracks the most recent chunk for the final flush. async for chunk in response: @@ -1429,11 +533,11 @@ class LLMShieldProxyGuardrail(CustomGuardrail): # Chunks are restored as copies, for the reason the non-streaming hook copies: # LiteLLM keeps the chunks it yielded to assemble the reply it caches and logs, # so restoring them in place would cache this caller's plaintext. - if _responses_event_type(chunk) is not None: - for event in await events.restore(_detached(chunk)): + if responses_event_type(chunk) is not None: + for event in await events.restore(detached(chunk)): yield event continue - restored_chunk = _detached(chunk) + restored_chunk = detached(chunk) last_chunk = restored_chunk for choice in getattr(restored_chunk, "choices", None) or (): await self._restore_choice(choice, carries, session_id) @@ -1448,7 +552,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): async for trailing in self._flush_trailing(last_chunk, carries, session_id): yield trailing - async def _restore_choice(self, choice: object, carries: _CarryWindows, session_id: str) -> None: + async def _restore_choice(self, choice: object, carries: CarryWindows, session_id: str) -> None: """Restores one choice's delta, advancing that choice's own windows. Content and each tool call are separate token streams, so each gets its own @@ -1457,7 +561,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): held back for one stream onto another. """ delta: Final = getattr(choice, "delta", None) - index: Final = _choice_index(choice) + index: Final = choice_index(choice) is_final: Final = bool(getattr(choice, "finish_reason", None)) if isinstance(choice, TextChoices): # A Completions stream carries its text on the choice itself, with no delta @@ -1481,27 +585,27 @@ class LLMShieldProxyGuardrail(CustomGuardrail): async def _restore_text_window( self, choice: object, - key: _CarryKey, - carries: _CarryWindows, + key: CarryKey, + carries: CarryWindows, session_id: str, is_final: bool, ) -> None: """Restores a Completions stream choice's `text` through its window.""" carry: Final = carries.get(key, "") - text: Final = _read_field(choice, "text") + text: Final = read_field(choice, "text") if not isinstance(text, str) or not text: if not (is_final and carry): return emitted, remaining = await self._stream_step(text if isinstance(text, str) else "", carry, is_final, session_id) carries[key] = remaining # rebind-ok: this stream's window advances. if emitted or text: - _write_field(choice, "text", emitted) + write_field(choice, "text", emitted) async def _restore_content_window( self, delta: Any, - key: _CarryKey, - carries: _CarryWindows, + key: CarryKey, + carries: CarryWindows, session_id: str, is_final: bool, ) -> None: @@ -1526,7 +630,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): self, tool_call: Any, choice_index: int, - carries: _CarryWindows, + carries: CarryWindows, session_id: str, ) -> None: """Restores one streamed tool call's argument fragment. @@ -1535,26 +639,26 @@ class LLMShieldProxyGuardrail(CustomGuardrail): concatenate per tool-call index, so each index gets a window of its own rather than sharing the content stream's. """ - tool_index: Final = _read_field(tool_call, "index") + tool_index: Final = read_field(tool_call, "index") if not isinstance(tool_index, int): return - function: Final = _read_field(tool_call, "function") + function: Final = read_field(tool_call, "function") if function is None: return - arguments: Final = _read_field(function, "arguments") + arguments: Final = read_field(function, "arguments") if not isinstance(arguments, str) or not arguments: return key: Final = (choice_index, tool_index) emitted, remaining = await self._stream_step(arguments, carries.get(key, ""), False, session_id) carries[key] = remaining # rebind-ok: this tool call's window advances. - _write_field(function, "arguments", emitted) + write_field(function, "arguments", emitted) async def _flush_finished_choice( self, delta: Any, choice_index: int, - carries: _CarryWindows, + carries: CarryWindows, session_id: str, ) -> None: """Emits everything this finishing choice still holds, into this chunk. @@ -1567,7 +671,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): even when this chunk already carried a fragment for that tool call. """ continuations: Final[list[dict[str, object]]] = [] # mutable-ok: built into this chunk's delta. - for key in sorted((held for held in carries if held[0] == choice_index), key=_carry_sort_key): + for key in sorted((held for held in carries if held[0] == choice_index), key=carry_sort_key): carry = carries[key] if not carry: continue @@ -1579,13 +683,13 @@ class LLMShieldProxyGuardrail(CustomGuardrail): if tool_index is None: delta.content = text else: - continuations.extend(_continuation_delta(tool_index, text)) + continuations.extend(continuation_delta(tool_index, text)) if continuations: existing: Final = tuple(getattr(delta, "tool_calls", None) or ()) delta.tool_calls = [*existing, *continuations] async def _flush_trailing( - self, last_chunk: Any, carries: _CarryWindows, session_id: str + self, last_chunk: Any, carries: CarryWindows, session_id: str ) -> AsyncGenerator[Any, None]: """Empties every window still holding text, one chunk per window. @@ -1598,7 +702,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): finished earlier is not present in the terminal chunk, and flushing only what that chunk carries would drop its held text and truncate its answer. """ - for key in sorted(carries, key=_carry_sort_key): + for key in sorted(carries, key=carry_sort_key): carry = carries[key] if not carry: continue @@ -1611,17 +715,17 @@ class LLMShieldProxyGuardrail(CustomGuardrail): if chunk is None: continue choice: object = chunk.choices[0] - delta = _read_field(choice, "delta") + delta = read_field(choice, "delta") if isinstance(choice, TextChoices): choice.text = text elif tool_index is None: - _write_field(delta, "content", text) + write_field(delta, "content", text) else: # The copy carried this chunk's own content and tool calls, both already # delivered. Replace rather than append, and drop the content, or the # client sees them twice. - _write_field(delta, "content", None) - _write_field(delta, "tool_calls", _continuation_delta(tool_index, text)) + write_field(delta, "content", None) + write_field(delta, "tool_calls", continuation_delta(tool_index, text)) yield chunk @staticmethod @@ -1636,7 +740,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): if not raw_choices: return None choices: Final[tuple[object, ...]] = tuple(raw_choices) - position: Final = next((at for at, choice in enumerate(choices) if _choice_index(choice) == index), 0) + position: Final = next((at for at, choice in enumerate(choices) if choice_index(choice) == index), 0) kept: Final = raw_choices[position] if getattr(kept, "delta", None) is None and not isinstance(kept, TextChoices): return None @@ -1697,11 +801,11 @@ class LLMShieldProxyGuardrail(CustomGuardrail): spans: Final[list[str]] = list(text_list) # mutable-ok: ordered batch, frozen before the call. writers: Final[list[Callable[[str], None]]] = [] # mutable-ok: one per span appended below. for call in restored_calls: - function = _read_field(call, "function") - arguments = _read_field(function, "arguments") if function is not None else None + function = read_field(call, "function") + arguments = read_field(function, "arguments") if function is not None else None if isinstance(arguments, str) and arguments: spans.append(arguments) - writers.append(functools.partial(_write_field, function, "arguments")) + writers.append(functools.partial(write_field, function, "arguments")) replaced: Final = ( await self._redact(tuple(spans), self._mint_session_id(request_data)) diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/payload.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/payload.py new file mode 100644 index 00000000000..ac4adbfc481 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/payload.py @@ -0,0 +1,206 @@ +import copy +import functools +from collections.abc import Awaitable, Callable, Sequence +from typing import ( + Final, + TypeAlias, +) + +# The proxy's own request dict. Mutable by design: a pre-call guardrail rewrites +# the caller's payload in place, which is the entire point of the hook. +MutableRequest: TypeAlias = dict[str, object] + +# A JSON body on its way to httpx, which requires a real dict rather than a view. +JsonBody: TypeAlias = dict[str, object] + +# How far a JSON value -- a tool input, a parameter schema -- is followed on the request +# side. Legitimate JSON nests far deeper than content blocks do, so the bound is +# generous; past it the request is refused, since text past the bound would reach the +# provider unredacted. +MAX_JSON_DEPTH: Final = 64 + +# One redactable span: the text as it stands, and the write that puts the +# replacement back where it came from. +Slot: TypeAlias = tuple[str, Callable[[str], None]] + +# One incremental rehydration step for a stream the caller has already bound to its +# vault: (new text, carried window, final) -> (text safe to emit, window still held). +StreamStep: TypeAlias = Callable[[str, str, bool], Awaitable[tuple[str, str]]] + +# A batch rehydration already bound to the request's vault. +Rehydrate: TypeAlias = Callable[[Sequence[str]], Awaitable[Sequence[str]]] + +# The accumulator the collectors below append into. It never escapes +# _locate_request_texts, which freezes it into a tuple before returning. +SlotSink: TypeAlias = list[Slot] + +# A caller-owned list whose entries are rewritten in place, such as a Completions +# `prompt` sent as an array of strings. +MutableSeq: TypeAlias = list[object] + + +def as_object(value: object) -> MutableRequest | None: + """`value` as a JSON object, or None. + + `isinstance(value, dict)` alone leaves the keys and values unknown to the type + checker. A JSON object's keys are strings, so the type is stated once, here. + """ + return value if isinstance(value, dict) else None + + +def as_array(value: object) -> MutableSeq | None: + """`value` as a JSON array, or None. See `as_object`.""" + return value if isinstance(value, list) else None + + +def detached(value: object) -> object: + """A deep copy of a reply or chunk, for restoring without touching LiteLLM's own object. + + LiteLLM keeps the object it handed the hooks to fill its response cache and its + logs, so writing restored plaintext into that object would put it there too. + """ + return copy.deepcopy(value) + + +def is_container(value: object) -> bool: + """Whether `value` is a JSON object or array, without narrowing it to unknown types.""" + return isinstance(value, (dict, list)) + + +def collect(container: MutableRequest, key: str, slots: SlotSink) -> None: + """Records the string at `key`, along with the write that replaces it.""" + value: Final = container.get(key) + if isinstance(value, str) and value: + slots.append((value, functools.partial(container.__setitem__, key))) + + +def collect_entry(entries: MutableSeq, index: int, slots: SlotSink) -> None: + """Records a string held directly in a list, rather than under a key.""" + value: Final = entries[index] + if isinstance(value, str) and value: + slots.append((value, functools.partial(entries.__setitem__, index))) + + +class RequestTooDeep(Exception): + """A request nests text past a walk's bound. + + Skipping the rest would forward it unredacted while the guardrail reports as + enabled, so the pre-call hook refuses the request instead. + """ + + +def collect_text_parts(container: MutableRequest, key: str, slots: SlotSink) -> None: + """Collects the `text` of every part in the list held at `key`.""" + for entry in as_array(container.get(key)) or (): + part = as_object(entry) + if part is not None: + collect(part, "text", slots) + + +def choice_index(choice: object) -> int: + """Streaming choices are matched across chunks by their index.""" + index: Final = getattr(choice, "index", 0) + return index if isinstance(index, int) else 0 + + +def read_field(holder: object, name: str) -> object: + """Reads one field from a dict or from an object. + + LiteLLM's replies arrive as Pydantic models on some paths and as plain dicts on + others, depending how far they have been deserialised, so every response walk here + has to handle 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. + + The entries are the reply's own objects, so writing through them edits the reply. + """ + value: Final = read_field(holder, name) + if isinstance(value, tuple): + return value + return tuple(as_array(value) or ()) + + +def write_field(holder: object, name: str, value: object) -> None: + """Writes one string field back into a dict or an object. Pairs with read_field.""" + if isinstance(holder, dict): + holder[name] = value + else: + setattr(holder, name, value) + + +def collect_json_leaves(node: object, slots: SlotSink, *, strict: bool = False) -> None: + """Collects every string leaf of a JSON-ish structure, with a write-back per leaf. + + An Anthropic `tool_use` block carries `input`, an arbitrary JSON object rather than a + string, so a value worth restoring can sit at any depth. Bounded by `MAX_JSON_DEPTH`: + the shape is caller or model controlled, and the bound is what stops a crafted one from + becoming an unbounded descent. Walked with an explicit stack rather than recursively, + so a deeply nested value cannot spend stack frames proportional to attacker-chosen + depth. + + `strict` is for the request side, where a leaf left behind would reach the provider + unredacted: past the bound it raises `RequestTooDeep`. On the reply side a leaf past + the bound just keeps its placeholder, which leaks nothing, so it is skipped. + """ + pending: Final[list[tuple[object, int]]] = [(node, 0)] # mutable-ok: local walk stack. + while pending: + current, current_depth = pending.pop() + if current_depth > MAX_JSON_DEPTH: + if strict and is_container(current) and current: + raise RequestTooDeep("json") + continue + current_object = as_object(current) + if current_object is not None: + for key in tuple(current_object): + value = current_object[key] + if isinstance(value, str) and value: + slots.append((value, functools.partial(current_object.__setitem__, key))) + else: + pending.append((value, current_depth + 1)) + continue + entries = as_array(current) + if entries is not None: + for index, value in enumerate(entries): + if isinstance(value, str) and value: + slots.append((value, functools.partial(entries.__setitem__, index))) + else: + pending.append((value, current_depth + 1)) + + +def collect_response_item(item: object, slots: SlotSink) -> None: + """Restorable spans in one Responses API output item, dict or object. + + Mirrors `collect_responses_fields` on the request side -- a function_call or + mcp_call item holds `arguments`, their outputs `output`, a reasoning item `summary` + parts -- so the two directions stay symmetric. A custom tool call carries `input` and + a code interpreter call `code`, both model-written. + """ + for block in read_list(item, "content"): + for field in ("text", "refusal"): + text = read_field(block, field) + if isinstance(text, str) and text: + slots.append((text, lambda new, b=block, f=field: write_field(b, f, new))) + for part in read_list(item, "summary"): + text = read_field(part, "text") + if isinstance(text, str) and text: + slots.append((text, lambda new, p=part: write_field(p, "text", new))) + for field in ("arguments", "output", "input", "code"): + value = read_field(item, field) + if isinstance(value, str) and value: + slots.append((value, lambda new, i=item, f=field: write_field(i, f, new))) + + +async def rehydrate_slots(slots: Sequence[Slot], rehydrate: Rehydrate) -> None: + """Restores every span in `slots` in one batch and writes each result back.""" + if not slots: + return + restored: Final = await rehydrate(tuple(text for text, _ in slots)) + for (_, write), replacement in zip(slots, restored): + write(replacement) diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/request_walk.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/request_walk.py new file mode 100644 index 00000000000..224ea840ccb --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/request_walk.py @@ -0,0 +1,402 @@ +from collections.abc import Sequence +from typing import ( + Final, +) + +from .payload import ( + MAX_JSON_DEPTH, + MutableRequest, + RequestTooDeep, + Slot, + SlotSink, + as_array, + as_object, + collect, + collect_entry, + collect_json_leaves, + collect_text_parts, + is_container, + read_list, +) + +# Roles whose text the application author wrote and the caller never sees. Their +# PII is still redacted outbound, but it is not restorable from the reply. +PRIVILEGED_ROLES: Final = frozenset({"system", "developer"}) + +# How far a tool_result chain is followed. Real payloads nest one or two deep; the +# bound is what stops a crafted one from becoming an unbounded walk. A request that +# nests deeper is refused rather than forwarded, because text past the bound would +# otherwise reach the provider unredacted. +MAX_CONTENT_DEPTH: Final = 8 + +# JSON Schema keywords whose value has to reach the model or a validator verbatim, so the +# schema walk leaves them alone: types, formats, patterns, references and +# required-property lists. Everything else is scanned. +SCHEMA_STRUCTURAL_KEYWORDS: Final = frozenset( + ( + "type", + "format", + "pattern", + "required", + "dependentRequired", + "propertyOrdering", + "discriminator", + "contentEncoding", + "contentMediaType", + "$ref", + "$id", + "$schema", + "$anchor", + "$dynamicRef", + "$dynamicAnchor", + "$recursiveRef", + "$recursiveAnchor", + "$vocabulary", + ) +) + +# Keywords holding JSON values rather than schemas: every string in them is collected, +# whatever the keys around it are called. +SCHEMA_VALUE_KEYWORDS: Final = frozenset(("examples", "default")) + +# Keywords holding the literal values the model must reproduce. These go to the CALLER's +# vault, not the privileged one: the model emits the stand-in in its tool arguments or +# structured output, and restoring the reply turns it back into the value the schema +# allows, so the call still routes. In the non-restorable vault it would come back as a +# stand-in no validator accepts. +SCHEMA_LITERAL_KEYWORDS: Final = frozenset(("enum", "const")) + +# Keywords whose value maps names to subschemas. Their keys are property names, not +# keywords, so a property called `type` or `enum` is walked like any other subschema. +SCHEMA_MAP_KEYWORDS: Final = frozenset( + ("properties", "patternProperties", "$defs", "definitions", "dependentSchemas", "dependencies") +) + + +def collect_prompt(data: MutableRequest, slots: SlotSink) -> None: + """The Completions API sends its text in `prompt`, and its tail in `suffix`.""" + collect(data, "suffix", slots) + prompt: Final = data.get("prompt") + if isinstance(prompt, str): + collect(data, "prompt", slots) + return + prompt_object: Final = as_object(prompt) + if prompt_object is not None: + # A Responses API PromptObject. `variables` are substituted into the stored + # prompt on the provider side, so they are caller text. `id` and `version` + # identify which prompt to use and must arrive unchanged. A variable is a string + # or a typed input such as `{"type": "input_text", "text": ...}`. + variables: Final = as_object(prompt_object.get("variables")) + if variables is not None: + for name in tuple(variables): + collect(variables, name, slots) + typed = as_object(variables[name]) + if typed is not None: + collect(typed, "text", slots) + return + entries: Final = as_array(prompt) + if entries is None: + return + for index in range(len(entries)): + collect_entry(entries, index, slots) + + +def collect_content(container: MutableRequest, slots: SlotSink) -> None: + """Collects `content`, a string or a list of typed parts. + + An Anthropic tool_result nests its own content, so this has to descend. It walks + with an explicit stack and a depth bound rather than by recursion: the nesting is + caller controlled, and an unbounded descent is a JSON bomb. Content nested past the + bound raises `RequestTooDeep` rather than being skipped. + """ + # Walked in document order: the shield maps its replies back by position, so the + # order spans are collected in is part of the contract. + pending: Final[list[tuple[MutableRequest, int]]] = [(container, 0)] # mutable-ok: local queue, never escapes. + cursor = 0 # rebind-ok: advances through the queue. + while cursor < len(pending): + node, depth = pending[cursor] + cursor += 1 + content = node.get("content") + if isinstance(content, str): + collect(node, "content", slots) + continue + if depth >= MAX_CONTENT_DEPTH and content: + raise RequestTooDeep("content") + for item in as_array(content) or (): + part = as_object(item) + if part is None: + continue + # Image and audio parts have no text and fall through untouched. + collect(part, "text", slots) + if part.get("type") == "tool_use": + # A replayed Anthropic tool call. Its `input` is a JSON object rather than + # a string, so a value can sit at any depth -- the reply side walks the + # same leaves when it restores one. Its own JSON bound applies, not the + # content one, and past it the request is refused. + collect_json_leaves(part.get("input"), slots, strict=True) + source = as_object(part.get("source")) if part.get("type") == "document" else None + if source is not None: + # An Anthropic document carries text inline: a `text` source holds it in + # `data`, a `content` source as a string or blocks, walked like any other + # content. Base64, URL and file sources are binary or remote, and pass + # untouched. Its `title` and `context` are caller text too. + collect(part, "title", slots) + collect(part, "context", slots) + if source.get("type") == "text": + collect(source, "data", slots) + elif source.get("type") == "content": + pending.append((source, depth + 1)) + if "content" in part: + pending.append((part, depth + 1)) + + +def collect_participant_name(message: MutableRequest, slots: SlotSink) -> None: + """Redacts `name` where it identifies a person, never where it names a function. + + On a user or assistant turn `name` is the participant, which is personal data. + On a tool or function turn the same field carries the function's name and has + to reach the provider unchanged, or the call no longer routes. + """ + if message.get("role") in ("tool", "function"): + return + collect(message, "name", slots) + + +def collect_tool_arguments(message: MutableRequest, slots: SlotSink) -> None: + """Tool arguments carry the values a user asked the model to act on.""" + for tool_call in read_list(message, "tool_calls"): + tool_call_object = as_object(tool_call) + function = as_object(tool_call_object.get("function")) if tool_call_object is not None else None + if function is not None: + collect(function, "arguments", slots) + legacy: Final = as_object(message.get("function_call")) + if legacy is not None: + collect(legacy, "arguments", slots) + + +def collect_system(data: MutableRequest, slots: SlotSink) -> None: + """Anthropic's /v1/messages carries its system prompt at the top level.""" + system: Final = data.get("system") + if isinstance(system, str): + collect(data, "system", slots) + return + collect_text_parts(data, "system", slots) + + +def collect_responses_fields(data: MutableRequest, slots: SlotSink, privileged: SlotSink) -> None: + """The Responses API sends text outside `messages`, in `instructions` and `input`. + + `instructions` is written by the application, not by the caller, so it is + collected into the privileged sink; `input` is the caller's own text, except for + system and developer items in it, which go to the privileged sink like their Chat + counterparts. + """ + collect(data, "instructions", privileged) + request_input: Final = data.get("input") + if isinstance(request_input, str): + collect(data, "input", slots) + return + entries: Final = as_array(request_input) + if entries is None: + return + for index, entry in enumerate(entries): + if isinstance(entry, str): + # The embeddings and moderations shape: `input` as an array of strings. + collect_entry(entries, index, slots) + continue + item = as_object(entry) + if item is None: + continue + collect_content(item, privileged if item.get("role") in PRIVILEGED_ROLES else slots) + # A function_call item holds `arguments`; a function_call_output holds `output`, + # as a string or as a list of input_text parts. A custom_tool_call holds `input` + # and a code_interpreter_call `code` -- the fields the reply side restores. + collect(item, "arguments", slots) + collect(item, "output", slots) + collect_text_parts(item, "output", slots) + collect(item, "input", slots) + collect(item, "code", slots) + # A replayed reasoning item carries the model's summary of its own reasoning, + # which quotes whatever the conversation contained. + collect_text_parts(item, "summary", slots) + + +def collect_tool_definitions(data: MutableRequest, slots: SlotSink, privileged: SlotSink) -> None: + """Tool definitions are application-authored free text bound for the provider. + + A tool's description and the free text in its parameter schema are where callers put + examples and customer context, so they carry PII as often as a prompt does. They are + collected into the privileged sink, like a system prompt: redacted outbound, and never + restorable from the reply. `enum` and `const` values are the exception, and go to the + caller's vault -- see `SCHEMA_LITERAL_KEYWORDS`. Names and types are left as sent. + + Covers Chat `tools[].function`, the legacy `functions[]`, and the flat tool shape the + Responses API and Anthropic share, whose schema is `parameters` or `input_schema`. + """ + for key in ("tools", "functions"): + for entry in as_array(data.get(key)) or (): + tool = as_object(entry) + if tool is None: + continue + function = as_object(tool.get("function")) + for holder in (tool, function) if function is not None else (tool,): + collect(holder, "description", privileged) + collect_schema_text(holder.get("parameters"), slots, privileged) + collect_schema_text(holder.get("input_schema"), slots, privileged) + + +def collect_schema_text(schema: object, slots: SlotSink, privileged: SlotSink) -> None: + """Collects the text in a JSON Schema, at any depth. + + Scan by default: every string is collected except under the keywords in + `SCHEMA_STRUCTURAL_KEYWORDS`, whose values must go out verbatim. A list of keywords + *to* collect would leak every one it forgot -- draft-07 `dependencies`, a `$comment`, + a vendor `x-` extension -- which is how this walk started out. Free text goes to the + privileged sink; `enum` / `const` literals go to the caller's, so the model's use of + them is restored. + + Structure matters in two places. Under `properties` and the other name -> subschema + maps, keys are property names rather than keywords, so a property called `type` is a + subschema to walk, not a keyword to skip. And `examples` / `default` hold JSON values, + so all their strings are collected whatever the keys around them are called. Nested + past `MAX_JSON_DEPTH`, the request is refused. + """ + pending: Final[list[tuple[object, int]]] = [(schema, 0)] # mutable-ok: local walk stack. + while pending: + node, depth = pending.pop() + if depth > MAX_JSON_DEPTH: + if is_container(node) and node: + raise RequestTooDeep("schema") + continue + entries = as_array(node) + if entries is not None: + for index, item in enumerate(entries): + collect_entry(entries, index, privileged) + if is_container(item): + pending.append((item, depth + 1)) + continue + schema_object = as_object(node) + if schema_object is None: + continue + for keyword, value in tuple(schema_object.items()): + if keyword in SCHEMA_STRUCTURAL_KEYWORDS: + continue + subschemas = as_object(value) if keyword in SCHEMA_MAP_KEYWORDS else None + if keyword in SCHEMA_LITERAL_KEYWORDS: + collect(schema_object, keyword, slots) + collect_json_leaves(value, slots, strict=True) + elif keyword in SCHEMA_VALUE_KEYWORDS: + collect(schema_object, keyword, privileged) + collect_json_leaves(value, privileged, strict=True) + elif subschemas is not None: + pending.extend((child, depth + 1) for child in subschemas.values()) + elif isinstance(value, str): + collect(schema_object, keyword, privileged) + elif is_container(value): + pending.append((value, depth + 1)) + + +def collect_output_contracts(data: MutableRequest, slots: SlotSink, privileged: SlotSink) -> None: + """Text the caller sends to shape the reply rather than to prompt it. + + A predicted output (`prediction.content`) is the caller's own draft of the answer, so + it goes with their text: the model largely repeats it, and it has to come back. A + structured-output schema -- Chat `response_format.json_schema`, Responses + `text.format` -- is application-authored like a tool schema, so its free text goes + to the privileged sink, and its names and types stay as sent. + """ + prediction: Final = as_object(data.get("prediction")) + if prediction is not None: + collect(prediction, "content", slots) + collect_text_parts(prediction, "content", slots) + response_format: Final = as_object(data.get("response_format")) + text_options: Final = as_object(data.get("text")) + for declared in ( + response_format.get("json_schema") if response_format is not None else None, + text_options.get("format") if text_options is not None else None, + ): + wrapper = as_object(declared) + if wrapper is not None: + collect(wrapper, "description", privileged) + collect_schema_text(wrapper.get("schema"), slots, privileged) + + +def collect_user_locations(data: MutableRequest, privileged: SlotSink) -> None: + """Web search forwards the user's approximate location, whose `city` and `region` + are free text and can hold a street address. + + Chat carries it in `web_search_options.user_location.approximate`; the Responses + and Anthropic web-search tools carry it flat on the tool's `user_location`. Nothing + restores it from a reply, hence the privileged sink. + """ + options: Final = data.get("web_search_options") + tools: Final = as_array(data.get("tools")) or () + for declared in (options, *tools): + holder = as_object(declared) + location = as_object(holder.get("user_location")) if holder is not None else None + if location is None: + continue + approximate = as_object(location.get("approximate")) + for container in (location, approximate) if approximate is not None else (location,): + collect(container, "city", privileged) + collect(container, "region", privileged) + + +def collect_end_user_ids(data: MutableRequest, privileged: SlotSink) -> None: + """`user` and `safety_identifier` are forwarded to the provider and often hold an email. + + Only detected PII is replaced, so an opaque id reaches the provider unchanged. LiteLLM's + own end-user spend tracking reads the id resolved at authentication, before this hook + runs, so rewriting the field here does not move spend. Nothing restores these from a + reply, hence the privileged sink. + """ + collect(data, "user", privileged) + collect(data, "safety_identifier", privileged) + + +def locate_request_texts( + data: MutableRequest, +) -> tuple[Sequence[Slot], Sequence[Slot]]: + """Finds every redactable span, split by whether the caller can see it. + + Anything missed here reaches the provider in the clear while the guardrail + still reports as enabled, so the walk covers every request shape that + carries text. + + The split exists because the response is restored against one vault only. + Server-authored spans -- system and developer turns, Anthropic's top-level + `system`, the Responses API `instructions`, tool and output schemas -- go into a + vault nothing is ever restored against, so a caller who gets the model to + echo one of their placeholders back receives the placeholder, not the value + behind it. End-user identifiers go there too: nothing in a reply needs them. + + Tool *results* stay on the caller's side deliberately. The model reads them in + order to answer, so it can already repeat anything in them; restoring the + placeholder gives the caller the answer they would have had without this + guardrail, and an agent that reads a file and quotes an address from it needs + that address back. + + `extra_body` is walked the same way as the request itself. LiteLLM merges it over + the transformed request just before sending, so a field there -- `input`, + `messages`, `system` -- replaces the redacted one on the wire. + """ + slots: Final[SlotSink] = [] + privileged: Final[SlotSink] = [] + for payload in (data, as_object(data.get("extra_body"))): + if payload is None: + continue + for entry in read_list(payload, "messages"): + message = as_object(entry) + if message is not None: + sink = privileged if message.get("role") in PRIVILEGED_ROLES else slots + collect_content(message, sink) + collect_participant_name(message, sink) + collect_tool_arguments(message, sink) + collect_responses_fields(payload, slots, privileged) + collect_prompt(payload, slots) + collect_system(payload, privileged) + collect_tool_definitions(payload, slots, privileged) + collect_output_contracts(payload, slots, privileged) + collect_user_locations(payload, privileged) + collect_end_user_ids(payload, privileged) + return tuple(slots), tuple(privileged) diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py new file mode 100644 index 00000000000..feab9b01ef1 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py @@ -0,0 +1,368 @@ +import copy +import functools +import itertools +import json +import re +from enum import Enum +from types import MappingProxyType +from typing import ( + Final, + TypeAlias, +) + +from .payload import ( + JsonBody, + MutableRequest, + Rehydrate, + SlotSink, + StreamStep, + as_object, + collect_response_item, + read_field, + read_list, + rehydrate_slots, + write_field, +) + +# Anthropic /v1/messages delta types that carry restorable text, and the field holding +# it. `thinking_delta` is left out on purpose: a thinking block is signed, and one +# rewritten here fails verification when the client sends it back on the next turn. +ANTHROPIC_DELTA_FIELDS: Final = MappingProxyType({"text_delta": "text", "input_json_delta": "partial_json"}) + +# A blank line ends an SSE event. Frames are cut there, never inside an event, so a +# `data:` line split across two network chunks is parsed only once it is whole. +SSE_EVENT_BOUNDARY: Final = re.compile(rb"(\r?\n\r?\n)") + +# What an SSE stream can open with: one of its fields, or a `:` comment. +SSE_OPENINGS: Final = (b"event:", b"data:", b"id:", b"retry:", b":") + +# Responses API delta events whose `delta` is not text. Audio arrives base64-encoded; +# sending it through the shield would cost a round trip per chunk to restore nothing. +RESPONSES_BINARY_DELTAS: Final = frozenset(("response.audio.delta",)) + +# Fields on a Responses API event that identify something rather than say something. +# Every other string field on a `.done` event is model text and is restored, so an event +# type added upstream is covered by default instead of leaking a placeholder. +RESPONSES_STRUCTURAL_FIELDS: Final = frozenset( + ("type", "id", "item_id", "call_id", "name", "server_label", "status", "obfuscation") +) + +# Terminal Responses API events that repeat the whole reply under `response`. +RESPONSES_TERMINAL_EVENTS: Final = frozenset(("response.completed", "response.incomplete")) + +# Sliding windows keyed by (choice index, tool-call index | None), threaded through one +# stream. `None` is the content channel; an int is one tool call's accumulating +# `arguments`. Content and each tool call are separate token streams, so each needs its +# own window -- one shared window would splice one stream's held-back tail onto another. +CarryKey: TypeAlias = tuple[int, int | None] +CarryWindows: TypeAlias = dict[CarryKey, str] + +# A Responses API delta stream: (event family, item id, output index, part index). +ResponsesStreamKey: TypeAlias = tuple[str, object, object, object] + + +def carry_sort_key(key: CarryKey) -> tuple[int, int]: + """Orders streaming windows without ever comparing None to an int. + + `sorted()` over the raw keys raises as soon as one choice holds both a content window + and a tool-call window, because `None < 0` is not orderable. Content sorts first, then + tool calls by their index. + """ + choice_index, tool_index = key + return (choice_index, -1 if tool_index is None else tool_index) + + +def continuation_delta(tool_index: int, text: str) -> list[dict[str, object]]: + """A `tool_calls` delta carrying `text` as an index-only continuation. + + Clients concatenate tool-call fragments by index, so no id or name is needed. + """ + return [{"index": tool_index, "function": {"arguments": text}}] + + +def opens_like_sse(head: bytes) -> bool | None: + """Whether a raw stream is SSE, judged by its opening bytes; None while undecidable. + + An SSE stream opens with a field name or a `:` comment. Anything else -- a JSON array + streamed in pieces, say -- has no event boundaries to wait for. A chunk that ends + partway through a field name decides nothing yet, so that case waits for more. + """ + opening: Final = head.lstrip() + if not opening: + return None + if opening.startswith(SSE_OPENINGS): + return True + if any(field.startswith(opening) for field in SSE_OPENINGS): + return None + return False + + +def responses_event_type(chunk: object) -> str | None: + """The event type of a Responses API stream event, or None for any other chunk. + + The type arrives as a plain string on dicts and as a str-valued Enum on LiteLLM's + event models. The Enum is unwrapped because it does not hash like its value, so it + would miss every lookup in the event tables above. + """ + if isinstance(chunk, (bytes, str)): + return None + kind: Final = read_field(chunk, "type") + value: Final = kind.value if isinstance(kind, Enum) else kind + return value if isinstance(value, str) and value.startswith("response.") else None + + +class AnthropicSSERestorer: + """Restores an Anthropic `/v1/messages` stream, which reaches the hook as raw SSE. + + Each content block is its own token stream with its own window, keyed by the block's + `index`: `text_delta` carries prose and `input_json_delta` a tool call's arguments. + When a block stops, whatever its window still holds is emitted as one more delta for + that block, just ahead of the `content_block_stop` frame, so the client has the whole + block before it is told the block is complete. + + Frames are processed whole. A network chunk can end in the middle of an event, so the + unfinished tail is kept until the rest arrives; that delays one partial event, never + a completed one. A frame that is not an Anthropic event -- another endpoint's SSE, or + anything that fails to parse -- is passed through byte for byte, and a raw stream that + does not open like SSE at all is passed through chunk by chunk, never buffered. + """ + + def __init__(self, step: StreamStep) -> None: + self._step: Final = step + self._carries: Final[dict[int, str]] = {} # mutable-ok: per-block windows advanced in place. + self._delta_types: Final[dict[int, str]] = {} # mutable-ok: each block's delta type, for its flush. + self._pending = b"" + self._as_text = False + self._is_sse: bool | None = None + + async def feed(self, chunk: bytes | str) -> tuple[bytes | str, ...]: + """Restores every event this chunk completes; holds back an unfinished tail.""" + if isinstance(chunk, str): + self._as_text = True + if self._is_sse is False: + return (chunk,) + raw: Final = chunk.encode("utf-8") if isinstance(chunk, str) else chunk + buffered: Final = self._pending + raw + if self._is_sse is None: + self._is_sse = opens_like_sse(buffered) + if self._is_sse is None: + # Too little has arrived to tell -- `b"eve"` could still become `event:`. + self._pending = buffered + return () + if not self._is_sse: + self._pending = b"" + return self._emit(buffered) + boundaries: Final = tuple(SSE_EVENT_BOUNDARY.finditer(buffered)) + if not boundaries: + self._pending = buffered + return () + cut: Final = boundaries[-1].end() + self._pending = buffered[cut:] + # With a capturing group, split alternates event, separator, ..., and ends in the + # empty remainder after the last separator. + parts: Final = SSE_EVENT_BOUNDARY.split(buffered[:cut]) + restored: Final = tuple( + [await self._restore_event(parts[index]) + parts[index + 1] for index in range(0, len(parts) - 1, 2)] + ) + return self._emit(b"".join(restored)) + + async def finish(self) -> tuple[bytes | str, ...]: + """Emits an unterminated final event and any window a block never closed.""" + held: Final = self._pending + self._pending = b"" + if not self._is_sse: + # The stream ended before it could be told apart from SSE: hand it back as is. + return self._emit(held) + tail: Final = await self._restore_event(held) if held.strip() else held + flushed: Final = await self._flush_all() + # The tail had no blank line after it; one is needed before another frame follows. + separator: Final = b"\n\n" if tail.strip() and flushed else b"" + return self._emit(tail + separator + flushed) + + def _emit(self, frames: bytes) -> tuple[bytes | str, ...]: + if not frames: + return () + return (frames.decode("utf-8") if self._as_text else frames,) + + async def _restore_event(self, block: bytes) -> bytes: + """Rewrites one SSE event, or returns it untouched if it carries nothing to restore.""" + try: + lines: Final = block.decode("utf-8").split("\n") + except UnicodeDecodeError: + return block + data_lines: Final = tuple(index for index, line in enumerate(lines) if line.startswith("data:")) + if len(data_lines) != 1: + return block + line: Final = lines[data_lines[0]] + try: + parsed: Final[object] = json.loads(line[len("data:") :]) + except ValueError: + return block + event: Final = as_object(parsed) + if event is None: + return block + kind: Final = event.get("type") + index: Final = event.get("index") + if kind == "content_block_stop" and isinstance(index, int): + return await self._flush(index) + block + if kind == "message_stop": + return await self._flush_all() + block + if kind != "content_block_delta" or not await self._restore_delta(event): + return block + ending: Final = "\r" if line.endswith("\r") else "" + rewritten: Final = ( + *lines[: data_lines[0]], + f"data: {json.dumps(event, ensure_ascii=False)}{ending}", + *lines[data_lines[0] + 1 :], + ) + return "\n".join(rewritten).encode("utf-8") + + async def _restore_delta(self, event: MutableRequest) -> bool: + """Advances one block's window through this delta. False if it holds no text.""" + index: Final = event.get("index") + delta: Final = event.get("delta") + if not isinstance(index, int) or not isinstance(delta, dict): + return False + delta_type: Final = delta.get("type") + if not isinstance(delta_type, str): + return False + field: Final = ANTHROPIC_DELTA_FIELDS.get(delta_type) + text: Final = delta.get(field) if field is not None else None + if field is None or not isinstance(text, str) or not text: + return False + emitted, remaining = await self._step(text, self._carries.get(index, ""), False) + self._carries[index] = remaining + self._delta_types[index] = delta_type + delta[field] = emitted + return True + + async def _flush(self, index: int) -> bytes: + """One synthetic delta frame carrying whatever `index`'s window still holds.""" + carry: Final = self._carries.pop(index, "") + delta_type: Final = self._delta_types.pop(index, None) + field: Final = ANTHROPIC_DELTA_FIELDS.get(delta_type) if isinstance(delta_type, str) else None + if not carry or field is None: + return b"" + text, _ = await self._step("", carry, True) + if not text: + return b"" + event: Final[JsonBody] = { + "type": "content_block_delta", + "index": index, + "delta": {"type": delta_type, field: text}, + } + return f"event: content_block_delta\ndata: {json.dumps(event, ensure_ascii=False)}\n\n".encode() + + async def _flush_all(self) -> bytes: + flushed: Final = tuple([await self._flush(index) for index in tuple(self._carries)]) + return b"".join(flushed) + + +class ResponsesStreamRestorer: + """Restores a Responses API event stream. + + The event families are matched by shape rather than listed, so a text stream the + API adds later is restored by default instead of leaking a placeholder: + + - Any `*.delta` event whose `delta` is a string is a token stream (output_text, + refusal, function-call and MCP arguments, reasoning summaries, ...). Each gets its + own window, keyed by the family, the item id and the part index. + - Any `*.done` event closes the stream of the same family. Whatever its window still + holds goes out first, as a copy of that stream's last delta event -- so it carries + the stream's own ids, and repeats that event's `sequence_number`. Then every text + field on the done event is restored in full: its string fields other than + identifiers, plus any `part` or `item` it repeats. + - `response.completed` / `response.incomplete` repeat the whole reply, and are + restored the same way the non-streaming reply is. + """ + + def __init__(self, step: StreamStep, rehydrate: Rehydrate) -> None: + self._step: Final = step + self._rehydrate: Final = rehydrate + self._carries: Final[dict[ResponsesStreamKey, str]] = {} # mutable-ok: per-stream windows advanced in place. + self._last_deltas: Final[dict[ResponsesStreamKey, object]] = {} # mutable-ok: newest delta per stream. + + async def restore(self, event: object) -> tuple[object, ...]: + """The events to emit in place of `event`: any flush, then the event itself.""" + kind: Final = responses_event_type(event) + if kind is None: + return (event,) + if kind.endswith(".delta") and kind not in RESPONSES_BINARY_DELTAS: + await self._restore_delta(event, kind) + return (event,) + slots: Final[SlotSink] = [] + flushed: Final = await self._flush(responses_stream_key(event, kind)) if kind.endswith(".done") else () + if kind.endswith(".done"): + collect_event_text(event, slots) + part: Final = read_field(event, "part") + if part is not None: + collect_response_item({"content": [part]}, slots) + collect_response_item(read_field(event, "item"), slots) + elif kind in RESPONSES_TERMINAL_EVENTS: + for item in read_list(read_field(event, "response"), "output"): + collect_response_item(item, slots) + await rehydrate_slots(slots, self._rehydrate) + return (*flushed, event) + + async def finish(self) -> tuple[object, ...]: + """Flushes every stream the provider never closed, e.g. a truncated reply.""" + flushed: Final = tuple([await self._flush(key) for key in tuple(self._carries)]) + return tuple(itertools.chain.from_iterable(flushed)) + + async def _restore_delta(self, event: object, kind: str) -> None: + text: Final = read_field(event, "delta") + if not isinstance(text, str) or not text: + return + key: Final = responses_stream_key(event, kind) + emitted, remaining = await self._step(text, self._carries.get(key, ""), False) + self._carries[key] = remaining + self._last_deltas[key] = event + write_field(event, "delta", emitted) + + async def _flush(self, key: ResponsesStreamKey) -> tuple[object, ...]: + carry: Final = self._carries.pop(key, "") + template: Final = self._last_deltas.pop(key, None) + if not carry or template is None: + return () + text, _ = await self._step("", carry, True) + if not text: + return () + flush: Final = copy.deepcopy(template) + write_field(flush, "delta", text) + return (flush,) + + +def responses_stream_key(event: object, kind: str) -> ResponsesStreamKey: + """Identifies the delta stream an event belongs to, the same for its delta and done. + + The family is the event type without its `.delta` / `.done` suffix, so an output_text + stream and a refusal stream on the same part never share a window. + """ + family: Final = kind.rsplit(".", 1)[0] + part_index: Final = read_field(event, "content_index") + summary_index: Final = read_field(event, "summary_index") + return ( + family, + read_field(event, "item_id"), + read_field(event, "output_index"), + part_index if part_index is not None else summary_index, + ) + + +def collect_event_text(event: object, slots: SlotSink) -> None: + """Collects every top-level text field of a Responses API event, dict or model. + + Scan by default, with identifiers excluded, rather than a list of known fields: the + `.done` event of each stream family names its text differently (`text`, `refusal`, + `arguments`, ...), and a family added upstream would otherwise leak a placeholder. + """ + # A model's fields live in its `__dict__`; an empty dict has none either way. + attributes: Final[object] = getattr(event, "__dict__", None) + fields: Final = as_object(event) or as_object(attributes) + if fields is None: + return + for name, value in tuple(fields.items()): + if name in RESPONSES_STRUCTURAL_FIELDS or name.endswith("_id"): + continue + if isinstance(value, str) and value: + slots.append((value, functools.partial(write_field, event, name)))