From db419125402af8476b604eea9f26bdec34ed601a Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Mon, 5 Oct 2026 14:31:05 +0000 Subject: [PATCH] style(guardrails): drop routine comments from llm_shield_proxy AGENTS.md keeps source comments to tool directives and genuinely complex logic; the rationale stays in the docstrings. --- .../llm_shield_proxy/llm_shield_proxy.py | 82 ------------------- .../llm_shield_proxy/payload.py | 16 ---- .../llm_shield_proxy/request_walk.py | 39 --------- .../llm_shield_proxy/stream_restorers.py | 23 ------ 4 files changed, 160 deletions(-) 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 0291dba2ac3..fa23a110a74 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 @@ -1,10 +1,3 @@ -# +-------------------------------------------------------------+ -# -# Use LLM Shield Proxy for reversible PII redaction -# https://github.com/ninadphalak/LLM-Shield-Proxy -# -# +-------------------------------------------------------------+ - import copy import functools import os @@ -79,24 +72,10 @@ _REDACT_PATH: Final = "/v1/guard/redact" _REHYDRATE_PATH: Final = "/v1/guard/rehydrate" _REHYDRATE_STREAM_PATH: Final = "/v1/guard/rehydrate/stream" -# The session id ties a redact call to the rehydrate calls that undo it. It is -# stored on the request dict rather than on the guardrail instance: the proxy -# registers one instance process-wide, so instance attributes would be shared -# across concurrent requests. _SESSION_METADATA_KEY: Final = "llm_shield_session_id" -# Set when the deployment pre-call hook redacted the request -- model-level `guardrails` -# outside the proxy -- to that request's vault id. Only then is the reply restored at the -# deployment, because only then does no later hook restore it. Matching it against the -# minted id, which carries the unguessable per-process prefix, means a caller cannot opt -# a proxy request into deployment-level restoration by sending the key themselves. _DEPLOYMENT_RESTORE_KEY: Final = "llm_shield_restore_at_deployment" -# 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 -# placeholder echoed back. The per-process prefix means a caller cannot even name -# a vault this process uses. _VAULT_PREFIX: Final = f"litellm-{uuid.uuid4().hex}" _DEFAULT_TIMEOUT_SECONDS: Final = 10.0 @@ -116,9 +95,6 @@ class LLMShieldProxyGuardrail(CustomGuardrail): split across two chunks is never emitted in fragments. """ - # Our redaction and restoration run in the native lifecycle hooks below. Without - # this the proxy would route every event through the unified apply_guardrail path - # and the streaming hook would never fire. use_native_lifecycle_hooks: ClassVar[bool] = True def __init__( @@ -207,8 +183,6 @@ class LLMShieldProxyGuardrail(CustomGuardrail): return None return await super().async_post_call_success_deployment_hook(request_data, response, call_type) - # --- transport --------------------------------------------------------------- - def _headers(self, session_id: str) -> JsonBody: headers: Final[JsonBody] = { "Content-Type": "application/json", @@ -261,7 +235,6 @@ class LLMShieldProxyGuardrail(CustomGuardrail): """Guards the positional mapping the callers rely on to write results back.""" 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): raise GuardrailRaisedException( guardrail_name=self.guardrail_name, @@ -269,8 +242,6 @@ class LLMShieldProxyGuardrail(CustomGuardrail): ) return texts - # --- session ------------------------------------------------------------------ - @staticmethod def _mint_session_id(data: MutableRequest) -> str: """Mints a vault id for this request, overwriting anything already there. @@ -280,10 +251,6 @@ class LLMShieldProxyGuardrail(CustomGuardrail): caller from reaching another caller's vault. """ session_id: Final = f"{_VAULT_PREFIX}-{uuid.uuid4().hex}" - # `litellm_metadata` is proxy-private; `metadata` is forwarded to the provider on - # /v1/responses. The session id is a capability against the vault's rehydrate - # endpoint, so handing it to the provider alongside the placeholders would let the - # provider read back exactly what this guardrail exists to withhold. metadata: Final = data.setdefault("litellm_metadata", {}) if isinstance(metadata, dict): metadata[_SESSION_METADATA_KEY] = session_id @@ -311,14 +278,10 @@ class LLMShieldProxyGuardrail(CustomGuardrail): existing: Final = LLMShieldProxyGuardrail._minted_session_id(data) return existing if existing is not None else f"{_VAULT_PREFIX}-{uuid.uuid4().hex}" - # --- request traversal -------------------------------------------------------- - @staticmethod def _locate_request_texts(data: MutableRequest) -> tuple[Sequence[Slot], Sequence[Slot]]: return locate_request_texts(data) - # --- hooks -------------------------------------------------------------------- - @log_guardrail_information async def async_pre_call_hook( self, @@ -343,9 +306,6 @@ class LLMShieldProxyGuardrail(CustomGuardrail): session_id: Final = self._mint_session_id(data) if privileged: - # A vault of its own, whose id is deliberately never stored: the - # response is restored against `session_id` alone, so nothing the - # model emits can turn one of these placeholders back into plaintext. await self._redact_into(privileged, f"{_VAULT_PREFIX}-{uuid.uuid4().hex}") if slots: await self._redact_into(slots, session_id) @@ -357,15 +317,6 @@ class LLMShieldProxyGuardrail(CustomGuardrail): for (_, write), replacement in zip(slots, redacted): write(replacement) - # KNOWN LIMIT: a tool call's `arguments` is a JSON *string*, and a restored value is - # spliced into it as raw text. If the original value contained a double quote, a - # backslash or a newline, the reassembled document is no longer valid JSON for a - # strict parser. The proxy's own tool-argument rehydration has the same property - # (`_rehydrate_json_response` in api/main.py), so this is a pre-existing limit of the - # product rather than one introduced here. Escaping is deliberately NOT applied as a - # fix: a fragment is an arbitrary slice of a JSON document, so the code cannot tell - # whether the position it writes is inside a string literal, and escaping - # unconditionally would corrupt the values that are not. async def async_post_call_success_hook( self, data: MutableRequest, @@ -397,17 +348,10 @@ class LLMShieldProxyGuardrail(CustomGuardrail): if not choices: return response - # One batch for every restorable span in the reply, collected in document order: - # the shield maps its answers back by position. A second round trip is not an - # option here -- /v1/guard/rehydrate caps a batch at 256 texts and 1,000,000 - # 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] = [] 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") if isinstance(text, str) and text: pending.append((text, functools.partial(write_field, choice, "text"))) @@ -415,9 +359,6 @@ class LLMShieldProxyGuardrail(CustomGuardrail): content = getattr(message, "content", None) if isinstance(content, str) and content: pending.append((content, functools.partial(setattr, message, "content"))) - # A tool call's `arguments` is model-generated text and the request path - # redacts it, so leaving it unrestored hands the application a placeholder to - # invoke a tool with. These are Pydantic objects on this path, not dicts. for tool_call in getattr(message, "tool_calls", None) or (): function = getattr(tool_call, "function", None) arguments = getattr(function, "arguments", None) if function is not None else None @@ -530,9 +471,6 @@ class LLMShieldProxyGuardrail(CustomGuardrail): for frames in await sse.feed(chunk): yield frames continue - # 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)): yield event @@ -543,7 +481,6 @@ class LLMShieldProxyGuardrail(CustomGuardrail): await self._restore_choice(choice, carries, session_id) yield restored_chunk - # A stream that ended early can still leave text held back, in any shape. for frames in await sse.finish(): yield frames for event in await events.finish(): @@ -564,8 +501,6 @@ class LLMShieldProxyGuardrail(CustomGuardrail): 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 - # and no tool calls: one window, the content one. await self._restore_text_window(choice, (index, None), carries, session_id, is_final) return if delta is None: @@ -577,9 +512,6 @@ class LLMShieldProxyGuardrail(CustomGuardrail): await self._restore_tool_call_window(tool_call, index, carries, session_id) if is_final: - # A client parses a tool call's arguments when it sees the finish_reason, so - # every window this choice still holds has to land in *this* chunk. Flushing - # after it produces argument JSON the client has already stopped waiting for. await self._flush_finished_choice(delta, index, carries, session_id) async def _restore_text_window( @@ -614,7 +546,6 @@ class LLMShieldProxyGuardrail(CustomGuardrail): text: Final = getattr(delta, "content", None) if not isinstance(text, str) or not text: - # Nothing to restore here, but a final chunk still has to flush the window. if is_final and carry: flushed, remaining = await self._stream_step("", carry, True, session_id) carries[key] = remaining # rebind-ok: this stream's window advances. @@ -721,9 +652,6 @@ class LLMShieldProxyGuardrail(CustomGuardrail): elif tool_index is None: 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)) yield chunk @@ -745,12 +673,8 @@ class LLMShieldProxyGuardrail(CustomGuardrail): if getattr(kept, "delta", None) is None and not isinstance(kept, TextChoices): return None kept.index = index - # The terminal signal, if there was one, already went out with the real chunk. kept.finish_reason = None chunk.choices = [kept] - # So did the usage, which `stream_options.include_usage` puts on that last chunk. A - # client that sums usage across chunks would count the request twice; a mid-stream - # chunk carries no `usage` attribute at all, so the copy drops it. if hasattr(chunk, "usage"): del chunk.usage return chunk @@ -771,8 +695,6 @@ class LLMShieldProxyGuardrail(CustomGuardrail): ) return emitted, remaining - # --- unified API (powers the UI "Test guardrail" button) ----------------------- - @log_guardrail_information async def apply_guardrail( self, @@ -795,8 +717,6 @@ class LLMShieldProxyGuardrail(CustomGuardrail): if not text_list and not tool_calls: return inputs - # Copied rather than mutated: the caller's tool calls are theirs to own, and this - # method's contract is to hand back a new mapping. restored_calls: Final[list[object]] = [copy.deepcopy(call) for call in tool_calls] # mutable-ok: a new list. 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. @@ -816,8 +736,6 @@ class LLMShieldProxyGuardrail(CustomGuardrail): for write, replacement in zip(writers, restored_values[len(text_list) :]): write(replacement) - # Return a new mapping rather than rewriting the caller's, so this stays a - # pure transform of the inputs it was handed. merged: Final[JsonBody] = {**inputs} if text_list: merged["texts"] = restored_values[: len(text_list)] diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/payload.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/payload.py index ac4adbfc481..644ac76efb8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/payload.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/payload.py @@ -6,36 +6,20 @@ from typing import ( 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] 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 index 224ea840ccb..c9c0d2ccc47 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/request_walk.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/request_walk.py @@ -19,19 +19,10 @@ from .payload import ( 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", @@ -55,19 +46,10 @@ SCHEMA_STRUCTURAL_KEYWORDS: Final = frozenset( ) ) -# 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") ) @@ -82,10 +64,6 @@ def collect_prompt(data: MutableRequest, slots: SlotSink) -> None: 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): @@ -109,8 +87,6 @@ def collect_content(container: MutableRequest, slots: SlotSink) -> None: 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): @@ -126,20 +102,11 @@ def collect_content(container: MutableRequest, slots: SlotSink) -> None: 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": @@ -201,23 +168,17 @@ def collect_responses_fields(data: MutableRequest, slots: SlotSink, privileged: 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) 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 index feab9b01ef1..4ef8bbdbb3e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py @@ -24,40 +24,23 @@ from .payload import ( 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] @@ -146,7 +129,6 @@ class AnthropicSSERestorer: 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: @@ -158,8 +140,6 @@ class AnthropicSSERestorer: 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)] @@ -171,11 +151,9 @@ class AnthropicSSERestorer: 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) @@ -356,7 +334,6 @@ def collect_event_text(event: object, slots: SlotSink) -> None: `.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: