scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-11 03:38:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 82

style(guardrails): drop routine comments from llm_shield_proxy

Browse source 
AGENTS.md keeps source comments to tool directives and genuinely complex logic;
the rationale stays in the docstrings.
This commit is contained in:
Ninad Phalak 2026-10-05 14:31:05 +00:00
parent 464516a188
commit db41912540
No known key found for this signature in database
4 changed files with 0 additions and 160 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

82
litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py
View file

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

16
litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/payload.py
View file

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

39
litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/request_walk.py
View file

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

23
litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py
View file

@ -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: