From 4509c874701f4b8243479a76b0144995120e380a Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Thu, 3 Sep 2026 11:47:57 -0500 Subject: [PATCH 01/22] feat(guardrails): add llm shield pii redaction and rehydration guardrail LLM Shield is a self-hosted PII gateway. This adds it as a guardrail so a proxy operator can redact personal data out of outbound requests and have the original values restored in the model's reply. The substitution is reversible, which is the difference from a masking guardrail. Outbound text is replaced with placeholders held in a session vault inside the operator's own LLM Shield deployment, and the reply is restored before it reaches the caller, so the end user still sees real values while the provider never received them. Streaming responses are restored incrementally. LLM Shield holds back only the trailing characters that could still turn out to be part of a placeholder, so tokens are forwarded as they arrive rather than the whole response being collected first. A placeholder split across two chunks is never emitted in fragments. The integration talks to LLM Shield over HTTP and adds no dependency. Notes for reviewers: - The guardrail sets use_native_lifecycle_hooks, since redaction and restoration need the native pre-call, post-call and streaming hooks rather than the unified path. - Per-request state lives on the request dict, never on the guardrail instance, because the proxy registers a single instance process-wide. The streaming carry-over is a local of the generator for the same reason. - Every failure blocks the request. A redaction guardrail that fails open would send the exact data it exists to protect to the provider. --- .../guardrail_hooks/llm_shield/__init__.py | 33 ++ .../guardrail_hooks/llm_shield/llm_shield.py | 351 ++++++++++++++++++ litellm/types/guardrails.py | 1 + .../guardrails/guardrail_hooks/llm_shield.py | 24 ++ .../guardrail_hooks/test_llm_shield.py | 261 +++++++++++++ 5 files changed, 670 insertions(+) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/llm_shield/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/llm_shield.py create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/__init__.py new file mode 100644 index 00000000000..a6cc54d5408 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/__init__.py @@ -0,0 +1,33 @@ +from typing import TYPE_CHECKING, Final + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .llm_shield import LLMShieldGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): + import litellm + + _llm_shield_guardrail_callback: Final = LLMShieldGuardrail( + api_key=litellm_params.api_key, + api_base=litellm_params.api_base, + guardrail_name=guardrail.get("guardrail_name", ""), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + + litellm.logging_callback_manager.add_litellm_callback(_llm_shield_guardrail_callback) + return _llm_shield_guardrail_callback + + +guardrail_initializer_registry: Final = { # mutable-ok: module-level registry, built once and never mutated + SupportedGuardrailIntegrations.LLM_SHIELD.value: initialize_guardrail, +} + + +guardrail_class_registry: Final = { # mutable-ok: module-level registry, built once and never mutated + SupportedGuardrailIntegrations.LLM_SHIELD.value: LLMShieldGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py new file mode 100644 index 00000000000..4199ad5ca65 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py @@ -0,0 +1,351 @@ +# +-------------------------------------------------------------+ +# +# Use LLM Shield for reversible PII redaction +# https://github.com/ninadphalak/LLM-Shield-Proxy +# +# +-------------------------------------------------------------+ + +import os +import uuid +from collections.abc import AsyncGenerator +from typing import ( + TYPE_CHECKING, + Any, # noqa: TID251 # **kwargs forwards verbatim to CustomGuardrail.__init__ + ClassVar, + Final, + Literal, + Optional, +) + +import httpx + +from litellm._logging import verbose_proxy_logger +from litellm.exceptions import GuardrailRaisedException +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + get_session_id_from_request_data, + log_guardrail_information, +) +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.caching.caching import DualCache + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +GUARDRAIL_NAME: Final = "llm_shield" + +_DEFAULT_API_BASE: Final = "http://localhost:8000" +_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" + +_DEFAULT_TIMEOUT_SECONDS: Final = 10.0 + + +class LLMShieldGuardrail(CustomGuardrail): + """Redacts PII before it leaves the proxy and restores it in the response. + + Unlike a masking guardrail, the substitution is reversible. Outbound text is + replaced with placeholders held in a session vault inside the user's own LLM + Shield deployment; the model's reply is then restored so the end user sees the + original values while the provider never received them. + + Streaming is restored incrementally rather than by buffering the response. LLM + Shield holds back only the trailing characters that could still turn out to be + part of a placeholder, so tokens are forwarded as they arrive and a placeholder + 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__( + self, + guardrail_name: str = GUARDRAIL_NAME, + api_base: str | None = None, + api_key: str | None = None, + **kwargs: Any, + ) -> None: + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + self.api_base: Final = (api_base or os.environ.get("LLM_SHIELD_API_BASE") or _DEFAULT_API_BASE).rstrip("/") + self.api_key: Final = api_key or os.environ.get("LLM_SHIELD_API_KEY") + super().__init__(guardrail_name=guardrail_name, **kwargs) + + @classmethod + def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: + return [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call] + + # --- transport --------------------------------------------------------------- + + def _headers(self, session_id: str) -> dict: + headers = {"Content-Type": "application/json", "X-Session-ID": session_id} + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + return headers + + async def _call_shield(self, path: str, session_id: str, payload: dict) -> dict: + """Posts to LLM Shield, failing closed on any transport or status error. + + A redaction guardrail that fails open sends the very data it exists to + protect to a third-party provider, so an unreachable or erroring shield + blocks the request instead of passing it through. + """ + try: + response = await self.async_handler.post( + f"{self.api_base}{path}", + headers=self._headers(session_id), + json=payload, + timeout=_DEFAULT_TIMEOUT_SECONDS, + ) + response.raise_for_status() + return response.json() + except httpx.HTTPStatusError as exc: + verbose_proxy_logger.exception("LLM Shield returned %s for %s", exc.response.status_code, path) + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=f"LLM Shield returned {exc.response.status_code}; blocking the request.", + ) from exc + except Exception as exc: + verbose_proxy_logger.exception("LLM Shield call to %s failed", path) + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message="LLM Shield is unreachable; blocking the request.", + ) from exc + + async def _redact(self, texts: list, session_id: str) -> list: + body = await self._call_shield(_REDACT_PATH, session_id, {"texts": texts}) + return self._same_length_or_raise(body.get("texts"), texts, "redact") + + async def _rehydrate(self, texts: list, session_id: str) -> list: + body = await self._call_shield(_REHYDRATE_PATH, session_id, {"texts": texts}) + return self._same_length_or_raise(body.get("texts"), texts, "rehydrate") + + def _same_length_or_raise(self, returned: Any, sent: list, operation: str) -> list: + """Guards the positional mapping the callers rely on to write results back.""" + if not isinstance(returned, list) or len(returned) != len(sent): + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=f"LLM Shield {operation} returned an unexpected payload; blocking the request.", + ) + return returned + + # --- session ------------------------------------------------------------------ + + def _session_id(self, data: dict) -> str: + """Returns a session id stable across this request's hooks.""" + metadata = data.setdefault("metadata", {}) + if not isinstance(metadata, dict): + return f"litellm-{uuid.uuid4().hex}" + existing = metadata.get(_SESSION_METADATA_KEY) + if isinstance(existing, str) and existing: + return existing + session_id = get_session_id_from_request_data(data) or f"litellm-{uuid.uuid4().hex}" + metadata[_SESSION_METADATA_KEY] = session_id + return session_id + + # --- message traversal -------------------------------------------------------- + + @staticmethod + def _locate_texts(messages: list) -> list: + """Finds every text span in a message list. + + Returns ``(message_index, part_index_or_None, text)``. The list form is the + multimodal shape, where only ``text`` parts carry redactable content. + """ + located = [] + for message_index, message in enumerate(messages): + if not isinstance(message, dict): + continue + content = message.get("content") + if isinstance(content, str) and content: + located.append((message_index, None, content)) + elif isinstance(content, list): + for part_index, part in enumerate(content): + if not isinstance(part, dict) or part.get("type") != "text": + continue + text = part.get("text") + if isinstance(text, str) and text: + located.append((message_index, part_index, text)) + return located + + @staticmethod + def _write_back(messages: list, located: list, replacements: list) -> None: + for (message_index, part_index, _), replacement in zip(located, replacements): + if part_index is None: + messages[message_index]["content"] = replacement + else: + messages[message_index]["content"][part_index]["text"] = replacement + + # --- hooks -------------------------------------------------------------------- + + @log_guardrail_information + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: "DualCache", + data: dict, + call_type: str, + ) -> dict | None: + """Replaces PII in the outbound messages with vault placeholders.""" + if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is not True: + return data + + messages = data.get("messages") + if not isinstance(messages, list): + return data + + located = self._locate_texts(messages) + if not located: + return data + + redacted = await self._redact([text for _, _, text in located], self._session_id(data)) + self._write_back(messages, located, redacted) + return data + + @log_guardrail_information + async def async_post_call_success_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + ) -> Any: + """Restores the original values in a non-streaming response.""" + if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) is not True: + return response + + choices = getattr(response, "choices", None) + if not choices: + return response + + pending = [] + for choice in choices: + message = getattr(choice, "message", None) + content = getattr(message, "content", None) + if isinstance(content, str) and content: + pending.append((message, content)) + + if not pending: + return response + + restored = await self._rehydrate([text for _, text in pending], self._session_id(data)) + for (message, _), replacement in zip(pending, restored): + message.content = replacement + return response + + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + request_data: dict, + ) -> AsyncGenerator[Any, None]: + """Restores original values incrementally, without buffering the stream. + + The carry-over window is a local of this generator, so it is scoped to one + stream and cannot leak between concurrent requests. LLM Shield returns the + text that is safe to emit now plus the trailing characters it is still + holding, which are sent back with the next delta. + """ + if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True: + async for chunk in response: + yield chunk + return + + session_id = self._session_id(request_data) + carry = "" + last_chunk = None + + async for chunk in response: + last_chunk = chunk + delta = self._stream_delta(chunk) + text = getattr(delta, "content", None) if delta is not None else None + is_final = self._is_final_chunk(chunk) + + if not isinstance(text, str) or not text: + # Nothing to restore in this chunk, but a final chunk still has to + # flush whatever the window is holding. + if is_final and carry: + body = await self._stream_step("", carry, True, session_id) + carry = body["carry"] + if body["text"] and delta is not None: + delta.content = body["text"] + yield chunk + continue + + body = await self._stream_step(text, carry, is_final, session_id) + carry = body["carry"] + delta.content = body["text"] + yield chunk + + # A stream that ended without a finish_reason can still leave text held back. + if carry and last_chunk is not None: + body = await self._stream_step("", carry, True, session_id) + if body["text"]: + trailing = last_chunk.model_copy(deep=True) + trailing_delta = self._stream_delta(trailing) + if trailing_delta is not None: + trailing_delta.content = body["text"] + yield trailing + + async def _stream_step(self, text: str, carry: str, final: bool, session_id: str) -> dict: + body = await self._call_shield( + _REHYDRATE_STREAM_PATH, + session_id, + {"text": text, "carry": carry, "final": final}, + ) + emitted = body.get("text") + remaining = body.get("carry") + if not isinstance(emitted, str) or not isinstance(remaining, str): + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message="LLM Shield stream rehydration returned an unexpected payload.", + ) + return {"text": emitted, "carry": remaining} + + @staticmethod + def _stream_delta(chunk: Any) -> Any: + choices = getattr(chunk, "choices", None) + if not choices: + return None + return getattr(choices[0], "delta", None) + + @staticmethod + def _is_final_chunk(chunk: Any) -> bool: + choices = getattr(chunk, "choices", None) + if not choices: + return False + return bool(getattr(choices[0], "finish_reason", None)) + + # --- unified API (powers the UI "Test guardrail" button) ----------------------- + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + texts = inputs.get("texts") + if not texts: + return inputs + + session_id = self._session_id(request_data) + if input_type == "request": + inputs["texts"] = await self._redact(list(texts), session_id) + else: + inputs["texts"] = await self._rehydrate(list(texts), session_id) + return inputs diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index c17103da890..42ab6034069 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -137,6 +137,7 @@ class SupportedGuardrailIntegrations(Enum): COMPRESR = "compresr" STRAIKER = "straiker" ALICE = "alice" + LLM_SHIELD = "llm_shield" class Role(Enum): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/llm_shield.py b/litellm/types/proxy/guardrails/guardrail_hooks/llm_shield.py new file mode 100644 index 00000000000..8d7afd907b0 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/llm_shield.py @@ -0,0 +1,24 @@ +from pydantic import Field + +from .base import GuardrailConfigModel + + +class LLMShieldGuardrailConfigModel(GuardrailConfigModel): + api_key: str | None = Field( + default=None, + description=( + "The virtual key for the LLM Shield instance. If not provided, the " + "`LLM_SHIELD_API_KEY` environment variable is checked." + ), + ) + api_base: str | None = Field( + default=None, + description=( + "The base URL of the LLM Shield instance. If not provided, the `LLM_SHIELD_API_BASE` " + "environment variable is checked, then `http://localhost:8000`." + ), + ) + + @staticmethod + def ui_friendly_name() -> str: + return "LLM Shield" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py new file mode 100644 index 00000000000..4b39c8bf517 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py @@ -0,0 +1,261 @@ +from unittest.mock import AsyncMock + +import pytest +from httpx import Request, Response + +import litellm +from litellm.exceptions import GuardrailRaisedException +from litellm.proxy.guardrails.guardrail_hooks.llm_shield.llm_shield import ( + GUARDRAIL_NAME, + LLMShieldGuardrail, +) +from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 +from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + +def _guardrail(**overrides: object) -> LLMShieldGuardrail: + params: dict[str, object] = { + "api_key": "test-key", + "api_base": "http://shield.test", + "guardrail_name": GUARDRAIL_NAME, + "event_hook": "pre_call", + "default_on": True, + } + params.update(overrides) + return LLMShieldGuardrail(**params) + + +def _response(payload: dict, status_code: int = 200) -> Response: + return Response( + status_code=status_code, + json=payload, + request=Request("POST", "http://shield.test/v1/guard/redact"), + ) + + +def _mock_post(guardrail: LLMShieldGuardrail, *payloads: dict) -> AsyncMock: + """Queues one shield response per expected call.""" + mock = AsyncMock(side_effect=[_response(p) for p in payloads]) + guardrail.async_handler.post = mock # type: ignore[method-assign] + return mock + + +def _chunk(content: str | None, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=content), finish_reason=finish_reason)] + ) + + +async def _drain(generator) -> list: + return [chunk async for chunk in generator] + + +def test_llm_shield_guardrail_config(monkeypatch: pytest.MonkeyPatch): + """Should register through init_guardrails_v2 like any other provider.""" + monkeypatch.setattr(litellm, "guardrail_name_config_map", {}) + monkeypatch.setenv("LLM_SHIELD_API_KEY", "test-key") + + init_guardrails_v2( + all_guardrails=[ + { + "guardrail_name": "llm_shield", + "litellm_params": {"guardrail": "llm_shield", "mode": "pre_call", "default_on": True}, + } + ], + config_file_path="", + ) + + registered = [cb for cb in litellm.callbacks if isinstance(cb, LLMShieldGuardrail)] + assert len(registered) == 1 + assert registered[0].guardrail_name == "llm_shield" + + +class TestLLMShieldInitialization: + def test_api_base_defaults_to_localhost(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("LLM_SHIELD_API_BASE", raising=False) + assert _guardrail(api_base=None).api_base == "http://localhost:8000" + + def test_api_base_reads_environment(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("LLM_SHIELD_API_BASE", "http://shield.internal:9000") + assert _guardrail(api_base=None).api_base == "http://shield.internal:9000" + + def test_trailing_slash_is_stripped(self): + assert _guardrail(api_base="http://shield.test/").api_base == "http://shield.test" + + +class TestRedaction: + @pytest.mark.asyncio + async def test_string_content_is_redacted(self): + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["Email [EMAIL_1] about it"]}) + + data = {"messages": [{"role": "user", "content": "Email a@b.com about it"}]} + result = await guardrail.async_pre_call_hook( + user_api_key_dict=None, cache=None, data=data, call_type="completion" + ) + + assert result["messages"][0]["content"] == "Email [EMAIL_1] about it" + + @pytest.mark.asyncio + async def test_multimodal_text_parts_are_redacted(self): + """The list content shape is a historical bypass; text parts must be covered.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["call [PHONE_1]"]}) + + data = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "call 555-0100"}, + {"type": "image_url", "image_url": {"url": "http://x/y.png"}}, + ], + } + ] + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert data["messages"][0]["content"][0]["text"] == "call [PHONE_1]" + assert data["messages"][0]["content"][1]["image_url"]["url"] == "http://x/y.png" + + @pytest.mark.asyncio + async def test_request_without_messages_is_untouched(self): + guardrail = _guardrail() + mock = _mock_post(guardrail) + data = {"input": "no messages here"} + + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + mock.assert_not_called() + + @pytest.mark.asyncio + async def test_session_id_is_reused_across_hooks(self): + """Rehydration can only resolve tokens minted under the same session.""" + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}, {"texts": ["a@b.com"]}) + + data = {"messages": [{"role": "user", "content": "a@b.com"}]} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + await guardrail._rehydrate(["[EMAIL_1]"], guardrail._session_id(data)) + + sessions = {call.kwargs["headers"]["X-Session-ID"] for call in mock.call_args_list} + assert len(sessions) == 1 + + +class TestFailClosed: + @pytest.mark.asyncio + async def test_unreachable_shield_blocks_the_request(self): + """Failing open would send the PII upstream, defeating the guardrail.""" + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock(side_effect=ConnectionError("refused")) + + with pytest.raises(GuardrailRaisedException): + await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data={"messages": [{"role": "user", "content": "a@b.com"}]}, + call_type="completion", + ) + + @pytest.mark.asyncio + async def test_error_status_blocks_the_request(self): + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_response({"error": "nope"}, status_code=500)) + + with pytest.raises(GuardrailRaisedException): + await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data={"messages": [{"role": "user", "content": "a@b.com"}]}, + call_type="completion", + ) + + @pytest.mark.asyncio + async def test_short_payload_blocks_the_request(self): + """A response that loses an entry would silently misalign the write-back.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": []}) + + with pytest.raises(GuardrailRaisedException): + await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data={"messages": [{"role": "user", "content": "a@b.com"}]}, + call_type="completion", + ) + + +class TestStreamingRehydration: + @pytest.mark.asyncio + async def test_split_placeholder_is_not_emitted_in_fragments(self): + """The window holds back a partial placeholder and releases it once complete.""" + guardrail = _guardrail(event_hook="post_call") + # Shield holds "[EMAIL" back, then releases the restored value. + _mock_post( + guardrail, + {"text": "Email ", "carry": "[EMAIL"}, + {"text": "a@b.com about it", "carry": ""}, + ) + + async def stream(): + yield _chunk("Email [EMAIL") + yield _chunk("_1] about it", finish_reason="stop") + + chunks = await _drain( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, response=stream(), request_data={"messages": []} + ) + ) + + emitted = [c.choices[0].delta.content for c in chunks] + assert emitted == ["Email ", "a@b.com about it"] + # No fragment of the placeholder ever reached the client. + assert not any("[EMAIL" in (text or "") for text in emitted) + + @pytest.mark.asyncio + async def test_carry_is_returned_to_the_next_call(self): + guardrail = _guardrail(event_hook="post_call") + mock = _mock_post( + guardrail, + {"text": "", "carry": "hold"}, + {"text": "held-and-more", "carry": ""}, + ) + + async def stream(): + yield _chunk("hold") + yield _chunk("-and-more", finish_reason="stop") + + await _drain( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, response=stream(), request_data={"messages": []} + ) + ) + + assert mock.call_args_list[0].kwargs["json"]["carry"] == "" + assert mock.call_args_list[1].kwargs["json"]["carry"] == "hold" + assert mock.call_args_list[1].kwargs["json"]["final"] is True + + @pytest.mark.asyncio + async def test_chunks_are_forwarded_as_they_arrive(self): + """Restoration must not buffer the stream into a single terminal chunk.""" + guardrail = _guardrail(event_hook="post_call") + _mock_post( + guardrail, + {"text": "one ", "carry": ""}, + {"text": "two ", "carry": ""}, + {"text": "three", "carry": ""}, + ) + + async def stream(): + yield _chunk("one ") + yield _chunk("two ") + yield _chunk("three", finish_reason="stop") + + chunks = await _drain( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, response=stream(), request_data={"messages": []} + ) + ) + + assert len(chunks) == 3 + assert [c.choices[0].delta.content for c in chunks] == ["one ", "two ", "three"] From 8ab969d56caf5a125b88a1ada46b51b21cf4eb60 Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Thu, 3 Sep 2026 12:02:48 -0500 Subject: [PATCH 02/22] feat(ui): list llm shield in the guardrail garden Adds the card, preset and logo so operators can pick LLM Shield from the guardrails page the same way as the other partner guardrails. --- .../public/assets/logos/llm_shield.svg | 5 +++++ .../guardrails/_components/guardrail_garden_configs.ts | 6 ++++++ .../_components/guardrail_garden_data.test.ts | 1 + .../guardrails/_components/guardrail_garden_data.ts | 10 ++++++++++ .../guardrails/_components/guardrail_info_helpers.tsx | 3 +++ 5 files changed, 25 insertions(+) create mode 100644 ui/litellm-dashboard/public/assets/logos/llm_shield.svg diff --git a/ui/litellm-dashboard/public/assets/logos/llm_shield.svg b/ui/litellm-dashboard/public/assets/logos/llm_shield.svg new file mode 100644 index 00000000000..d61edff4473 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/llm_shield.svg @@ -0,0 +1,5 @@ + + + + + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts index 7785a8e44ab..b445cc9c5ad 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts @@ -318,4 +318,10 @@ export const GUARDRAIL_PRESETS: Record = { mode: "pre_call", defaultOn: false, }, + llm_shield: { + provider: "LLM Shield", + guardrailNameSuggestion: "LLM Shield", + mode: "pre_call", + defaultOn: false, + }, }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts index 1e486639840..d3212d66737 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts @@ -28,6 +28,7 @@ const EXPECTED_PARTNER_LOGO_FILES: Record = { repelloai: "repelloai.png", straiker: "straiker.svg", alice: "alice.svg", + llm_shield: "llm_shield.svg", }; describe("guardrail_garden_data logos", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts index 931b3a111d8..d46847a5a80 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts @@ -474,6 +474,16 @@ export const PARTNER_GUARDRAIL_CARDS: GuardrailCardInfo[] = [ tags: ["Content Moderation", "Prompt Injection", "PII", "Policy"], providerKey: "Alice", }, + { + id: "llm_shield", + name: "LLM Shield", + description: + "Self-hosted PII redaction that puts the original values back into the model's response, so the provider never receives personal data while the end user still sees it.", + category: "partner", + logo: guardrailLogoMap["LLM Shield"], + tags: ["PII", "Data Privacy", "Compliance", "Streaming"], + providerKey: "LLM Shield", + }, ]; export const ALL_CARDS = [...LITELLM_CONTENT_FILTER_CARDS, ...PARTNER_GUARDRAIL_CARDS]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx index f686ff5644a..f620ea6dcd0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx @@ -1,6 +1,7 @@ import aimSecurityLogo from "../../../../../public/assets/logos/aim_security.jpeg"; import aktoLogo from "../../../../../public/assets/logos/akto.svg"; import aliceLogo from "../../../../../public/assets/logos/alice.svg"; +import llmShieldLogo from "../../../../../public/assets/logos/llm_shield.svg"; import aporiaLogo from "../../../../../public/assets/logos/aporia.png"; import bedrockLogo from "../../../../../public/assets/logos/bedrock.svg"; import catoNetworksLogo from "../../../../../public/assets/logos/cato_networks.svg"; @@ -85,6 +86,7 @@ export const guardrail_provider_map: Record = { QostodianNexus: "qostodian_nexus", Repelloai: "repelloai", Alice: "alice", + "LLM Shield": "llm_shield", }; // Function to populate provider map from API response - updates the original map @@ -208,6 +210,7 @@ export const guardrailLogoMap = { "RepelloAI Argus": repelloAiLogo.src, Straiker: straikerLogo.src, Alice: aliceLogo.src, + "LLM Shield": llmShieldLogo.src, } satisfies Record; export const getGuardrailLogo = (displayName: string): string | undefined => From 2da630debd3e28be0ef5c2840b6941e4c9733bba Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Thu, 3 Sep 2026 12:15:38 -0500 Subject: [PATCH 03/22] docs(guardrails): add llm shield example config Shows both modes on one entry. Listing only pre_call redacts the request and then hands the placeholders back to the end user, so the test asserts both hooks are enabled. --- .../llm_shield/example_config.yaml | 57 +++++++++++++++++++ .../guardrail_hooks/test_llm_shield.py | 14 +++++ 2 files changed, 71 insertions(+) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/llm_shield/example_config.yaml diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/example_config.yaml b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/example_config.yaml new file mode 100644 index 00000000000..3a4b43d5432 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/example_config.yaml @@ -0,0 +1,57 @@ +# Example LiteLLM Proxy configuration for LLM Shield +# LLM Shield is a self-hosted PII gateway: https://github.com/ninadphalak/LLM-Shield-Proxy +# +# Unlike a masking guardrail, LLM Shield's substitution is reversible. Personal data is +# replaced with placeholders before the request goes to the provider, and the original +# values are put back into the model's reply, so the end user still sees real data while +# the provider never received it. + +model_list: + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY + +guardrails: + # Both modes belong on ONE entry. pre_call redacts the outbound request and post_call + # restores the reply; listing only pre_call would send placeholders back to the user. + - guardrail_name: "llm-shield" + litellm_params: + guardrail: llm_shield + mode: ["pre_call", "post_call"] + default_on: true + # Your own LLM Shield deployment. Defaults to http://localhost:8000, and also reads + # LLM_SHIELD_API_BASE from the environment. + api_base: "http://localhost:8000" + # A virtual key configured on that deployment. Also reads LLM_SHIELD_API_KEY. + api_key: os.environ/LLM_SHIELD_API_KEY + +# Usage: +# +# 1. Run LLM Shield somewhere the proxy can reach: +# pip install llm-shield-proxy +# llm-shield-proxy serve +# +# 2. Point this config at it and start the proxy: +# export LLM_SHIELD_API_KEY="your-virtual-key" +# litellm --config example_config.yaml +# +# 3. Send a request containing personal data: +# curl http://localhost:4000/v1/chat/completions \ +# -H "Authorization: Bearer sk-1234" \ +# -H "Content-Type: application/json" \ +# -d '{"model":"gpt-4o","messages":[{"role":"user","content":"Email jane.doe@example.com the invoice"}]}' +# +# The provider receives a stand-in value in place of the address. The reply you get +# back carries the real address again. +# +# Notes: +# +# - Requests are refused if LLM Shield is unreachable or returns an error, rather than +# being forwarded. Sending them on would hand the provider exactly the data this +# guardrail exists to withhold. +# - Restoring a value requires the request and the reply to share a session. LiteLLM's +# session id is used when present; otherwise one is generated per request. +# - Streaming replies are restored as chunks arrive. A placeholder split across two +# chunks is held back until it is complete, so partial values are never emitted. +# - Only text is redacted; images and audio pass through untouched. diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py index 4b39c8bf517..160c2300690 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py @@ -10,6 +10,7 @@ from litellm.proxy.guardrails.guardrail_hooks.llm_shield.llm_shield import ( LLMShieldGuardrail, ) from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices @@ -82,6 +83,19 @@ class TestLLMShieldInitialization: def test_trailing_slash_is_stripped(self): assert _guardrail(api_base="http://shield.test/").api_base == "http://shield.test" + def test_both_modes_can_be_enabled_on_one_entry(self): + """Redaction and restoration are two halves of one config entry. + + A deployment that lists only pre_call would redact the request and then hand + the placeholders straight back to the end user. + """ + guardrail = _guardrail(event_hook=["pre_call", "post_call"]) + data: dict = {"messages": []} + + assert guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is True + assert guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) is True + assert guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.during_call) is False + class TestRedaction: @pytest.mark.asyncio From 5fccbfe49f83913270116b28d18646412b9f87af Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Thu, 3 Sep 2026 12:28:18 -0500 Subject: [PATCH 04/22] feat(ui): use the llm shield brand mark for the guardrail logo --- ui/litellm-dashboard/public/assets/logos/llm_shield.svg | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/public/assets/logos/llm_shield.svg b/ui/litellm-dashboard/public/assets/logos/llm_shield.svg index d61edff4473..0dd78b078c9 100644 --- a/ui/litellm-dashboard/public/assets/logos/llm_shield.svg +++ b/ui/litellm-dashboard/public/assets/logos/llm_shield.svg @@ -1,5 +1,6 @@ - - - + + + + From ee0bac5148e56d400d699ef707ade6f8dabe55a7 Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Thu, 3 Sep 2026 14:23:14 -0500 Subject: [PATCH 05/22] fix(guardrails): restore llm shield values in anthropic replies The /v1/messages reply is a plain dict with a content block list and no choices, so it fell through the restore path and went back to the caller still carrying placeholders. The request was redacted correctly, which is what made this easy to miss. Found by running all three endpoints against a live provider; the mocked tests all passed because they only built the OpenAI shape. Adds tests for the message shape and for leaving non-text blocks alone. --- .../guardrail_hooks/llm_shield/llm_shield.py | 31 ++++++++++ .../guardrail_hooks/test_llm_shield.py | 58 ++++++++++++++++++- 2 files changed, 88 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py index 4199ad5ca65..125bea5590e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py @@ -227,6 +227,9 @@ class LLMShieldGuardrail(CustomGuardrail): if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) is not True: return response + if self._is_anthropic_message_response(response): + return await self._restore_anthropic_response(response, data) + choices = getattr(response, "choices", None) if not choices: return response @@ -246,6 +249,34 @@ class LLMShieldGuardrail(CustomGuardrail): message.content = replacement return response + @staticmethod + def _is_anthropic_message_response(response: Any) -> bool: + """Anthropic's native /v1/messages reply arrives as a plain dict.""" + return ( + isinstance(response, dict) + and response.get("type") == "message" + and isinstance(response.get("content"), list) + ) + + async def _restore_anthropic_response(self, response: dict, data: dict) -> dict: + """Restores text blocks in an Anthropic native message reply. + + This shape has no `choices`, so without its own branch the reply would go + back to the caller still carrying placeholders. + """ + blocks = [ + block + for block in response["content"] + if isinstance(block, dict) and block.get("type") == "text" and isinstance(block.get("text"), str) + ] + if not blocks: + return response + + restored = await self._rehydrate([block["text"] for block in blocks], self._session_id(data)) + for block, replacement in zip(blocks, restored): + block["text"] = replacement + return response + async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py index 160c2300690..c2d50bf5fa7 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py @@ -11,7 +11,7 @@ from litellm.proxy.guardrails.guardrail_hooks.llm_shield.llm_shield import ( ) from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices +from litellm.types.utils import Choices, Delta, Message, ModelResponse, ModelResponseStream, StreamingChoices def _guardrail(**overrides: object) -> LLMShieldGuardrail: @@ -156,6 +156,62 @@ class TestRedaction: assert len(sessions) == 1 +class TestRestoration: + @pytest.mark.asyncio + async def test_openai_shape_is_restored(self): + guardrail = _guardrail(event_hook="post_call") + _mock_post(guardrail, {"texts": ["a@b.com"]}) + + response = ModelResponse(choices=[Choices(index=0, message=Message(role="assistant", content="[EMAIL_1]"))]) + result = await guardrail.async_post_call_success_hook( + data={"messages": []}, user_api_key_dict=None, response=response + ) + + assert result.choices[0].message.content == "a@b.com" + + @pytest.mark.asyncio + async def test_anthropic_message_shape_is_restored(self): + """The /v1/messages reply is a plain dict with no choices. + + Measured against a live provider: without its own branch the reply went + back to the caller still carrying the placeholder, even though the + request had been redacted correctly. + """ + guardrail = _guardrail(event_hook="post_call") + _mock_post(guardrail, {"texts": ["a@b.com"]}) + + response = { + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "[EMAIL_1]"}], + } + result = await guardrail.async_post_call_success_hook( + data={"messages": []}, user_api_key_dict=None, response=response + ) + + assert result["content"][0]["text"] == "a@b.com" + + @pytest.mark.asyncio + async def test_anthropic_non_text_blocks_are_left_alone(self): + guardrail = _guardrail(event_hook="post_call") + _mock_post(guardrail, {"texts": ["a@b.com"]}) + + response = { + "type": "message", + "role": "assistant", + "content": [ + {"type": "text", "text": "[EMAIL_1]"}, + {"type": "tool_use", "id": "t1", "name": "lookup", "input": {}}, + ], + } + result = await guardrail.async_post_call_success_hook( + data={"messages": []}, user_api_key_dict=None, response=response + ) + + assert result["content"][0]["text"] == "a@b.com" + assert result["content"][1] == {"type": "tool_use", "id": "t1", "name": "lookup", "input": {}} + + class TestFailClosed: @pytest.mark.asyncio async def test_unreachable_shield_blocks_the_request(self): From 0de4b8b9a83ade05a7b3d6b82932568ddb18f537 Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Thu, 3 Sep 2026 14:46:28 -0500 Subject: [PATCH 06/22] docs(guardrails): correct the llm shield start command --- .../guardrails/guardrail_hooks/llm_shield/example_config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/example_config.yaml b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/example_config.yaml index 3a4b43d5432..aa63fa9d252 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/example_config.yaml +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/example_config.yaml @@ -30,7 +30,7 @@ guardrails: # # 1. Run LLM Shield somewhere the proxy can reach: # pip install llm-shield-proxy -# llm-shield-proxy serve +# llm-shield-proxy --port 8000 # # 2. Point this config at it and start the proxy: # export LLM_SHIELD_API_KEY="your-virtual-key" From b6e3e6decd82c249255e6dc7dbdd2d9b20992237 Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Thu, 3 Sep 2026 18:07:32 -0500 Subject: [PATCH 07/22] fix(guardrails): redact every request shape and restore every reply shape Three gaps, all of which let an enabled guardrail hand data to the provider or hand placeholders to the caller. Requests only walked `messages`. The Responses API `input` and tool call `arguments` went out untouched. Measured against a live provider: a request sent through `/v1/responses` reached the model with the real address in it while the guardrail reported as enabled. Request traversal now covers chat content (string and multimodal), tool call arguments, and `input` as a bare string or a list of items. Fixing that exposed the matching gap on the way back: the Responses API reply carries `output` items rather than `choices`, so it returned to the caller still holding placeholders. It now gets its own walk, handling text blocks as dicts or objects. The dashboard preset seeded only pre_call, so a guardrail created from the UI would redact the request and return the placeholders to the user. Presets can now seed both modes; the form already normalised either shape. Adds tests for each request shape, for both Responses API reply forms, and replaces a test that had asserted the `input` bypass as correct behaviour. --- .../guardrail_hooks/llm_shield/llm_shield.py | 291 +++++++++++------- ruff-strict.toml | 4 + .../guardrail_hooks/test_llm_shield.py | 126 +++++++- .../_components/add_guardrail_form.tsx | 4 +- .../_components/guardrail_garden_configs.ts | 8 +- 5 files changed, 323 insertions(+), 110 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py index 125bea5590e..fe8b70c08b8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py @@ -7,7 +7,7 @@ import os import uuid -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Callable, Mapping, Sequence from typing import ( TYPE_CHECKING, Any, # noqa: TID251 # **kwargs forwards verbatim to CustomGuardrail.__init__ @@ -15,6 +15,7 @@ from typing import ( Final, Literal, Optional, + TypeAlias, ) import httpx @@ -53,6 +54,19 @@ _SESSION_METADATA_KEY: Final = "llm_shield_session_id" _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. +# mutable-ok: the shape is fixed by CustomLogger's hook signatures. +MutableRequest: TypeAlias = dict + +# A JSON body on its way to httpx, which requires a real dict rather than a view. +# mutable-ok: handed straight to the HTTP client. +JsonBody: TypeAlias = dict + +# 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]] # mutable-ok: Callable's param list. + class LLMShieldGuardrail(CustomGuardrail): """Redacts PII before it leaves the proxy and restores it in the response. @@ -78,7 +92,7 @@ class LLMShieldGuardrail(CustomGuardrail): guardrail_name: str = GUARDRAIL_NAME, api_base: str | None = None, api_key: str | None = None, - **kwargs: Any, + **kwargs: Any, # noqa: LIT008 # kwargs-ok: forwarded verbatim to CustomGuardrail.__init__ ) -> None: self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.api_base: Final = (api_base or os.environ.get("LLM_SHIELD_API_BASE") or _DEFAULT_API_BASE).rstrip("/") @@ -86,18 +100,21 @@ class LLMShieldGuardrail(CustomGuardrail): super().__init__(guardrail_name=guardrail_name, **kwargs) @classmethod - def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: - return [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call] + def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: # mutable-ok: parent's signature. + return [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call] # mutable-ok: parent's signature. # --- transport --------------------------------------------------------------- - def _headers(self, session_id: str) -> dict: - headers = {"Content-Type": "application/json", "X-Session-ID": session_id} + def _headers(self, session_id: str) -> JsonBody: + headers: Final[JsonBody] = { # mutable-ok: httpx requires a real dict. + "Content-Type": "application/json", + "X-Session-ID": session_id, + } if self.api_key: headers["Authorization"] = f"Bearer {self.api_key}" return headers - async def _call_shield(self, path: str, session_id: str, payload: dict) -> dict: + async def _call_shield(self, path: str, session_id: str, payload: JsonBody) -> Mapping[str, object]: """Posts to LLM Shield, failing closed on any transport or status error. A redaction guardrail that fails open sends the very data it exists to @@ -105,7 +122,7 @@ class LLMShieldGuardrail(CustomGuardrail): blocks the request instead of passing it through. """ try: - response = await self.async_handler.post( + response: Final = await self.async_handler.post( f"{self.api_base}{path}", headers=self._headers(session_id), json=payload, @@ -126,69 +143,89 @@ class LLMShieldGuardrail(CustomGuardrail): message="LLM Shield is unreachable; blocking the request.", ) from exc - async def _redact(self, texts: list, session_id: str) -> list: - body = await self._call_shield(_REDACT_PATH, session_id, {"texts": texts}) + async def _redact(self, texts: Sequence[str], session_id: str) -> Sequence[str]: + payload: Final[JsonBody] = {"texts": list(texts)} # mutable-ok: JSON body for httpx. + body: Final = await self._call_shield(_REDACT_PATH, session_id, payload) return self._same_length_or_raise(body.get("texts"), texts, "redact") - async def _rehydrate(self, texts: list, session_id: str) -> list: - body = await self._call_shield(_REHYDRATE_PATH, session_id, {"texts": texts}) + async def _rehydrate(self, texts: Sequence[str], session_id: str) -> Sequence[str]: + payload: Final[JsonBody] = {"texts": list(texts)} # mutable-ok: JSON body for httpx. + body: Final = await self._call_shield(_REHYDRATE_PATH, session_id, payload) return self._same_length_or_raise(body.get("texts"), texts, "rehydrate") - def _same_length_or_raise(self, returned: Any, sent: list, operation: str) -> list: + 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.""" if not isinstance(returned, list) or len(returned) != len(sent): raise GuardrailRaisedException( guardrail_name=self.guardrail_name, message=f"LLM Shield {operation} returned an unexpected payload; blocking the request.", ) - return returned + return tuple(returned) # --- session ------------------------------------------------------------------ - def _session_id(self, data: dict) -> str: + def _session_id(self, data: MutableRequest) -> str: """Returns a session id stable across this request's hooks.""" - metadata = data.setdefault("metadata", {}) + metadata: Final = data.setdefault("metadata", {}) # mutable-ok: per-request store. if not isinstance(metadata, dict): return f"litellm-{uuid.uuid4().hex}" - existing = metadata.get(_SESSION_METADATA_KEY) + existing: Final = metadata.get(_SESSION_METADATA_KEY) if isinstance(existing, str) and existing: return existing - session_id = get_session_id_from_request_data(data) or f"litellm-{uuid.uuid4().hex}" + session_id: Final = get_session_id_from_request_data(data) or f"litellm-{uuid.uuid4().hex}" metadata[_SESSION_METADATA_KEY] = session_id return session_id - # --- message traversal -------------------------------------------------------- + # --- request traversal -------------------------------------------------------- @staticmethod - def _locate_texts(messages: list) -> list: - """Finds every text span in a message list. + def _locate_request_texts(data: MutableRequest) -> Sequence[_Slot]: + """Finds every redactable span in an outbound request. - Returns ``(message_index, part_index_or_None, text)``. The list form is the - multimodal shape, where only ``text`` parts carry redactable content. + Returns ``(text, write)`` pairs. Any shape missed here reaches the provider + in the clear, so this walks all of the request shapes that carry caller text: + + - chat ``messages``, both string and multimodal list ``content`` + - tool call ``arguments``, which routinely carry the values a user asked + the model to look up + - the Responses API ``input``, as a bare string or a list of items """ - located = [] - for message_index, message in enumerate(messages): - if not isinstance(message, dict): - continue - content = message.get("content") - if isinstance(content, str) and content: - located.append((message_index, None, content)) - elif isinstance(content, list): - for part_index, part in enumerate(content): - if not isinstance(part, dict) or part.get("type") != "text": - continue - text = part.get("text") - if isinstance(text, str) and text: - located.append((message_index, part_index, text)) - return located + slots: Final[list[_Slot]] = [] # mutable-ok: accumulator, frozen on return. - @staticmethod - def _write_back(messages: list, located: list, replacements: list) -> None: - for (message_index, part_index, _), replacement in zip(located, replacements): - if part_index is None: - messages[message_index]["content"] = replacement - else: - messages[message_index]["content"][part_index]["text"] = replacement + def add(container: MutableRequest, key: str, value: object) -> None: + if isinstance(value, str) and value: + slots.append((value, lambda new, c=container, k=key: c.__setitem__(k, new))) + + def add_content(container: MutableRequest) -> None: + """Adds `content`, which is either a string or a list of typed parts.""" + content: Final = container.get("content") + if isinstance(content, str): + add(container, "content", content) + return + for part in content if isinstance(content, list) else (): + if isinstance(part, dict): + add(part, "text", part.get("text")) + + def add_tool_calls(message: MutableRequest) -> None: + for tool_call in message.get("tool_calls") or (): + function = tool_call.get("function") if isinstance(tool_call, dict) else None + if isinstance(function, dict): + add(function, "arguments", function.get("arguments")) + + for message in data.get("messages") or (): + if isinstance(message, dict): + add_content(message) + add_tool_calls(message) + + request_input: Final = data.get("input") + if isinstance(request_input, str): + add(data, "input", request_input) + else: + for item in request_input if isinstance(request_input, list) else (): + if isinstance(item, dict): + add_content(item) + + return tuple(slots) # --- hooks -------------------------------------------------------------------- @@ -197,29 +234,26 @@ class LLMShieldGuardrail(CustomGuardrail): self, user_api_key_dict: UserAPIKeyAuth, cache: "DualCache", - data: dict, + data: MutableRequest, call_type: str, - ) -> dict | None: - """Replaces PII in the outbound messages with vault placeholders.""" + ) -> MutableRequest | None: + """Replaces PII anywhere in the outbound request with vault placeholders.""" if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is not True: return data - messages = data.get("messages") - if not isinstance(messages, list): + slots: Final = self._locate_request_texts(data) + if not slots: return data - located = self._locate_texts(messages) - if not located: - return data - - redacted = await self._redact([text for _, _, text in located], self._session_id(data)) - self._write_back(messages, located, redacted) + redacted: Final = await self._redact(tuple(text for text, _ in slots), self._session_id(data)) + for (_, write), replacement in zip(slots, redacted): + write(replacement) return data @log_guardrail_information async def async_post_call_success_hook( self, - data: dict, + data: MutableRequest, user_api_key_dict: UserAPIKeyAuth, response: Any, ) -> Any: @@ -230,27 +264,31 @@ class LLMShieldGuardrail(CustomGuardrail): if self._is_anthropic_message_response(response): return await self._restore_anthropic_response(response, data) - choices = getattr(response, "choices", None) + text_blocks: Final = self._responses_api_text_blocks(response) + if text_blocks: + return await self._restore_responses_api_response(response, text_blocks, data) + + choices: Final = getattr(response, "choices", None) if not choices: return response - pending = [] - for choice in choices: - message = getattr(choice, "message", None) - content = getattr(message, "content", None) - if isinstance(content, str) and content: - pending.append((message, content)) - + pending: Final = tuple( + (choice.message, choice.message.content) + for choice in choices + if getattr(choice, "message", None) is not None + and isinstance(getattr(choice.message, "content", None), str) + and choice.message.content + ) if not pending: return response - restored = await self._rehydrate([text for _, text in pending], self._session_id(data)) + restored: Final = await self._rehydrate(tuple(text for _, text in pending), self._session_id(data)) for (message, _), replacement in zip(pending, restored): message.content = replacement return response @staticmethod - def _is_anthropic_message_response(response: Any) -> bool: + def _is_anthropic_message_response(response: object) -> bool: """Anthropic's native /v1/messages reply arrives as a plain dict.""" return ( isinstance(response, dict) @@ -258,30 +296,68 @@ class LLMShieldGuardrail(CustomGuardrail): and isinstance(response.get("content"), list) ) - async def _restore_anthropic_response(self, response: dict, data: dict) -> dict: + async def _restore_anthropic_response(self, response: MutableRequest, data: MutableRequest) -> MutableRequest: """Restores text blocks in an Anthropic native message reply. This shape has no `choices`, so without its own branch the reply would go back to the caller still carrying placeholders. """ - blocks = [ + blocks: Final = tuple( block for block in response["content"] if isinstance(block, dict) and block.get("type") == "text" and isinstance(block.get("text"), str) - ] + ) if not blocks: return response - restored = await self._rehydrate([block["text"] for block in blocks], self._session_id(data)) + restored: Final = await self._rehydrate(tuple(block["text"] for block in blocks), self._session_id(data)) for block, replacement in zip(blocks, restored): block["text"] = replacement return response + @staticmethod + def _responses_api_text_blocks(response: object) -> Sequence[object]: + """Text blocks in a Responses API reply. + + That shape carries `output` items rather than `choices`, so it needs its own + walk; without one the reply goes back to the caller still holding + placeholders even though the request was redacted correctly. Blocks come + through as dicts or as objects depending on how far the reply has been + deserialised, so both are handled. + """ + blocks: Final[list[object]] = [] # mutable-ok: accumulator, frozen on return. + for item in getattr(response, "output", None) or (): + for block in getattr(item, "content", None) or (): + if isinstance(block, dict): + if isinstance(block.get("text"), str) and block["text"]: + blocks.append(block) + elif isinstance(getattr(block, "text", None), str) and block.text: + blocks.append(block) + return tuple(blocks) + + @staticmethod + def _block_text(block: object) -> str: + return block["text"] if isinstance(block, dict) else block.text + + async def _restore_responses_api_response( + self, response: Any, blocks: Sequence[object], data: MutableRequest + ) -> Any: + """Puts the original values back into a Responses API reply.""" + restored: Final = await self._rehydrate( + tuple(self._block_text(block) for block in blocks), self._session_id(data) + ) + for block, replacement in zip(blocks, restored): + if isinstance(block, dict): + block["text"] = replacement + else: + block.text = replacement + return response + async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, response: Any, - request_data: dict, + request_data: MutableRequest, ) -> AsyncGenerator[Any, None]: """Restores original values incrementally, without buffering the stream. @@ -295,9 +371,9 @@ class LLMShieldGuardrail(CustomGuardrail): yield chunk return - session_id = self._session_id(request_data) - carry = "" - last_chunk = None + session_id: Final = self._session_id(request_data) + carry = "" # rebind-ok: the sliding window advances with every delta. + last_chunk = None # rebind-ok: tracks the most recent chunk for the final flush. async for chunk in response: last_chunk = chunk @@ -309,53 +385,54 @@ class LLMShieldGuardrail(CustomGuardrail): # Nothing to restore in this chunk, but a final chunk still has to # flush whatever the window is holding. if is_final and carry: - body = await self._stream_step("", carry, True, session_id) - carry = body["carry"] - if body["text"] and delta is not None: - delta.content = body["text"] + emitted, carry = await self._stream_step("", carry, True, session_id) + if emitted and delta is not None: + delta.content = emitted yield chunk continue - body = await self._stream_step(text, carry, is_final, session_id) - carry = body["carry"] - delta.content = body["text"] + emitted, carry = await self._stream_step(text, carry, is_final, session_id) + delta.content = emitted yield chunk # A stream that ended without a finish_reason can still leave text held back. if carry and last_chunk is not None: - body = await self._stream_step("", carry, True, session_id) - if body["text"]: - trailing = last_chunk.model_copy(deep=True) - trailing_delta = self._stream_delta(trailing) + flushed: Final = await self._stream_step("", carry, True, session_id) + trailing_text, carry = flushed # rebind-ok: window advances. + if trailing_text: + trailing: Final = last_chunk.model_copy(deep=True) + trailing_delta: Final = self._stream_delta(trailing) if trailing_delta is not None: - trailing_delta.content = body["text"] + trailing_delta.content = trailing_text yield trailing - async def _stream_step(self, text: str, carry: str, final: bool, session_id: str) -> dict: - body = await self._call_shield( + async def _stream_step(self, text: str, carry: str, final: bool, session_id: str) -> tuple[str, str]: + """Returns ``(text safe to emit now, window still being held)``.""" + body: Final = await self._call_shield( _REHYDRATE_STREAM_PATH, session_id, - {"text": text, "carry": carry, "final": final}, + # mutable-ok: JSON request body for httpx. + {"text": text, "carry": carry, "final": final}, # mutable-ok: JSON request body for httpx. ) - emitted = body.get("text") - remaining = body.get("carry") + emitted: Final = body.get("text") + remaining: Final = body.get("carry") if not isinstance(emitted, str) or not isinstance(remaining, str): raise GuardrailRaisedException( guardrail_name=self.guardrail_name, message="LLM Shield stream rehydration returned an unexpected payload.", ) - return {"text": emitted, "carry": remaining} + return emitted, remaining @staticmethod - def _stream_delta(chunk: Any) -> Any: - choices = getattr(chunk, "choices", None) + def _stream_delta(chunk: object) -> Any: + choices: Final = getattr(chunk, "choices", None) if not choices: return None return getattr(choices[0], "delta", None) @staticmethod - def _is_final_chunk(chunk: Any) -> bool: - choices = getattr(chunk, "choices", None) + def _is_final_chunk(chunk: object) -> bool: + choices: Final = getattr(chunk, "choices", None) if not choices: return False return bool(getattr(choices[0], "finish_reason", None)) @@ -366,17 +443,21 @@ class LLMShieldGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: MutableRequest, input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: - texts = inputs.get("texts") + texts: Final = inputs.get("texts") if not texts: return inputs - session_id = self._session_id(request_data) - if input_type == "request": - inputs["texts"] = await self._redact(list(texts), session_id) - else: - inputs["texts"] = await self._rehydrate(list(texts), session_id) - return inputs + session_id: Final = self._session_id(request_data) + replaced: Final = ( + await self._redact(tuple(texts), session_id) + if input_type == "request" + else await self._rehydrate(tuple(texts), session_id) + ) + # 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, "texts": list(replaced)} # mutable-ok: TypedDict. + return merged diff --git a/ruff-strict.toml b/ruff-strict.toml index ae092bdde7d..f9d026c6011 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -30,6 +30,10 @@ external = [ # grows over time; typing it concretely (`object`) broke that forwarding call outright — # basedpyright turned every named param into a reportArgumentType error. Any is correct here. "litellm/proxy/guardrails/guardrail_hooks/alice/alice.py" = ["ANN401"] +# Same reason: `**kwargs` forwards verbatim to CustomGuardrail.__init__, and the lifecycle +# hook signatures inherit `Any` for `response` from CustomLogger, so narrowing them here +# would break the override rather than describe it. +"litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py" = ["ANN401"] [lint.mccabe] max-complexity = 15 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py index c2d50bf5fa7..7a07c9761a8 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py @@ -1,3 +1,4 @@ +from types import SimpleNamespace from unittest.mock import AsyncMock import pytest @@ -133,10 +134,16 @@ class TestRedaction: assert data["messages"][0]["content"][1]["image_url"]["url"] == "http://x/y.png" @pytest.mark.asyncio - async def test_request_without_messages_is_untouched(self): + async def test_request_without_text_is_untouched(self): + """No text to redact means no call to LLM Shield. + + This deliberately uses a request with no caller text at all. An earlier + version used a Responses-API `input`, which asserted the very bypass that + let `input` reach the provider unredacted. + """ guardrail = _guardrail() mock = _mock_post(guardrail) - data = {"input": "no messages here"} + data = {"model": "gpt-4o", "temperature": 0.2} await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") @@ -156,6 +163,91 @@ class TestRedaction: assert len(sessions) == 1 +class TestRequestCoverage: + """Every request shape that carries caller text must be redacted. + + A shape missed here is not a cosmetic gap: the guardrail reports as enabled + while the raw value goes to the provider. + """ + + @pytest.mark.asyncio + async def test_responses_api_string_input_is_redacted(self): + """Measured against a live provider: `input` reached the model unredacted.""" + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["Email [EMAIL_1] the invoice"]}) + + data = {"input": "Email jane.doe@example.com the invoice"} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="aresponses") + + assert mock.call_args_list[0].kwargs["json"]["texts"] == ["Email jane.doe@example.com the invoice"] + assert data["input"] == "Email [EMAIL_1] the invoice" + + @pytest.mark.asyncio + async def test_responses_api_list_input_is_redacted(self): + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["[EMAIL_1]", "[PHONE_1]"]}) + + data = { + "input": [ + {"role": "user", "content": "jane.doe@example.com"}, + {"role": "user", "content": [{"type": "input_text", "text": "555-0100"}]}, + ] + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="aresponses") + + assert data["input"][0]["content"] == "[EMAIL_1]" + assert data["input"][1]["content"][0]["text"] == "[PHONE_1]" + + @pytest.mark.asyncio + async def test_tool_call_arguments_are_redacted(self): + """Tool arguments carry the values the user asked the model to act on.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ['{"email": "[EMAIL_1]"}']}) + + data = { + "messages": [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "send", "arguments": '{"email": "jane.doe@example.com"}'}, + } + ], + } + ] + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert data["messages"][0]["tool_calls"][0]["function"]["arguments"] == '{"email": "[EMAIL_1]"}' + + @pytest.mark.asyncio + async def test_every_shape_in_one_request_is_redacted(self): + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["a", "b", "c", "d"]}) + + data = { + "messages": [ + {"role": "user", "content": "one"}, + {"role": "user", "content": [{"type": "text", "text": "two"}]}, + { + "role": "assistant", + "tool_calls": [{"function": {"name": "f", "arguments": "three"}}], + }, + ], + "input": "four", + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert mock.call_args_list[0].kwargs["json"]["texts"] == ["one", "two", "three", "four"] + assert data["messages"][0]["content"] == "a" + assert data["messages"][1]["content"][0]["text"] == "b" + assert data["messages"][2]["tool_calls"][0]["function"]["arguments"] == "c" + assert data["input"] == "d" + + class TestRestoration: @pytest.mark.asyncio async def test_openai_shape_is_restored(self): @@ -169,6 +261,36 @@ class TestRestoration: assert result.choices[0].message.content == "a@b.com" + @pytest.mark.asyncio + async def test_responses_api_shape_is_restored(self): + """The Responses API reply carries output items, not choices. + + Measured against a live provider: once the request side was fixed the reply + came back still holding the placeholder, because this shape has no choices + to walk. + """ + guardrail = _guardrail(event_hook="post_call") + _mock_post(guardrail, {"texts": ["a@b.com"]}) + + response = SimpleNamespace(output=[SimpleNamespace(content=[{"type": "output_text", "text": "[EMAIL_1]"}])]) + result = await guardrail.async_post_call_success_hook( + data={"messages": []}, user_api_key_dict=None, response=response + ) + + assert result.output[0].content[0]["text"] == "a@b.com" + + @pytest.mark.asyncio + async def test_responses_api_object_blocks_are_restored(self): + """Blocks arrive as objects too, depending on how far the reply is parsed.""" + guardrail = _guardrail(event_hook="post_call") + _mock_post(guardrail, {"texts": ["a@b.com"]}) + + block = SimpleNamespace(text="[EMAIL_1]") + response = SimpleNamespace(output=[SimpleNamespace(content=[block])]) + await guardrail.async_post_call_success_hook(data={"messages": []}, user_api_key_dict=None, response=response) + + assert block.text == "a@b.com" + @pytest.mark.asyncio async def test_anthropic_message_shape_is_restored(self): """The /v1/messages reply is a plain dict with no choices. diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx index 29df7c8bf3d..a02cae097a7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx @@ -73,7 +73,9 @@ interface GuardrailPreset { provider: string; categoryName?: string; guardrailNameSuggestion: string; - mode: string; + // A guardrail that both rewrites the request and repairs the response needs two + // modes seeded, not one; the form already normalises either shape. + mode: string | string[]; defaultOn: boolean; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts index b445cc9c5ad..3579457bcd5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts @@ -2,7 +2,9 @@ export interface GuardrailPreset { provider: string; categoryName?: string; guardrailNameSuggestion: string; - mode: string; + // A guardrail that both rewrites the request and repairs the response needs two + // modes seeded, not one; the form already normalises either shape. + mode: string | string[]; defaultOn: boolean; } @@ -321,7 +323,9 @@ export const GUARDRAIL_PRESETS: Record = { llm_shield: { provider: "LLM Shield", guardrailNameSuggestion: "LLM Shield", - mode: "pre_call", + // Both halves are required. With only pre_call the request is redacted and the + // placeholders are handed straight back to the caller. + mode: ["pre_call", "post_call"], defaultOn: false, }, }; From 295ad527d50540a3eb66cb5a29933ef4cac07102 Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Thu, 3 Sep 2026 19:13:31 -0500 Subject: [PATCH 08/22] fix(guardrails): narrow the stream delta before writing to it basedpyright could not prove the delta was non-None on the write path, and reportOptionalMemberAccess has a zero budget. The guard is also clearer than relying on the text check to imply it. --- .../guardrails/guardrail_hooks/llm_shield/llm_shield.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py index fe8b70c08b8..c126bf0de37 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py @@ -381,12 +381,12 @@ class LLMShieldGuardrail(CustomGuardrail): text = getattr(delta, "content", None) if delta is not None else None is_final = self._is_final_chunk(chunk) - if not isinstance(text, str) or not text: + if delta is None or not isinstance(text, str) or not text: # Nothing to restore in this chunk, but a final chunk still has to # flush whatever the window is holding. - if is_final and carry: + if is_final and carry and delta is not None: emitted, carry = await self._stream_step("", carry, True, session_id) - if emitted and delta is not None: + if emitted: delta.content = emitted yield chunk continue From 47421b7541c6684db3f91c8ba627ceed15a40468 Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Thu, 3 Sep 2026 19:24:12 -0500 Subject: [PATCH 09/22] fix(guardrails): mint the vault id instead of trusting the caller's The vault id was taken from caller-supplied session metadata, and every caller shares one LLM Shield key. Someone who knew or guessed another caller's session id could send a placeholder, have the model echo it back, and get that caller's plaintext restored into their own reply. Vault ids are now minted per request behind a per-process prefix, so a caller cannot name a vault this process uses. Redaction mints, restoration reads back, and a reply whose id does not match is left holding its placeholders rather than resolved against some other vault. Also covers two more request fields that were reaching the provider intact: the Responses API `instructions`, and the legacy `function_call.arguments` alongside `tool_calls`. The collectors move to module level, which drops the traversal back under the complexity limit and lets the code carry its own explanation instead of the comments that were restating it. --- .../guardrail_hooks/llm_shield/llm_shield.py | 141 +++++++++++------- .../guardrail_hooks/test_llm_shield.py | 82 ++++++++++ 2 files changed, 169 insertions(+), 54 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py index c126bf0de37..26c6774f48d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py @@ -24,7 +24,6 @@ from litellm._logging import verbose_proxy_logger from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_guardrail import ( CustomGuardrail, - get_session_id_from_request_data, log_guardrail_information, ) from litellm.llms.custom_httpx.http_handler import ( @@ -52,6 +51,13 @@ _REHYDRATE_STREAM_PATH: Final = "/v1/guard/rehydrate/stream" # across concurrent requests. _SESSION_METADATA_KEY: Final = "llm_shield_session_id" +# 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 # The proxy's own request dict. Mutable by design: a pre-call guardrail rewrites @@ -67,6 +73,51 @@ JsonBody: TypeAlias = dict # replacement back where it came from. _Slot: TypeAlias = tuple[str, Callable[[str], None]] # mutable-ok: Callable's param list. +# 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] # mutable-ok: accumulator passed between collectors. + + +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, lambda new, c=container, k=key: c.__setitem__(k, new))) + + +def _collect_content(container: MutableRequest, slots: _SlotSink) -> None: + """`content` is either a string or the multimodal list of typed parts.""" + content: Final = container.get("content") + if isinstance(content, str): + _collect(container, "content", slots) + return + for part in content if isinstance(content, list) else (): + if isinstance(part, dict): + _collect(part, "text", 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 message.get("tool_calls") or (): + function: Final = tool_call.get("function") if isinstance(tool_call, dict) else None + if isinstance(function, dict): + _collect(function, "arguments", slots) + legacy: Final = message.get("function_call") + if isinstance(legacy, dict): + _collect(legacy, "arguments", slots) + + +def _collect_responses_fields(data: MutableRequest, slots: _SlotSink) -> None: + """The Responses API sends text outside `messages`, in `instructions` and `input`.""" + _collect(data, "instructions", slots) + request_input: Final = data.get("input") + if isinstance(request_input, str): + _collect(data, "input", slots) + return + for item in request_input if isinstance(request_input, list) else (): + if isinstance(item, dict): + _collect_content(item, slots) + class LLMShieldGuardrail(CustomGuardrail): """Redacts PII before it leaves the proxy and restores it in the response. @@ -164,67 +215,50 @@ class LLMShieldGuardrail(CustomGuardrail): # --- session ------------------------------------------------------------------ - def _session_id(self, data: MutableRequest) -> str: - """Returns a session id stable across this request's hooks.""" + @staticmethod + def _mint_session_id(data: MutableRequest) -> str: + """Mints a vault id for this request, overwriting anything already there. + + Redaction and restoration both happen inside one request/response pair, so + a fresh id per request is all that is needed, and it is what keeps one + caller from reaching another caller's vault. + """ + session_id: Final = f"{_VAULT_PREFIX}-{uuid.uuid4().hex}" metadata: Final = data.setdefault("metadata", {}) # mutable-ok: per-request store. - if not isinstance(metadata, dict): - return f"litellm-{uuid.uuid4().hex}" - existing: Final = metadata.get(_SESSION_METADATA_KEY) - if isinstance(existing, str) and existing: - return existing - session_id: Final = get_session_id_from_request_data(data) or f"litellm-{uuid.uuid4().hex}" - metadata[_SESSION_METADATA_KEY] = session_id + if isinstance(metadata, dict): + metadata[_SESSION_METADATA_KEY] = session_id return session_id + @staticmethod + def _session_id(data: MutableRequest) -> str: + """Reads back the vault id minted while redacting this request. + + Falls back to an unused id rather than to anything the caller supplied: a + reply that cannot be restored is a visible placeholder, while trusting a + caller-supplied id would hand them someone else's plaintext. + """ + metadata: Final = data.get("metadata") + existing: Final = metadata.get(_SESSION_METADATA_KEY) if isinstance(metadata, dict) else None + if isinstance(existing, str) and existing.startswith(_VAULT_PREFIX): + return existing + return f"{_VAULT_PREFIX}-{uuid.uuid4().hex}" + # --- request traversal -------------------------------------------------------- @staticmethod def _locate_request_texts(data: MutableRequest) -> Sequence[_Slot]: """Finds every redactable span in an outbound request. - Returns ``(text, write)`` pairs. Any shape missed here reaches the provider - in the clear, so this walks all of the request shapes that carry caller text: - - - chat ``messages``, both string and multimodal list ``content`` - - tool call ``arguments``, which routinely carry the values a user asked - the model to look up - - the Responses API ``input``, as a bare string or a list of items + 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 caller text. """ - slots: Final[list[_Slot]] = [] # mutable-ok: accumulator, frozen on return. - - def add(container: MutableRequest, key: str, value: object) -> None: - if isinstance(value, str) and value: - slots.append((value, lambda new, c=container, k=key: c.__setitem__(k, new))) - - def add_content(container: MutableRequest) -> None: - """Adds `content`, which is either a string or a list of typed parts.""" - content: Final = container.get("content") - if isinstance(content, str): - add(container, "content", content) - return - for part in content if isinstance(content, list) else (): - if isinstance(part, dict): - add(part, "text", part.get("text")) - - def add_tool_calls(message: MutableRequest) -> None: - for tool_call in message.get("tool_calls") or (): - function = tool_call.get("function") if isinstance(tool_call, dict) else None - if isinstance(function, dict): - add(function, "arguments", function.get("arguments")) - + slots: Final[_SlotSink] = [] # mutable-ok: accumulator, frozen on return. for message in data.get("messages") or (): if isinstance(message, dict): - add_content(message) - add_tool_calls(message) - - request_input: Final = data.get("input") - if isinstance(request_input, str): - add(data, "input", request_input) - else: - for item in request_input if isinstance(request_input, list) else (): - if isinstance(item, dict): - add_content(item) - + _collect_content(message, slots) + _collect_tool_arguments(message, slots) + _collect_responses_fields(data, slots) return tuple(slots) # --- hooks -------------------------------------------------------------------- @@ -245,7 +279,7 @@ class LLMShieldGuardrail(CustomGuardrail): if not slots: return data - redacted: Final = await self._redact(tuple(text for text, _ in slots), self._session_id(data)) + redacted: Final = await self._redact(tuple(text for text, _ in slots), self._mint_session_id(data)) for (_, write), replacement in zip(slots, redacted): write(replacement) return data @@ -451,11 +485,10 @@ class LLMShieldGuardrail(CustomGuardrail): if not texts: return inputs - session_id: Final = self._session_id(request_data) replaced: Final = ( - await self._redact(tuple(texts), session_id) + await self._redact(tuple(texts), self._mint_session_id(request_data)) if input_type == "request" - else await self._rehydrate(tuple(texts), session_id) + else await self._rehydrate(tuple(texts), self._session_id(request_data)) ) # Return a new mapping rather than rewriting the caller's, so this stays a # pure transform of the inputs it was handed. diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py index 7a07c9761a8..45d45301858 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py @@ -223,6 +223,35 @@ class TestRequestCoverage: assert data["messages"][0]["tool_calls"][0]["function"]["arguments"] == '{"email": "[EMAIL_1]"}' + @pytest.mark.asyncio + async def test_responses_api_instructions_are_redacted(self): + """`instructions` is provider-bound text that sits outside `messages`.""" + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["contact [EMAIL_1]"]}) + + data = {"instructions": "contact jane.doe@example.com", "input": ""} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="aresponses") + + assert mock.call_args_list[0].kwargs["json"]["texts"] == ["contact jane.doe@example.com"] + assert data["instructions"] == "contact [EMAIL_1]" + + @pytest.mark.asyncio + async def test_legacy_function_call_arguments_are_redacted(self): + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ['{"email": "[EMAIL_1]"}']}) + + data = { + "messages": [ + { + "role": "assistant", + "function_call": {"name": "send", "arguments": '{"email": "jane.doe@example.com"}'}, + } + ] + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert data["messages"][0]["function_call"]["arguments"] == '{"email": "[EMAIL_1]"}' + @pytest.mark.asyncio async def test_every_shape_in_one_request_is_redacted(self): guardrail = _guardrail() @@ -334,6 +363,59 @@ class TestRestoration: assert result["content"][1] == {"type": "tool_use", "id": "t1", "name": "lookup", "input": {}} +class TestVaultIsolation: + """The vault id must never be something a caller can choose. + + The vault holds the plaintext behind every placeholder. If a caller could name + the vault, they could send a placeholder, have the model echo it back, and get + another caller's value restored into their own reply. + """ + + @pytest.mark.asyncio + async def test_caller_supplied_session_id_is_not_used(self): + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}) + + data = { + "messages": [{"role": "user", "content": "a@b.com"}], + "metadata": {"llm_shield_session_id": "victim-session"}, + "litellm_session_id": "victim-session", + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + used = mock.call_args_list[0].kwargs["headers"]["X-Session-ID"] + assert used != "victim-session" + assert data["metadata"]["llm_shield_session_id"] == used + + @pytest.mark.asyncio + async def test_restore_ignores_a_foreign_session_id(self): + """A reply is left unrestored rather than resolved against another vault.""" + guardrail = _guardrail(event_hook="post_call") + mock = _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}) + + data = {"metadata": {"llm_shield_session_id": "victim-session"}} + response = ModelResponse(choices=[Choices(index=0, message=Message(role="assistant", content="[EMAIL_1]"))]) + await guardrail.async_post_call_success_hook(data=data, user_api_key_dict=None, response=response) + + assert mock.call_args_list[0].kwargs["headers"]["X-Session-ID"] != "victim-session" + + @pytest.mark.asyncio + async def test_each_request_gets_its_own_vault(self): + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}, {"texts": ["[EMAIL_1]"]}) + + for _ in range(2): + await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data={"messages": [{"role": "user", "content": "a@b.com"}]}, + call_type="completion", + ) + + seen = {call.kwargs["headers"]["X-Session-ID"] for call in mock.call_args_list} + assert len(seen) == 2 + + class TestFailClosed: @pytest.mark.asyncio async def test_unreachable_shield_blocks_the_request(self): From f3eb108f86a5b02063f35a6d2f50d69ceb57cf91 Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Thu, 3 Sep 2026 19:41:46 -0500 Subject: [PATCH 10/22] fix(guardrails): drop Final from a loop-assigned local basedpyright rejects a Final assigned inside a loop, and reportGeneralTypeIssues sits one over its budget ceiling. --- .../proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py index 26c6774f48d..172de67020b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py @@ -99,7 +99,7 @@ def _collect_content(container: MutableRequest, slots: _SlotSink) -> None: 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 message.get("tool_calls") or (): - function: Final = tool_call.get("function") if isinstance(tool_call, dict) else None + function = tool_call.get("function") if isinstance(tool_call, dict) else None # rebind-ok: loop variable. if isinstance(function, dict): _collect(function, "arguments", slots) legacy: Final = message.get("function_call") From 8d0b1881d7b028e1eb6b09c1bb76f1527923768f Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Thu, 3 Sep 2026 19:59:36 -0500 Subject: [PATCH 11/22] fix(guardrails): redact completion prompts and responses tool items Two more provider-bound request shapes were reaching the model intact while the guardrail reported as enabled. /v1/completions carries its text in a top-level `prompt`, which the traversal never looked at. It is handled as a string and as the array form, where each entry is rewritten in place. Responses input items hold tool data outside `content`: a function_call item in `arguments`, a function_call_output item in `output`. Both are now collected alongside the item's content. Adds a test per shape. --- .../guardrail_hooks/llm_shield/llm_shield.py | 30 +++++++++++++- .../guardrail_hooks/test_llm_shield.py | 40 +++++++++++++++++++ 2 files changed, 68 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py index 172de67020b..8989fabc321 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py @@ -77,6 +77,10 @@ _Slot: TypeAlias = tuple[str, Callable[[str], None]] # mutable-ok: Callable's p # _locate_request_texts, which freezes it into a tuple before returning. _SlotSink: TypeAlias = list[_Slot] # mutable-ok: accumulator passed between collectors. +# A caller-owned list whose entries are rewritten in place, such as a Completions +# `prompt` sent as an array of strings. +MutableSeq: TypeAlias = list # mutable-ok: the request payload's own list. + def _collect(container: MutableRequest, key: str, slots: _SlotSink) -> None: """Records the string at `key`, along with the write that replaces it.""" @@ -85,6 +89,23 @@ def _collect(container: MutableRequest, key: str, slots: _SlotSink) -> None: slots.append((value, lambda new, c=container, k=key: c.__setitem__(k, new))) +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, lambda new, e=entries, i=index: e.__setitem__(i, new))) + + +def _collect_prompt(data: MutableRequest, slots: _SlotSink) -> None: + """The Completions API sends its text in a top-level `prompt`.""" + prompt: Final = data.get("prompt") + if isinstance(prompt, str): + _collect(data, "prompt", slots) + return + for index in range(len(prompt)) if isinstance(prompt, list) else (): + _collect_entry(prompt, index, slots) + + def _collect_content(container: MutableRequest, slots: _SlotSink) -> None: """`content` is either a string or the multimodal list of typed parts.""" content: Final = container.get("content") @@ -115,8 +136,12 @@ def _collect_responses_fields(data: MutableRequest, slots: _SlotSink) -> None: _collect(data, "input", slots) return for item in request_input if isinstance(request_input, list) else (): - if isinstance(item, dict): - _collect_content(item, slots) + if not isinstance(item, dict): + continue + _collect_content(item, slots) + # A function_call item holds `arguments`; a function_call_output holds `output`. + _collect(item, "arguments", slots) + _collect(item, "output", slots) class LLMShieldGuardrail(CustomGuardrail): @@ -259,6 +284,7 @@ class LLMShieldGuardrail(CustomGuardrail): _collect_content(message, slots) _collect_tool_arguments(message, slots) _collect_responses_fields(data, slots) + _collect_prompt(data, slots) return tuple(slots) # --- hooks -------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py index 45d45301858..ffd6ede28b6 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py @@ -252,6 +252,46 @@ class TestRequestCoverage: assert data["messages"][0]["function_call"]["arguments"] == '{"email": "[EMAIL_1]"}' + @pytest.mark.asyncio + async def test_completions_prompt_is_redacted(self): + """/v1/completions puts its text in a top-level `prompt`, not in messages.""" + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["Email [EMAIL_1]"]}) + + data = {"prompt": "Email jane.doe@example.com"} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="atext_completion") + + assert mock.call_args_list[0].kwargs["json"]["texts"] == ["Email jane.doe@example.com"] + assert data["prompt"] == "Email [EMAIL_1]" + + @pytest.mark.asyncio + async def test_completions_prompt_array_is_redacted(self): + """`prompt` also accepts an array, and each entry is provider-bound.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["[EMAIL_1]", "[PHONE_1]"]}) + + data = {"prompt": ["jane.doe@example.com", "555-0100"]} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="atext_completion") + + assert data["prompt"] == ["[EMAIL_1]", "[PHONE_1]"] + + @pytest.mark.asyncio + async def test_responses_function_call_items_are_redacted(self): + """Responses input items hold tool data in `arguments` and `output`.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ['{"email": "[EMAIL_1]"}', "sent to [EMAIL_1]"]}) + + data = { + "input": [ + {"type": "function_call", "name": "send", "arguments": '{"email": "jane.doe@example.com"}'}, + {"type": "function_call_output", "call_id": "c1", "output": "sent to jane.doe@example.com"}, + ] + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="aresponses") + + assert data["input"][0]["arguments"] == '{"email": "[EMAIL_1]"}' + assert data["input"][1]["output"] == "sent to [EMAIL_1]" + @pytest.mark.asyncio async def test_every_shape_in_one_request_is_redacted(self): guardrail = _guardrail() From 46f13807a513801e107592f1eec79ea24c47f1a6 Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Thu, 3 Sep 2026 20:13:45 -0500 Subject: [PATCH 12/22] fix(guardrails): redact the anthropic system prompt and string-array input Two more provider-bound shapes, found by walking the request types rather than waiting for them to be reported. /v1/messages carries its system prompt at the top level, as a string or a list of text blocks. It is one of the endpoints this guardrail claims to cover, and a system prompt is a natural place to put a customer's details. `input` as an array of bare strings, the embeddings and moderations shape, was skipped because the loop only handled item dicts. Verified against a live provider: a system prompt holding an address now reaches the model as a stand-in and is restored in the reply. --- .../guardrail_hooks/llm_shield/llm_shield.py | 18 ++++++++- .../guardrail_hooks/test_llm_shield.py | 38 +++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py index 8989fabc321..e495f2e99d4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py @@ -128,6 +128,17 @@ def _collect_tool_arguments(message: MutableRequest, slots: _SlotSink) -> 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 + for part in system if isinstance(system, list) else (): + if isinstance(part, dict): + _collect(part, "text", slots) + + def _collect_responses_fields(data: MutableRequest, slots: _SlotSink) -> None: """The Responses API sends text outside `messages`, in `instructions` and `input`.""" _collect(data, "instructions", slots) @@ -135,7 +146,11 @@ def _collect_responses_fields(data: MutableRequest, slots: _SlotSink) -> None: if isinstance(request_input, str): _collect(data, "input", slots) return - for item in request_input if isinstance(request_input, list) else (): + for index, item in enumerate(request_input if isinstance(request_input, list) else ()): + if isinstance(item, str): + # The embeddings and moderations shape: `input` as an array of strings. + _collect_entry(request_input, index, slots) + continue if not isinstance(item, dict): continue _collect_content(item, slots) @@ -285,6 +300,7 @@ class LLMShieldGuardrail(CustomGuardrail): _collect_tool_arguments(message, slots) _collect_responses_fields(data, slots) _collect_prompt(data, slots) + _collect_system(data, slots) return tuple(slots) # --- hooks -------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py index ffd6ede28b6..f16584fbd1d 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py @@ -292,6 +292,44 @@ class TestRequestCoverage: assert data["input"][0]["arguments"] == '{"email": "[EMAIL_1]"}' assert data["input"][1]["output"] == "sent to [EMAIL_1]" + @pytest.mark.asyncio + async def test_anthropic_system_prompt_is_redacted(self): + """/v1/messages carries its system prompt at the top level, not in messages.""" + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["the user is [EMAIL_1]"]}) + + data = {"system": "the user is jane.doe@example.com", "messages": []} + await guardrail.async_pre_call_hook( + user_api_key_dict=None, cache=None, data=data, call_type="anthropic_messages" + ) + + assert mock.call_args_list[0].kwargs["json"]["texts"] == ["the user is jane.doe@example.com"] + assert data["system"] == "the user is [EMAIL_1]" + + @pytest.mark.asyncio + async def test_anthropic_system_blocks_are_redacted(self): + """`system` also accepts a list of text blocks.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}) + + data = {"system": [{"type": "text", "text": "jane.doe@example.com"}], "messages": []} + await guardrail.async_pre_call_hook( + user_api_key_dict=None, cache=None, data=data, call_type="anthropic_messages" + ) + + assert data["system"][0]["text"] == "[EMAIL_1]" + + @pytest.mark.asyncio + async def test_string_array_input_is_redacted(self): + """Embeddings and moderations send `input` as an array of bare strings.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["[EMAIL_1]", "[PHONE_1]"]}) + + data = {"input": ["jane.doe@example.com", "555-0100"]} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="aembedding") + + assert data["input"] == ["[EMAIL_1]", "[PHONE_1]"] + @pytest.mark.asyncio async def test_every_shape_in_one_request_is_redacted(self): guardrail = _guardrail() From a35c5028196c07880ae7000f5c99dcf40d069c59 Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Thu, 3 Sep 2026 20:26:57 -0500 Subject: [PATCH 13/22] fix(guardrails): narrow prompt and input to a list before iterating Guarding with a conditional iterable left the value un-narrowed, so passing it on was an argument-type error and the element checks read as unreachable. An early return narrows it properly and reads better. --- .../guardrails/guardrail_hooks/llm_shield/llm_shield.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py index e495f2e99d4..39c0930d7cf 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py @@ -102,7 +102,9 @@ def _collect_prompt(data: MutableRequest, slots: _SlotSink) -> None: if isinstance(prompt, str): _collect(data, "prompt", slots) return - for index in range(len(prompt)) if isinstance(prompt, list) else (): + if not isinstance(prompt, list): + return + for index in range(len(prompt)): _collect_entry(prompt, index, slots) @@ -146,7 +148,9 @@ def _collect_responses_fields(data: MutableRequest, slots: _SlotSink) -> None: if isinstance(request_input, str): _collect(data, "input", slots) return - for index, item in enumerate(request_input if isinstance(request_input, list) else ()): + if not isinstance(request_input, list): + return + for index, item in enumerate(request_input): if isinstance(item, str): # The embeddings and moderations shape: `input` as an array of strings. _collect_entry(request_input, index, slots) From 46438d7cf74913f8f6694f5f1adf1f48585d2d83 Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Thu, 3 Sep 2026 20:48:50 -0500 Subject: [PATCH 14/22] fix(guardrails): restore every streaming choice, not just the first Streaming rehydration read and rewrote choices[0] only, so with n>1 every later choice went back to the caller still holding its placeholders. Each choice is its own token stream, so the sliding window is now tracked per choice index rather than once per stream. A single shared window would have been worse than the bug: it would splice the characters held back for one choice onto the next one's delta. The final flush walks every choice the same way, and the two helpers that only ever looked at choices[0] are gone. Adds a test that both choices come back restored, and one that each choice gets its own window handed back rather than its neighbour's. --- .../guardrail_hooks/llm_shield/llm_shield.py | 105 ++++++++++-------- .../guardrail_hooks/test_llm_shield.py | 67 +++++++++++ 2 files changed, 128 insertions(+), 44 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py index 39c0930d7cf..f84f0f429aa 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py @@ -77,6 +77,9 @@ _Slot: TypeAlias = tuple[str, Callable[[str], None]] # mutable-ok: Callable's p # _locate_request_texts, which freezes it into a tuple before returning. _SlotSink: TypeAlias = list[_Slot] # mutable-ok: accumulator passed between collectors. +# Sliding windows keyed by streaming choice index, threaded through one stream. +_CarryWindows: TypeAlias = dict # mutable-ok: per-choice windows advanced in place. + # A caller-owned list whose entries are rewritten in place, such as a Completions # `prompt` sent as an array of strings. MutableSeq: TypeAlias = list # mutable-ok: the request payload's own list. @@ -163,6 +166,12 @@ def _collect_responses_fields(data: MutableRequest, slots: _SlotSink) -> None: _collect(item, "output", 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 + + class LLMShieldGuardrail(CustomGuardrail): """Redacts PII before it leaves the proxy and restores it in the response. @@ -441,10 +450,10 @@ class LLMShieldGuardrail(CustomGuardrail): ) -> AsyncGenerator[Any, None]: """Restores original values incrementally, without buffering the stream. - The carry-over window is a local of this generator, so it is scoped to one - stream and cannot leak between concurrent requests. LLM Shield returns the - text that is safe to emit now plus the trailing characters it is still - holding, which are sent back with the next delta. + Each choice is its own token stream, so the sliding window is tracked per + choice index. One shared window would splice the characters held back for + one choice onto the next. The windows are locals of this generator, so they + are scoped to a single stream and cannot leak between concurrent requests. """ if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True: async for chunk in response: @@ -452,39 +461,61 @@ class LLMShieldGuardrail(CustomGuardrail): return session_id: Final = self._session_id(request_data) - carry = "" # rebind-ok: the sliding window advances with every delta. + carries: Final[dict] = {} # mutable-ok: per-choice windows, local to this stream. last_chunk = None # rebind-ok: tracks the most recent chunk for the final flush. async for chunk in response: last_chunk = chunk - delta = self._stream_delta(chunk) - text = getattr(delta, "content", None) if delta is not None else None - is_final = self._is_final_chunk(chunk) - - if delta is None or not isinstance(text, str) or not text: - # Nothing to restore in this chunk, but a final chunk still has to - # flush whatever the window is holding. - if is_final and carry and delta is not None: - emitted, carry = await self._stream_step("", carry, True, session_id) - if emitted: - delta.content = emitted - yield chunk - continue - - emitted, carry = await self._stream_step(text, carry, is_final, session_id) - delta.content = emitted + for choice in getattr(chunk, "choices", None) or (): + await self._restore_choice(choice, carries, session_id) yield chunk # A stream that ended without a finish_reason can still leave text held back. - if carry and last_chunk is not None: - flushed: Final = await self._stream_step("", carry, True, session_id) - trailing_text, carry = flushed # rebind-ok: window advances. - if trailing_text: - trailing: Final = last_chunk.model_copy(deep=True) - trailing_delta: Final = self._stream_delta(trailing) - if trailing_delta is not None: - trailing_delta.content = trailing_text - yield trailing + if last_chunk is not None and any(carries.values()): + trailing: Final = last_chunk.model_copy(deep=True) + if await self._flush_trailing(trailing, carries, session_id): + yield trailing + + async def _restore_choice(self, choice: Any, carries: _CarryWindows, session_id: str) -> None: + """Restores one choice's delta, advancing that choice's own window.""" + delta: Final = getattr(choice, "delta", None) + if delta is None: + return + index: Final = _choice_index(choice) + carry: Final = carries.get(index, "") + text: Final = getattr(delta, "content", None) + is_final: Final = bool(getattr(choice, "finish_reason", 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, flushed_carry = await self._stream_step("", carry, True, session_id) + carries[index] = flushed_carry # rebind-ok: this choice's window advances. + if flushed: + delta.content = flushed + return + + emitted, remaining = await self._stream_step(text, carry, is_final, session_id) + carries[index] = remaining # rebind-ok: this choice's window advances. + delta.content = emitted + + async def _flush_trailing(self, trailing: Any, carries: _CarryWindows, session_id: str) -> bool: + """Empties every still-held window into a copy of the last chunk.""" + emitted_any = False # rebind-ok: set once any choice contributes text. + for choice in getattr(trailing, "choices", None) or (): + delta = getattr(choice, "delta", None) + if delta is None: + continue + index = _choice_index(choice) + carry = carries.get(index, "") + if not carry: + delta.content = None + continue + text, remaining = await self._stream_step("", carry, True, session_id) + carries[index] = remaining # rebind-ok: this choice's window advances. + delta.content = text or None + emitted_any = emitted_any or bool(text) + return emitted_any async def _stream_step(self, text: str, carry: str, final: bool, session_id: str) -> tuple[str, str]: """Returns ``(text safe to emit now, window still being held)``.""" @@ -503,20 +534,6 @@ class LLMShieldGuardrail(CustomGuardrail): ) return emitted, remaining - @staticmethod - def _stream_delta(chunk: object) -> Any: - choices: Final = getattr(chunk, "choices", None) - if not choices: - return None - return getattr(choices[0], "delta", None) - - @staticmethod - def _is_final_chunk(chunk: object) -> bool: - choices: Final = getattr(chunk, "choices", None) - if not choices: - return False - return bool(getattr(choices[0], "finish_reason", None)) - # --- unified API (powers the UI "Test guardrail" button) ----------------------- @log_guardrail_information diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py index f16584fbd1d..e9cc1a36b07 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py @@ -587,6 +587,73 @@ class TestStreamingRehydration: assert mock.call_args_list[1].kwargs["json"]["carry"] == "hold" assert mock.call_args_list[1].kwargs["json"]["final"] is True + @pytest.mark.asyncio + async def test_every_choice_is_restored(self): + """With n>1 a later choice must not be handed back still holding a placeholder.""" + guardrail = _guardrail(event_hook="post_call") + _mock_post( + guardrail, + {"text": "first@example.com", "carry": ""}, + {"text": "second@example.com", "carry": ""}, + ) + + async def stream(): + yield ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content="[EMAIL_1]"), finish_reason="stop"), + StreamingChoices(index=1, delta=Delta(content="[EMAIL_2]"), finish_reason="stop"), + ] + ) + + chunks = await _drain( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, response=stream(), request_data={"messages": []} + ) + ) + + restored = [choice.delta.content for choice in chunks[0].choices] + assert restored == ["first@example.com", "second@example.com"] + + @pytest.mark.asyncio + async def test_choice_windows_do_not_cross_contaminate(self): + """Each choice is its own token stream, so each carries its own window. + + One shared window would send the characters held back for choice 0 up + against choice 1's next delta and splice the two streams together. + """ + guardrail = _guardrail(event_hook="post_call") + mock = _mock_post( + guardrail, + {"text": "", "carry": "A-held"}, + {"text": "", "carry": "B-held"}, + {"text": "a-done", "carry": ""}, + {"text": "b-done", "carry": ""}, + ) + + async def stream(): + yield ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content="a1")), + StreamingChoices(index=1, delta=Delta(content="b1")), + ] + ) + yield ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content="a2"), finish_reason="stop"), + StreamingChoices(index=1, delta=Delta(content="b2"), finish_reason="stop"), + ] + ) + + await _drain( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, response=stream(), request_data={"messages": []} + ) + ) + + sent = [call.kwargs["json"] for call in mock.call_args_list] + assert sent[2]["carry"] == "A-held", "choice 0 must get its own window back" + assert sent[3]["carry"] == "B-held", "choice 1 must get its own window back" + @pytest.mark.asyncio async def test_chunks_are_forwarded_as_they_arrive(self): """Restoration must not buffer the stream into a single terminal chunk.""" From a0abb9a4992489b91deea3cf69ab3b354314c9fb Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Thu, 3 Sep 2026 21:09:48 -0500 Subject: [PATCH 15/22] refactor(guardrails): name the guardrail llm_shield_proxy throughout The integration was called llm_shield in code, llm-shield in the example config, and LLM Shield in the dashboard, while the product and its PyPI package are both llm-shield-proxy. An operator who saw the guardrail in LiteLLM could not tell what to install. One identifier now: llm_shield_proxy for the enum value, module, directory, class, config model, logo and environment variables, with LLM Shield Proxy as the display name. That matches `pip install llm-shield-proxy`. Renames only; no behaviour change. --- .../__init__.py | 10 +++---- .../example_config.yaml | 24 ++++++++-------- .../llm_shield_proxy.py} | 25 +++++++++-------- litellm/types/guardrails.py | 2 +- .../{llm_shield.py => llm_shield_proxy.py} | 10 +++---- ruff-strict.toml | 2 +- ...llm_shield.py => test_llm_shield_proxy.py} | 28 +++++++++---------- .../{llm_shield.svg => llm_shield_proxy.svg} | 0 .../_components/guardrail_garden_configs.ts | 6 ++-- .../_components/guardrail_garden_data.test.ts | 2 +- .../_components/guardrail_garden_data.ts | 8 +++--- .../_components/guardrail_info_helpers.tsx | 6 ++-- 12 files changed, 62 insertions(+), 61 deletions(-) rename litellm/proxy/guardrails/guardrail_hooks/{llm_shield => llm_shield_proxy}/__init__.py (72%) rename litellm/proxy/guardrails/guardrail_hooks/{llm_shield => llm_shield_proxy}/example_config.yaml (70%) rename litellm/proxy/guardrails/guardrail_hooks/{llm_shield/llm_shield.py => llm_shield_proxy/llm_shield_proxy.py} (95%) rename litellm/types/proxy/guardrails/guardrail_hooks/{llm_shield.py => llm_shield_proxy.py} (51%) rename tests/test_litellm/proxy/guardrails/guardrail_hooks/{test_llm_shield.py => test_llm_shield_proxy.py} (96%) rename ui/litellm-dashboard/public/assets/logos/{llm_shield.svg => llm_shield_proxy.svg} (100%) diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/__init__.py similarity index 72% rename from litellm/proxy/guardrails/guardrail_hooks/llm_shield/__init__.py rename to litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/__init__.py index a6cc54d5408..c8ca68a8967 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/__init__.py @@ -2,16 +2,16 @@ from typing import TYPE_CHECKING, Final from litellm.types.guardrails import SupportedGuardrailIntegrations -from .llm_shield import LLMShieldGuardrail +from .llm_shield_proxy import LLMShieldProxyGuardrail if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams -def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> LLMShieldProxyGuardrail: import litellm - _llm_shield_guardrail_callback: Final = LLMShieldGuardrail( + _llm_shield_guardrail_callback: Final = LLMShieldProxyGuardrail( api_key=litellm_params.api_key, api_base=litellm_params.api_base, guardrail_name=guardrail.get("guardrail_name", ""), @@ -24,10 +24,10 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" guardrail_initializer_registry: Final = { # mutable-ok: module-level registry, built once and never mutated - SupportedGuardrailIntegrations.LLM_SHIELD.value: initialize_guardrail, + SupportedGuardrailIntegrations.LLM_SHIELD_PROXY.value: initialize_guardrail, } guardrail_class_registry: Final = { # mutable-ok: module-level registry, built once and never mutated - SupportedGuardrailIntegrations.LLM_SHIELD.value: LLMShieldGuardrail, + SupportedGuardrailIntegrations.LLM_SHIELD_PROXY.value: LLMShieldProxyGuardrail, } diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/example_config.yaml b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/example_config.yaml similarity index 70% rename from litellm/proxy/guardrails/guardrail_hooks/llm_shield/example_config.yaml rename to litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/example_config.yaml index aa63fa9d252..b5c9f0f8b69 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/example_config.yaml +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/example_config.yaml @@ -1,7 +1,7 @@ -# Example LiteLLM Proxy configuration for LLM Shield -# LLM Shield is a self-hosted PII gateway: https://github.com/ninadphalak/LLM-Shield-Proxy +# Example LiteLLM Proxy configuration for LLM Shield Proxy +# LLM Shield Proxy is a self-hosted PII gateway: https://github.com/ninadphalak/LLM-Shield-Proxy # -# Unlike a masking guardrail, LLM Shield's substitution is reversible. Personal data is +# Unlike a masking guardrail, LLM Shield Proxy's substitution is reversible. Personal data is # replaced with placeholders before the request goes to the provider, and the original # values are put back into the model's reply, so the end user still sees real data while # the provider never received it. @@ -15,25 +15,25 @@ model_list: guardrails: # Both modes belong on ONE entry. pre_call redacts the outbound request and post_call # restores the reply; listing only pre_call would send placeholders back to the user. - - guardrail_name: "llm-shield" + - guardrail_name: "llm_shield_proxy" litellm_params: - guardrail: llm_shield + guardrail: llm_shield_proxy mode: ["pre_call", "post_call"] default_on: true - # Your own LLM Shield deployment. Defaults to http://localhost:8000, and also reads - # LLM_SHIELD_API_BASE from the environment. + # Your own LLM Shield Proxy deployment. Defaults to http://localhost:8000, and also reads + # LLM_SHIELD_PROXY_API_BASE from the environment. api_base: "http://localhost:8000" - # A virtual key configured on that deployment. Also reads LLM_SHIELD_API_KEY. - api_key: os.environ/LLM_SHIELD_API_KEY + # A virtual key configured on that deployment. Also reads LLM_SHIELD_PROXY_API_KEY. + api_key: os.environ/LLM_SHIELD_PROXY_API_KEY # Usage: # -# 1. Run LLM Shield somewhere the proxy can reach: +# 1. Run LLM Shield Proxy somewhere the proxy can reach: # pip install llm-shield-proxy # llm-shield-proxy --port 8000 # # 2. Point this config at it and start the proxy: -# export LLM_SHIELD_API_KEY="your-virtual-key" +# export LLM_SHIELD_PROXY_API_KEY="your-virtual-key" # litellm --config example_config.yaml # # 3. Send a request containing personal data: @@ -47,7 +47,7 @@ guardrails: # # Notes: # -# - Requests are refused if LLM Shield is unreachable or returns an error, rather than +# - Requests are refused if LLM Shield Proxy is unreachable or returns an error, rather than # being forwarded. Sending them on would hand the provider exactly the data this # guardrail exists to withhold. # - Restoring a value requires the request and the reply to share a session. LiteLLM's diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py similarity index 95% rename from litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py rename to litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py index f84f0f429aa..56bee894714 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py @@ -1,6 +1,6 @@ # +-------------------------------------------------------------+ # -# Use LLM Shield for reversible PII redaction +# Use LLM Shield Proxy for reversible PII redaction # https://github.com/ninadphalak/LLM-Shield-Proxy # # +-------------------------------------------------------------+ @@ -38,7 +38,7 @@ if TYPE_CHECKING: from litellm.caching.caching import DualCache from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -GUARDRAIL_NAME: Final = "llm_shield" +GUARDRAIL_NAME: Final = "llm_shield_proxy" _DEFAULT_API_BASE: Final = "http://localhost:8000" _REDACT_PATH: Final = "/v1/guard/redact" @@ -172,7 +172,7 @@ def _choice_index(choice: object) -> int: return index if isinstance(index, int) else 0 -class LLMShieldGuardrail(CustomGuardrail): +class LLMShieldProxyGuardrail(CustomGuardrail): """Redacts PII before it leaves the proxy and restores it in the response. Unlike a masking guardrail, the substitution is reversible. Outbound text is @@ -199,8 +199,9 @@ class LLMShieldGuardrail(CustomGuardrail): **kwargs: Any, # noqa: LIT008 # kwargs-ok: forwarded verbatim to CustomGuardrail.__init__ ) -> None: self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) - self.api_base: Final = (api_base or os.environ.get("LLM_SHIELD_API_BASE") or _DEFAULT_API_BASE).rstrip("/") - self.api_key: Final = api_key or os.environ.get("LLM_SHIELD_API_KEY") + env_base: Final = os.environ.get("LLM_SHIELD_PROXY_API_BASE") + self.api_base: Final = (api_base or env_base or _DEFAULT_API_BASE).rstrip("/") + self.api_key: Final = api_key or os.environ.get("LLM_SHIELD_PROXY_API_KEY") super().__init__(guardrail_name=guardrail_name, **kwargs) @classmethod @@ -219,7 +220,7 @@ class LLMShieldGuardrail(CustomGuardrail): return headers async def _call_shield(self, path: str, session_id: str, payload: JsonBody) -> Mapping[str, object]: - """Posts to LLM Shield, failing closed on any transport or status error. + """Posts to LLM Shield Proxy, failing closed on any transport or status error. A redaction guardrail that fails open sends the very data it exists to protect to a third-party provider, so an unreachable or erroring shield @@ -235,16 +236,16 @@ class LLMShieldGuardrail(CustomGuardrail): response.raise_for_status() return response.json() except httpx.HTTPStatusError as exc: - verbose_proxy_logger.exception("LLM Shield returned %s for %s", exc.response.status_code, path) + verbose_proxy_logger.exception("LLM Shield Proxy returned %s for %s", exc.response.status_code, path) raise GuardrailRaisedException( guardrail_name=self.guardrail_name, - message=f"LLM Shield returned {exc.response.status_code}; blocking the request.", + message=f"LLM Shield Proxy returned {exc.response.status_code}; blocking the request.", ) from exc except Exception as exc: - verbose_proxy_logger.exception("LLM Shield call to %s failed", path) + verbose_proxy_logger.exception("LLM Shield Proxy call to %s failed", path) raise GuardrailRaisedException( guardrail_name=self.guardrail_name, - message="LLM Shield is unreachable; blocking the request.", + message="LLM Shield Proxy is unreachable; blocking the request.", ) from exc async def _redact(self, texts: Sequence[str], session_id: str) -> Sequence[str]: @@ -262,7 +263,7 @@ class LLMShieldGuardrail(CustomGuardrail): if not isinstance(returned, list) or len(returned) != len(sent): raise GuardrailRaisedException( guardrail_name=self.guardrail_name, - message=f"LLM Shield {operation} returned an unexpected payload; blocking the request.", + message=f"LLM Shield Proxy {operation} returned an unexpected payload; blocking the request.", ) return tuple(returned) @@ -530,7 +531,7 @@ class LLMShieldGuardrail(CustomGuardrail): if not isinstance(emitted, str) or not isinstance(remaining, str): raise GuardrailRaisedException( guardrail_name=self.guardrail_name, - message="LLM Shield stream rehydration returned an unexpected payload.", + message="LLM Shield Proxy stream rehydration returned an unexpected payload.", ) return emitted, remaining diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 42ab6034069..41b33433d92 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -137,7 +137,7 @@ class SupportedGuardrailIntegrations(Enum): COMPRESR = "compresr" STRAIKER = "straiker" ALICE = "alice" - LLM_SHIELD = "llm_shield" + LLM_SHIELD_PROXY = "llm_shield_proxy" class Role(Enum): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/llm_shield.py b/litellm/types/proxy/guardrails/guardrail_hooks/llm_shield_proxy.py similarity index 51% rename from litellm/types/proxy/guardrails/guardrail_hooks/llm_shield.py rename to litellm/types/proxy/guardrails/guardrail_hooks/llm_shield_proxy.py index 8d7afd907b0..967d7ee75c2 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/llm_shield.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/llm_shield_proxy.py @@ -3,22 +3,22 @@ from pydantic import Field from .base import GuardrailConfigModel -class LLMShieldGuardrailConfigModel(GuardrailConfigModel): +class LLMShieldProxyGuardrailConfigModel(GuardrailConfigModel): api_key: str | None = Field( default=None, description=( - "The virtual key for the LLM Shield instance. If not provided, the " - "`LLM_SHIELD_API_KEY` environment variable is checked." + "The virtual key for the LLM Shield Proxy instance. If not provided, the " + "`LLM_SHIELD_PROXY_API_KEY` environment variable is checked." ), ) api_base: str | None = Field( default=None, description=( - "The base URL of the LLM Shield instance. If not provided, the `LLM_SHIELD_API_BASE` " + "The base URL of the LLM Shield Proxy instance. If not provided, the `LLM_SHIELD_PROXY_API_BASE` " "environment variable is checked, then `http://localhost:8000`." ), ) @staticmethod def ui_friendly_name() -> str: - return "LLM Shield" + return "LLM Shield Proxy" diff --git a/ruff-strict.toml b/ruff-strict.toml index f9d026c6011..21595556d28 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -33,7 +33,7 @@ external = [ # Same reason: `**kwargs` forwards verbatim to CustomGuardrail.__init__, and the lifecycle # hook signatures inherit `Any` for `response` from CustomLogger, so narrowing them here # would break the override rather than describe it. -"litellm/proxy/guardrails/guardrail_hooks/llm_shield/llm_shield.py" = ["ANN401"] +"litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py" = ["ANN401"] [lint.mccabe] max-complexity = 15 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py similarity index 96% rename from tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py rename to tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py index e9cc1a36b07..bc74b42a0e7 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py @@ -6,16 +6,16 @@ from httpx import Request, Response import litellm from litellm.exceptions import GuardrailRaisedException -from litellm.proxy.guardrails.guardrail_hooks.llm_shield.llm_shield import ( +from litellm.proxy.guardrails.guardrail_hooks.llm_shield_proxy.llm_shield_proxy import ( GUARDRAIL_NAME, - LLMShieldGuardrail, + LLMShieldProxyGuardrail, ) from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import Choices, Delta, Message, ModelResponse, ModelResponseStream, StreamingChoices -def _guardrail(**overrides: object) -> LLMShieldGuardrail: +def _guardrail(**overrides: object) -> LLMShieldProxyGuardrail: params: dict[str, object] = { "api_key": "test-key", "api_base": "http://shield.test", @@ -24,7 +24,7 @@ def _guardrail(**overrides: object) -> LLMShieldGuardrail: "default_on": True, } params.update(overrides) - return LLMShieldGuardrail(**params) + return LLMShieldProxyGuardrail(**params) def _response(payload: dict, status_code: int = 200) -> Response: @@ -35,7 +35,7 @@ def _response(payload: dict, status_code: int = 200) -> Response: ) -def _mock_post(guardrail: LLMShieldGuardrail, *payloads: dict) -> AsyncMock: +def _mock_post(guardrail: LLMShieldProxyGuardrail, *payloads: dict) -> AsyncMock: """Queues one shield response per expected call.""" mock = AsyncMock(side_effect=[_response(p) for p in payloads]) guardrail.async_handler.post = mock # type: ignore[method-assign] @@ -55,30 +55,30 @@ async def _drain(generator) -> list: def test_llm_shield_guardrail_config(monkeypatch: pytest.MonkeyPatch): """Should register through init_guardrails_v2 like any other provider.""" monkeypatch.setattr(litellm, "guardrail_name_config_map", {}) - monkeypatch.setenv("LLM_SHIELD_API_KEY", "test-key") + monkeypatch.setenv("LLM_SHIELD_PROXY_API_KEY", "test-key") init_guardrails_v2( all_guardrails=[ { - "guardrail_name": "llm_shield", - "litellm_params": {"guardrail": "llm_shield", "mode": "pre_call", "default_on": True}, + "guardrail_name": "llm_shield_proxy", + "litellm_params": {"guardrail": "llm_shield_proxy", "mode": "pre_call", "default_on": True}, } ], config_file_path="", ) - registered = [cb for cb in litellm.callbacks if isinstance(cb, LLMShieldGuardrail)] + registered = [cb for cb in litellm.callbacks if isinstance(cb, LLMShieldProxyGuardrail)] assert len(registered) == 1 - assert registered[0].guardrail_name == "llm_shield" + assert registered[0].guardrail_name == "llm_shield_proxy" -class TestLLMShieldInitialization: +class TestLLMShieldProxyInitialization: def test_api_base_defaults_to_localhost(self, monkeypatch: pytest.MonkeyPatch): - monkeypatch.delenv("LLM_SHIELD_API_BASE", raising=False) + monkeypatch.delenv("LLM_SHIELD_PROXY_API_BASE", raising=False) assert _guardrail(api_base=None).api_base == "http://localhost:8000" def test_api_base_reads_environment(self, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("LLM_SHIELD_API_BASE", "http://shield.internal:9000") + monkeypatch.setenv("LLM_SHIELD_PROXY_API_BASE", "http://shield.internal:9000") assert _guardrail(api_base=None).api_base == "http://shield.internal:9000" def test_trailing_slash_is_stripped(self): @@ -135,7 +135,7 @@ class TestRedaction: @pytest.mark.asyncio async def test_request_without_text_is_untouched(self): - """No text to redact means no call to LLM Shield. + """No text to redact means no call to LLM Shield Proxy. This deliberately uses a request with no caller text at all. An earlier version used a Responses-API `input`, which asserted the very bypass that diff --git a/ui/litellm-dashboard/public/assets/logos/llm_shield.svg b/ui/litellm-dashboard/public/assets/logos/llm_shield_proxy.svg similarity index 100% rename from ui/litellm-dashboard/public/assets/logos/llm_shield.svg rename to ui/litellm-dashboard/public/assets/logos/llm_shield_proxy.svg diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts index 3579457bcd5..38751cb1d43 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts @@ -320,9 +320,9 @@ export const GUARDRAIL_PRESETS: Record = { mode: "pre_call", defaultOn: false, }, - llm_shield: { - provider: "LLM Shield", - guardrailNameSuggestion: "LLM Shield", + llm_shield_proxy: { + provider: "LLM Shield Proxy", + guardrailNameSuggestion: "LLM Shield Proxy", // Both halves are required. With only pre_call the request is redacted and the // placeholders are handed straight back to the caller. mode: ["pre_call", "post_call"], diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts index d3212d66737..293c447604f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts @@ -28,7 +28,7 @@ const EXPECTED_PARTNER_LOGO_FILES: Record = { repelloai: "repelloai.png", straiker: "straiker.svg", alice: "alice.svg", - llm_shield: "llm_shield.svg", + llm_shield_proxy: "llm_shield_proxy.svg", }; describe("guardrail_garden_data logos", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts index d46847a5a80..74a13f2f611 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts @@ -475,14 +475,14 @@ export const PARTNER_GUARDRAIL_CARDS: GuardrailCardInfo[] = [ providerKey: "Alice", }, { - id: "llm_shield", - name: "LLM Shield", + id: "llm_shield_proxy", + name: "LLM Shield Proxy", description: "Self-hosted PII redaction that puts the original values back into the model's response, so the provider never receives personal data while the end user still sees it.", category: "partner", - logo: guardrailLogoMap["LLM Shield"], + logo: guardrailLogoMap["LLM Shield Proxy"], tags: ["PII", "Data Privacy", "Compliance", "Streaming"], - providerKey: "LLM Shield", + providerKey: "LLM Shield Proxy", }, ]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx index f620ea6dcd0..1bf7056bd5c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx @@ -1,7 +1,7 @@ import aimSecurityLogo from "../../../../../public/assets/logos/aim_security.jpeg"; import aktoLogo from "../../../../../public/assets/logos/akto.svg"; import aliceLogo from "../../../../../public/assets/logos/alice.svg"; -import llmShieldLogo from "../../../../../public/assets/logos/llm_shield.svg"; +import llmShieldProxyLogo from "../../../../../public/assets/logos/llm_shield_proxy.svg"; import aporiaLogo from "../../../../../public/assets/logos/aporia.png"; import bedrockLogo from "../../../../../public/assets/logos/bedrock.svg"; import catoNetworksLogo from "../../../../../public/assets/logos/cato_networks.svg"; @@ -86,7 +86,7 @@ export const guardrail_provider_map: Record = { QostodianNexus: "qostodian_nexus", Repelloai: "repelloai", Alice: "alice", - "LLM Shield": "llm_shield", + "LLM Shield Proxy": "llm_shield_proxy", }; // Function to populate provider map from API response - updates the original map @@ -210,7 +210,7 @@ export const guardrailLogoMap = { "RepelloAI Argus": repelloAiLogo.src, Straiker: straikerLogo.src, Alice: aliceLogo.src, - "LLM Shield": llmShieldLogo.src, + "LLM Shield Proxy": llmShieldProxyLogo.src, } satisfies Record; export const getGuardrailLogo = (displayName: string): string | undefined => From 6536d61adf80d668d4aebfbcce6c0b272e57de77 Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Thu, 3 Sep 2026 21:22:30 -0500 Subject: [PATCH 16/22] feat(guardrails): redact the participant name on a message `name` on a user or assistant turn identifies a person and was going to the provider intact. The proxy this integrates with already redacts it, so the integration was the weaker of the two. On a tool or function turn the same field carries the function's name, which has to arrive unchanged or the call stops routing. That case is skipped, and a test asserts the value is never even sent to the shield. --- .../llm_shield_proxy/llm_shield_proxy.py | 13 +++++++++ .../guardrail_hooks/test_llm_shield_proxy.py | 28 +++++++++++++++++++ 2 files changed, 41 insertions(+) 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 56bee894714..b1ad9563eb2 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 @@ -122,6 +122,18 @@ def _collect_content(container: MutableRequest, slots: _SlotSink) -> None: _collect(part, "text", slots) +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 message.get("tool_calls") or (): @@ -311,6 +323,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): for message in data.get("messages") or (): if isinstance(message, dict): _collect_content(message, slots) + _collect_participant_name(message, slots) _collect_tool_arguments(message, slots) _collect_responses_fields(data, slots) _collect_prompt(data, slots) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py index bc74b42a0e7..372bb6ed2ab 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py @@ -330,6 +330,34 @@ class TestRequestCoverage: assert data["input"] == ["[EMAIL_1]", "[PHONE_1]"] + @pytest.mark.asyncio + async def test_participant_name_is_redacted(self): + """`name` on a user turn identifies a person.""" + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["hi", "[PERSON_1]"]}) + + data = {"messages": [{"role": "user", "name": "Jane Doe", "content": "hi"}]} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert mock.call_args_list[0].kwargs["json"]["texts"] == ["hi", "Jane Doe"] + assert data["messages"][0]["name"] == "[PERSON_1]" + + @pytest.mark.asyncio + async def test_tool_function_name_is_left_alone(self): + """On a tool turn the same field is the function name. + + Redacting it would stop the call routing, so this asserts it is never sent + to the shield at all. + """ + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["result"]}) + + data = {"messages": [{"role": "tool", "name": "get_weather", "content": "result"}]} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert data["messages"][0]["name"] == "get_weather" + assert mock.call_args_list[0].kwargs["json"]["texts"] == ["result"] + @pytest.mark.asyncio async def test_every_shape_in_one_request_is_redacted(self): guardrail = _guardrail() From 403b06ea76b43ae0eeef635daea3cc1b08db1730 Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Thu, 3 Sep 2026 23:28:26 -0500 Subject: [PATCH 17/22] fix(guardrails): flush every held choice, and cover tool results and suffix Three review findings. The trailing flush walked the last chunk's choices, so a choice that finished earlier and stopped appearing lost whatever text was still held for it and its answer was truncated. It is now driven by the windows themselves and emits one chunk per choice, synthesising the choice when the terminal chunk omits it. That was data loss, not just under-redaction. An Anthropic tool_result carries its own content, as a string or as further blocks, and only each part's `text` was being collected. Handled recursively; image and audio parts still fall through untouched. The legacy completions `suffix` is forwarded to providers that support it and was never collected. Note the placement: it has to be gathered before the string-prompt early return, which is what the new test pins. --- .../llm_shield_proxy/llm_shield_proxy.py | 68 ++++++++++++----- .../guardrail_hooks/test_llm_shield_proxy.py | 75 +++++++++++++++++++ 2 files changed, 125 insertions(+), 18 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 b1ad9563eb2..47adc39aa9e 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 @@ -100,7 +100,8 @@ def _collect_entry(entries: MutableSeq, index: int, slots: _SlotSink) -> None: def _collect_prompt(data: MutableRequest, slots: _SlotSink) -> None: - """The Completions API sends its text in a top-level `prompt`.""" + """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) @@ -118,8 +119,13 @@ def _collect_content(container: MutableRequest, slots: _SlotSink) -> None: _collect(container, "content", slots) return for part in content if isinstance(content, list) else (): - if isinstance(part, dict): - _collect(part, "text", slots) + if not isinstance(part, dict): + continue + _collect(part, "text", slots) + # An Anthropic tool_result carries its own content, as a string or as more + # blocks. Image and audio parts have no text and fall through untouched. + if "content" in part: + _collect_content(part, slots) def _collect_participant_name(message: MutableRequest, slots: _SlotSink) -> None: @@ -486,8 +492,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): # A stream that ended without a finish_reason can still leave text held back. if last_chunk is not None and any(carries.values()): - trailing: Final = last_chunk.model_copy(deep=True) - if await self._flush_trailing(trailing, carries, session_id): + async for trailing in self._flush_trailing(last_chunk, carries, session_id): yield trailing async def _restore_choice(self, choice: Any, carries: _CarryWindows, session_id: str) -> None: @@ -513,23 +518,50 @@ class LLMShieldProxyGuardrail(CustomGuardrail): carries[index] = remaining # rebind-ok: this choice's window advances. delta.content = emitted - async def _flush_trailing(self, trailing: Any, carries: _CarryWindows, session_id: str) -> bool: - """Empties every still-held window into a copy of the last chunk.""" - emitted_any = False # rebind-ok: set once any choice contributes text. - for choice in getattr(trailing, "choices", None) or (): - delta = getattr(choice, "delta", None) - if delta is None: - continue - index = _choice_index(choice) - carry = carries.get(index, "") + async def _flush_trailing( + self, last_chunk: Any, carries: _CarryWindows, session_id: str + ) -> AsyncGenerator[Any, None]: + """Empties every window still holding text, one chunk per choice. + + Driven by the windows rather than by the last chunk's choices. A choice that + 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 index in sorted(carries): + carry = carries[index] if not carry: - delta.content = None continue text, remaining = await self._stream_step("", carry, True, session_id) carries[index] = remaining # rebind-ok: this choice's window advances. - delta.content = text or None - emitted_any = emitted_any or bool(text) - return emitted_any + if not text: + continue + chunk = self._chunk_for_choice(last_chunk, index) + if chunk is None: + continue + chunk.choices[0].delta.content = text + yield chunk + + @staticmethod + def _chunk_for_choice(last_chunk: Any, index: int) -> Any: + """A single-choice copy of the last chunk, carrying only `index`. + + Emitting one choice per chunk keeps a flush from reading as content on a + choice it does not belong to. + """ + chunk: Final = last_chunk.model_copy(deep=True) + raw_choices: Final = getattr(chunk, "choices", None) + if not raw_choices: + return None + choices: Final[tuple] = tuple(raw_choices) + matching: Final = tuple(choice for choice in choices if _choice_index(choice) == index) + kept: Final = matching[0] if matching else choices[0] + if getattr(kept, "delta", None) is None: + 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] # mutable-ok: the chunk model requires a list. + return chunk async def _stream_step(self, text: str, carry: str, final: bool, session_id: str) -> tuple[str, str]: """Returns ``(text safe to emit now, window still being held)``.""" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py index 372bb6ed2ab..472e2606608 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py @@ -358,6 +358,43 @@ class TestRequestCoverage: assert data["messages"][0]["name"] == "get_weather" assert mock.call_args_list[0].kwargs["json"]["texts"] == ["result"] + @pytest.mark.asyncio + async def test_anthropic_tool_result_content_is_redacted(self): + """A tool_result nests its own content, as a string or as more blocks.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["[EMAIL_1]", "[EMAIL_2]"]}) + + data = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "t1", "content": "found jane.doe@example.com"}, + { + "type": "tool_result", + "tool_use_id": "t2", + "content": [{"type": "text", "text": "also bob@example.com"}], + }, + ], + } + ] + } + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + + assert data["messages"][0]["content"][0]["content"] == "[EMAIL_1]" + assert data["messages"][0]["content"][1]["content"][0]["text"] == "[EMAIL_2]" + + @pytest.mark.asyncio + async def test_completions_suffix_is_redacted(self): + """LiteLLM forwards the legacy `suffix` to providers that support it.""" + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["signed [EMAIL_1]", "write to [EMAIL_1]"]}) + + data = {"prompt": "write to jane.doe@example.com", "suffix": "signed jane.doe@example.com"} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="atext_completion") + + assert data["suffix"] == "signed [EMAIL_1]" + @pytest.mark.asyncio async def test_every_shape_in_one_request_is_redacted(self): guardrail = _guardrail() @@ -682,6 +719,44 @@ class TestStreamingRehydration: assert sent[2]["carry"] == "A-held", "choice 0 must get its own window back" assert sent[3]["carry"] == "B-held", "choice 1 must get its own window back" + @pytest.mark.asyncio + async def test_a_choice_missing_from_the_last_chunk_still_flushes(self): + """Held text must not be dropped because its choice ended earlier. + + Choice 1 finishes and stops appearing, then the stream ends without a + finish_reason for choice 0. Flushing only the terminal chunk's choices would + discard whatever choice 1 was still holding and truncate its answer. + """ + guardrail = _guardrail(event_hook="post_call") + _mock_post( + guardrail, + {"text": "", "carry": "held-0"}, + {"text": "", "carry": "held-1"}, + {"text": "zero-done", "carry": ""}, + {"text": "one-done", "carry": ""}, + ) + + async def stream(): + yield ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content="a")), + StreamingChoices(index=1, delta=Delta(content="b")), + ] + ) + yield ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=None))]) + + chunks = await _drain( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, response=stream(), request_data={"messages": []} + ) + ) + + flushed = { + choice.index: choice.delta.content for chunk in chunks for choice in chunk.choices if choice.delta.content + } + assert flushed.get(1) == "one-done", "choice 1's held text was dropped" + assert flushed.get(0) == "zero-done" + @pytest.mark.asyncio async def test_chunks_are_forwarded_as_they_arrive(self): """Restoration must not buffer the stream into a single terminal chunk.""" From ce921f49e47ae06946d7484ca334a0b24ac650f1 Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Fri, 4 Sep 2026 02:31:20 -0500 Subject: [PATCH 18/22] fix(guardrails): walk nested tool results iteratively, with a depth bound CI flagged _collect_content as recursive. It was, and worse, it was unbounded: a tool_result nests its own content, the nesting is caller controlled, and the descent had nothing to stop it. That is a JSON bomb, not a style issue. Now an explicit queue with a depth bound of 8. Real payloads nest one or two deep. The queue is walked in document order because the shield maps its replies back by position, so collection order is part of the contract. --- .../llm_shield_proxy/llm_shield_proxy.py | 41 +++++++++++++------ .../guardrail_hooks/test_llm_shield_proxy.py | 17 ++++++++ 2 files changed, 46 insertions(+), 12 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 47adc39aa9e..347b70ecd75 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 @@ -71,6 +71,10 @@ JsonBody: TypeAlias = dict # 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. +_MAX_CONTENT_DEPTH: Final = 8 + _Slot: TypeAlias = tuple[str, Callable[[str], None]] # mutable-ok: Callable's param list. # The accumulator the collectors below append into. It never escapes @@ -113,19 +117,32 @@ def _collect_prompt(data: MutableRequest, slots: _SlotSink) -> None: def _collect_content(container: MutableRequest, slots: _SlotSink) -> None: - """`content` is either a string or the multimodal list of typed parts.""" - content: Final = container.get("content") - if isinstance(content, str): - _collect(container, "content", slots) - return - for part in content if isinstance(content, list) else (): - if not isinstance(part, dict): + """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. + """ + # 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] = [(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 - _collect(part, "text", slots) - # An Anthropic tool_result carries its own content, as a string or as more - # blocks. Image and audio parts have no text and fall through untouched. - if "content" in part: - _collect_content(part, slots) + if depth >= _MAX_CONTENT_DEPTH: + continue + for part in content if isinstance(content, list) else (): + if not isinstance(part, dict): + continue + # Image and audio parts have no text and fall through untouched. + _collect(part, "text", slots) + if "content" in part: + pending.append((part, depth + 1)) def _collect_participant_name(message: MutableRequest, slots: _SlotSink) -> None: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py index 472e2606608..a49bc5b1275 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py @@ -384,6 +384,23 @@ class TestRequestCoverage: assert data["messages"][0]["content"][0]["content"] == "[EMAIL_1]" assert data["messages"][0]["content"][1]["content"][0]["text"] == "[EMAIL_2]" + @pytest.mark.asyncio + async def test_deeply_nested_tool_results_are_bounded(self): + """Nesting is caller controlled, so the descent has to stop somewhere. + + The walk must terminate on a payload built to be pathological, rather than + following it as far as it goes. + """ + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["ok"] * 64}) + + deep: dict = {"type": "tool_result", "content": "jane.doe@example.com"} + for _ in range(200): + deep = {"type": "tool_result", "content": [deep]} + data = {"messages": [{"role": "user", "content": [deep]}]} + + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + @pytest.mark.asyncio async def test_completions_suffix_is_redacted(self): """LiteLLM forwards the legacy `suffix` to providers that support it.""" From 119ec629529a43a0e3504650d5458170a7a2dc17 Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Fri, 4 Sep 2026 02:35:32 -0500 Subject: [PATCH 19/22] fix(guardrails): redact Responses PromptObject variables A Responses request can send `prompt` as a PromptObject rather than a string. Its `variables` are substituted into the stored prompt on the provider side, so they are caller text, and the dict shape was falling through untouched. `id` and `version` pick which stored prompt to run and are left unchanged. --- .../llm_shield_proxy/llm_shield_proxy.py | 9 +++++++++ .../guardrail_hooks/test_llm_shield_proxy.py | 17 +++++++++++++++++ 2 files changed, 26 insertions(+) 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 347b70ecd75..dea06cafab4 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 @@ -110,6 +110,15 @@ def _collect_prompt(data: MutableRequest, slots: _SlotSink) -> None: if isinstance(prompt, str): _collect(data, "prompt", slots) return + if isinstance(prompt, dict): + # 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. + variables: Final = prompt.get("variables") + if isinstance(variables, dict): + for name in tuple(variables): + _collect(variables, name, slots) + return if not isinstance(prompt, list): return for index in range(len(prompt)): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py index a49bc5b1275..b056756ac15 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py @@ -401,6 +401,23 @@ class TestRequestCoverage: await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + @pytest.mark.asyncio + async def test_responses_prompt_object_variables_are_redacted(self): + """A PromptObject's variables are substituted into the prompt provider side. + + The id and version pick which stored prompt to run and have to arrive + unchanged; the variables are caller text. + """ + guardrail = _guardrail() + _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}) + + data = {"prompt": {"id": "pmpt_123", "version": "2", "variables": {"customer": "jane.doe@example.com"}}} + await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="aresponses") + + assert data["prompt"]["variables"]["customer"] == "[EMAIL_1]" + assert data["prompt"]["id"] == "pmpt_123" + assert data["prompt"]["version"] == "2" + @pytest.mark.asyncio async def test_completions_suffix_is_redacted(self): """LiteLLM forwards the legacy `suffix` to providers that support it.""" From 76ac63abd4eaf9e39d5a3e1e8b921e09a0e6f918 Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Fri, 4 Sep 2026 02:47:19 -0500 Subject: [PATCH 20/22] test(guardrails): assert the depth bound instead of only reaching the end The depth test asserted nothing, so it passed whether or not the bound held, and the test-quality gate counted it as a zero-assert test. It now sends a shallow value alongside a 200-deep chain and asserts the shallow one is collected while the value past the bound is not. --- .../guardrail_hooks/test_llm_shield_proxy.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py index b056756ac15..42781b441f7 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py @@ -392,15 +392,27 @@ class TestRequestCoverage: following it as far as it goes. """ guardrail = _guardrail() - _mock_post(guardrail, {"texts": ["ok"] * 64}) - deep: dict = {"type": "tool_result", "content": "jane.doe@example.com"} + captured: list = [] + + async def echo(url, headers, json, timeout): # noqa: ARG001 + captured.append(json["texts"]) + return _response({"texts": list(json["texts"])}) + + guardrail.async_handler.post = AsyncMock(side_effect=echo) # type: ignore[method-assign] + + deep: dict = {"type": "tool_result", "content": "past-the-bound@example.com"} for _ in range(200): deep = {"type": "tool_result", "content": [deep]} - data = {"messages": [{"role": "user", "content": [deep]}]} + data = {"messages": [{"role": "user", "content": [{"type": "text", "text": "shallow"}, deep]}]} await guardrail.async_pre_call_hook(user_api_key_dict=None, cache=None, data=data, call_type="completion") + sent = captured[0] + assert "shallow" in sent + assert "past-the-bound@example.com" not in sent, "the walk followed the chain past its bound" + assert len(sent) < 200 + @pytest.mark.asyncio async def test_responses_prompt_object_variables_are_redacted(self): """A PromptObject's variables are substituted into the prompt provider side. From 6a4fc88d21a4449eeb9e068a0482714b9b54c417 Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Sat, 5 Sep 2026 05:57:18 -0500 Subject: [PATCH 21/22] fix(guardrails): keep system-prompt values out of the restored reply Redaction put every span of a request into one vault, and the reply was restored against that same vault. System prompts are written by the application and the caller never sees them, so a caller who got the model to echo a placeholder back had its plaintext restored into their own reply -- a way to read a system prompt they were never shown. Server-authored spans now go into a vault of their own: system and developer turns, Anthropic's top-level `system`, and the Responses API `instructions`. Its id is deliberately never stored, so nothing restores against it. The reply is restored against the caller's vault alone, and an echoed placeholder from a system prompt comes back as the placeholder. Values the caller also wrote themselves are unaffected -- they are in the caller's vault too, and still restore. The extra round trip happens only when a request actually carries server-authored text. --- .../llm_shield_proxy/llm_shield_proxy.py | 63 +++++++++++----- .../guardrail_hooks/test_llm_shield_proxy.py | 72 +++++++++++++++++++ 2 files changed, 119 insertions(+), 16 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 dea06cafab4..a25b077cee1 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 @@ -51,6 +51,10 @@ _REHYDRATE_STREAM_PATH: Final = "/v1/guard/rehydrate/stream" # across concurrent requests. _SESSION_METADATA_KEY: Final = "llm_shield_session_id" +# 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 @@ -188,9 +192,15 @@ def _collect_system(data: MutableRequest, slots: _SlotSink) -> None: _collect(part, "text", slots) -def _collect_responses_fields(data: MutableRequest, slots: _SlotSink) -> None: - """The Responses API sends text outside `messages`, in `instructions` and `input`.""" - _collect(data, "instructions", 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. + """ + _collect(data, "instructions", privileged) request_input: Final = data.get("input") if isinstance(request_input, str): _collect(data, "input", slots) @@ -344,23 +354,33 @@ class LLMShieldProxyGuardrail(CustomGuardrail): # --- request traversal -------------------------------------------------------- @staticmethod - def _locate_request_texts(data: MutableRequest) -> Sequence[_Slot]: - """Finds every redactable span in an outbound request. + 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 caller text. + 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` -- 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. """ slots: Final[_SlotSink] = [] # mutable-ok: accumulator, frozen on return. + privileged: Final[_SlotSink] = [] # mutable-ok: accumulator, frozen on return. for message in data.get("messages") or (): if isinstance(message, dict): - _collect_content(message, slots) - _collect_participant_name(message, slots) - _collect_tool_arguments(message, slots) - _collect_responses_fields(data, slots) + 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(data, slots, privileged) _collect_prompt(data, slots) - _collect_system(data, slots) - return tuple(slots) + _collect_system(data, privileged) + return tuple(slots), tuple(privileged) # --- hooks -------------------------------------------------------------------- @@ -376,14 +396,25 @@ class LLMShieldProxyGuardrail(CustomGuardrail): if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is not True: return data - slots: Final = self._locate_request_texts(data) - if not slots: + slots, privileged = self._locate_request_texts(data) + if not slots and not privileged: return data - redacted: Final = await self._redact(tuple(text for text, _ in slots), self._mint_session_id(data)) + 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) + return data + + 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): write(replacement) - return data @log_guardrail_information async def async_post_call_success_hook( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py index 42781b441f7..70a29333297 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py @@ -1,3 +1,4 @@ +import json from types import SimpleNamespace from unittest.mock import AsyncMock @@ -605,6 +606,77 @@ class TestVaultIsolation: assert len(seen) == 2 + @pytest.mark.parametrize( + "data", + [ + pytest.param( + {"messages": [{"role": "system", "content": "S"}, {"role": "user", "content": "U"}]}, + id="system-turn", + ), + pytest.param( + {"messages": [{"role": "developer", "content": "S"}, {"role": "user", "content": "U"}]}, + id="developer-turn", + ), + pytest.param( + {"system": "S", "messages": [{"role": "user", "content": "U"}]}, + id="anthropic-top-level-system", + ), + pytest.param({"instructions": "S", "input": "U"}, id="responses-instructions"), + ], + ) + def test_server_authored_text_is_split_from_the_callers(self, data: dict): + """Every request shape must sort its server-authored spans out of the caller's.""" + caller, privileged = LLMShieldProxyGuardrail._locate_request_texts(data) + + assert [text for text, _ in caller] == ["U"] + assert [text for text, _ in privileged] == ["S"] + + @pytest.mark.asyncio + async def test_a_system_prompt_gets_a_vault_of_its_own(self): + """The reply is restored against the caller's vault, so the two cannot be one.""" + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}, {"texts": ["[EMAIL_2]"]}) + + data = { + "messages": [ + {"role": "system", "content": "escalate to admin@corp.internal"}, + {"role": "user", "content": "email a@b.com"}, + ] + } + await guardrail.async_pre_call_hook( + user_api_key_dict=None, cache=None, data=data, call_type="completion" + ) + + privileged_id, caller_id = ( + call.kwargs["headers"]["X-Session-ID"] for call in mock.call_args_list + ) + assert privileged_id != caller_id + assert guardrail._session_id(data) == caller_id + + @pytest.mark.asyncio + async def test_the_system_prompt_vault_id_is_never_stored(self): + """Nothing can restore against the system vault later, because its id is not kept. + + This is what stops a caller from having the model echo a placeholder out of a + system prompt they cannot see and receiving the plaintext behind it. + """ + guardrail = _guardrail() + mock = _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}, {"texts": ["[EMAIL_2]"]}) + + data = { + "messages": [ + {"role": "system", "content": "escalate to admin@corp.internal"}, + {"role": "user", "content": "email a@b.com"}, + ] + } + await guardrail.async_pre_call_hook( + user_api_key_dict=None, cache=None, data=data, call_type="completion" + ) + + privileged_id = mock.call_args_list[0].kwargs["headers"]["X-Session-ID"] + assert privileged_id not in json.dumps(data, default=str) + + class TestFailClosed: @pytest.mark.asyncio async def test_unreachable_shield_blocks_the_request(self): From 8b2ee5ac7a96d5833483d6262f1dc89cc63027db Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Sat, 5 Sep 2026 06:25:12 -0500 Subject: [PATCH 22/22] style(guardrails): satisfy ruff format and annotate the new tests `ruff format` wanted the widened `_collect_responses_fields` signature on one line, and the three tests added with the split-vault fix needed return annotations to keep ANN201 level with the base. --- .../guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py | 4 +--- .../guardrails/guardrail_hooks/test_llm_shield_proxy.py | 6 +++--- 2 files changed, 4 insertions(+), 6 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 a25b077cee1..3cab7b20cef 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 @@ -192,9 +192,7 @@ def _collect_system(data: MutableRequest, slots: _SlotSink) -> None: _collect(part, "text", slots) -def _collect_responses_fields( - data: MutableRequest, slots: _SlotSink, privileged: _SlotSink -) -> None: +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 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py index 70a29333297..700702f3bb2 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py @@ -624,7 +624,7 @@ class TestVaultIsolation: pytest.param({"instructions": "S", "input": "U"}, id="responses-instructions"), ], ) - def test_server_authored_text_is_split_from_the_callers(self, data: dict): + def test_server_authored_text_is_split_from_the_callers(self, data: dict) -> None: """Every request shape must sort its server-authored spans out of the caller's.""" caller, privileged = LLMShieldProxyGuardrail._locate_request_texts(data) @@ -632,7 +632,7 @@ class TestVaultIsolation: assert [text for text, _ in privileged] == ["S"] @pytest.mark.asyncio - async def test_a_system_prompt_gets_a_vault_of_its_own(self): + async def test_a_system_prompt_gets_a_vault_of_its_own(self) -> None: """The reply is restored against the caller's vault, so the two cannot be one.""" guardrail = _guardrail() mock = _mock_post(guardrail, {"texts": ["[EMAIL_1]"]}, {"texts": ["[EMAIL_2]"]}) @@ -654,7 +654,7 @@ class TestVaultIsolation: assert guardrail._session_id(data) == caller_id @pytest.mark.asyncio - async def test_the_system_prompt_vault_id_is_never_stored(self): + async def test_the_system_prompt_vault_id_is_never_stored(self) -> None: """Nothing can restore against the system vault later, because its id is not kept. This is what stops a caller from having the model echo a placeholder out of a