From a635d7be6a10d5126790cd10df7b67fcfbf1a2a7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 20:29:09 -0700 Subject: [PATCH 1/9] fix(guardrails): write per-message guardrail rewrites back onto Responses input items A guardrail that answers one rewritten text per message it saw no longer matches the texts the Responses handler extracted once the request carries instructions or tool items, so the rewrite was rejected with a 500. Spread such an answer over the structured messages' text slots and write it back through the structured path, have Prompt Security modify return structured_messages directly, and give the chat completions pairing the same named rejection instead of a silent misalignment when the counts differ. --- .../base_llm/guardrail_translation/utils.py | 77 ++++++++++++- .../chat/guardrail_translation/handler.py | 4 + .../guardrail_translation/handler.py | 27 ++++- .../prompt_security/prompt_security.py | 39 ++++++- .../test_openai_guardrail_handler.py | 43 ++++++++ ...test_openai_responses_guardrail_handler.py | 103 ++++++++++++++++++ .../test_prompt_security_guardrails.py | 89 +++++++++++++++ 7 files changed, 373 insertions(+), 9 deletions(-) diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 94a780f8148..383e668e45c 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -1,8 +1,10 @@ from __future__ import annotations import json -from collections.abc import Callable, Iterator, Sequence -from typing import Final, TypeVar +from collections.abc import Callable, Iterator, Mapping, Sequence +from itertools import accumulate +from types import MappingProxyType +from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles from pydantic import BaseModel @@ -364,3 +366,74 @@ def merge_guardrailed_scoped_messages( yield from appended return list(_merged()) + + +def _content_part_text(part: object) -> str | None: + if not isinstance(part, Mapping): + return None + text: Final = part.get("text") + return text if isinstance(text, str) else None + + +def message_text_slot_count(message: AllMessageValues) -> int: + content: Final = message.get("content") + if isinstance(content, str): + return 1 + if isinstance(content, list): + return sum(1 for part in content if _content_part_text(part) is not None) + return 0 + + +def _part_with_text(part: object, text: str) -> object: + if not isinstance(part, Mapping): + return part + return {**part, "text": text} # mutable-ok: content parts stay JSON-plain dicts + + +def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> list[object]: + text_part_indices: Final = tuple( + index for index, part in enumerate(content) if _content_part_text(part) is not None + ) + replacement_by_index: Final = MappingProxyType(dict(zip(text_part_indices, texts))) + return [ # mutable-ok: message content stays a JSON list + _part_with_text(part, replacement_by_index[index]) if index in replacement_by_index else part + for index, part in enumerate(content) + ] + + +def _message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues: + content: Final = message.get("content") + if not isinstance(content, (str, list)): + return message + rewritten_content: Final = texts[0] if isinstance(content, str) else _content_with_slot_texts(content, texts) + rewritten: Final = {**message, "content": rewritten_content} # mutable-ok: chat rows stay JSON-plain dicts + return cast("AllMessageValues", rewritten) # cast-ok: the same row with only its text slots swapped + + +def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues | None: + if message_text_slot_count(message) != len(texts): + return None + return _message_with_slot_texts(message, texts) + + +def messages_with_slot_texts( + messages: Sequence[AllMessageValues], + texts: Sequence[str], +) -> list[AllMessageValues] | None: + """Spread one flat list of rewritten texts over the messages' text slots, in order. + + A slot is a string ``content`` or one list part carrying a string ``text``; + images and other parts ride along untouched. A guardrail that answers one + text per message it saw produces exactly this shape, which stops matching + the endpoint's own per-text extraction as soon as the request carries + instructions or tool items. Returns None unless the counts line up exactly, + so a rewrite never lands on the wrong slot. + """ + slot_counts: Final = tuple(message_text_slot_count(message) for message in messages) + if sum(slot_counts) != len(texts): + return None + offsets: Final = tuple(accumulate(slot_counts, initial=0)) + return [ # mutable-ok: guardrail rows travel as a list + _message_with_slot_texts(message, texts[start:end]) + for message, start, end in zip(messages, offsets, offsets[1:]) + ] diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 58ff03e6a0d..56fda636e9a 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -196,6 +196,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): else: # Step 3: Map guardrail responses back to original message structure if guardrailed_texts and texts_to_check: + if len(guardrailed_texts) != len(text_task_mappings): + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + raise UnappliableRequestRewrite(guardrail_to_apply.guardrail_name or "unknown") await self._apply_guardrail_responses_to_input_texts( messages=messages, responses=guardrailed_texts, diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 2fe11d9f7bd..61903a54cd3 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -53,6 +53,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import ( ) from litellm.llms.base_llm.guardrail_translation.utils import ( blocked_responses_stream_usage, + messages_with_slot_texts, stream_item_field, stream_item_fingerprint, stream_item_items, @@ -395,6 +396,20 @@ def _patched_request_fields( ) +def _guardrailed_structured_messages( + structured_messages: Sequence[AllMessageValues] | None, + sent_text_count: int, + guardrailed_inputs: GenericGuardrailAPIInputs, +) -> Sequence[AllMessageValues] | None: + returned: Final = guardrailed_inputs.get("structured_messages") + if returned is not None and returned is not structured_messages: + return returned + rewritten_texts: Final = guardrailed_inputs.get("texts") + if not structured_messages or rewritten_texts is None or len(rewritten_texts) == sent_text_count: + return None + return messages_with_slot_texts(structured_messages, rewritten_texts) + + def _patch_or_convert_request_fields( raw_input: object, instructions: object, @@ -473,7 +488,8 @@ class OpenAIResponsesHandler(BaseTranslation): form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(original_tools) ) extracted: Final = self._extract_guardrail_inputs(data, input_data, flattened_tool_groups) - if not extracted.inputs.get("texts"): + sent_texts: Final = extracted.inputs.get("texts") + if not sent_texts: return data if structured_messages: extracted.inputs["structured_messages"] = structured_messages @@ -486,7 +502,9 @@ class OpenAIResponsesHandler(BaseTranslation): self._apply_guardrailed_tools_to_data( data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools") ) - written_back: Final = self._written_back_request_fields(data, structured_messages, guardrailed_inputs) + written_back: Final = self._written_back_request_fields( + data, structured_messages, len(sent_texts), guardrailed_inputs + ) if written_back is not None: data["input"] = list(written_back.input) # mutable-ok: JSON body if written_back.instructions is None: @@ -553,10 +571,11 @@ class OpenAIResponsesHandler(BaseTranslation): def _written_back_request_fields( data: Mapping[str, object], structured_messages: Sequence[AllMessageValues] | None, + sent_text_count: int, guardrailed_inputs: GenericGuardrailAPIInputs, ) -> _RequestFields | None: - guardrailed: Final = guardrailed_inputs.get("structured_messages") - if guardrailed is None or guardrailed is structured_messages: + guardrailed: Final = _guardrailed_structured_messages(structured_messages, sent_text_count, guardrailed_inputs) + if guardrailed is None: return None return _patch_or_convert_request_fields( data.get("input"), diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 0954fe1698a..72c87a793a5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -2,6 +2,7 @@ import asyncio import base64 import os from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Optional import httpx @@ -14,11 +15,13 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.llms.base_llm.guardrail_translation.utils import message_with_slot_texts from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -27,6 +30,7 @@ if TYPE_CHECKING: _SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS: Final = 30.0 +_PROTECT_ROLES: Final = frozenset({"system", "user", "assistant"}) class PromptSecurityGuardrailMissingSecrets(Exception): @@ -275,14 +279,44 @@ class PromptSecurityGuardrail(CustomGuardrail): detail="Blocked by Prompt Security, Violations: " + ", ".join(violations), ) elif action == "modify": - # Extract modified texts from modified_messages modified_messages: Final = result.get("modified_messages", []) modified_texts: Final = self._extract_texts_from_messages(modified_messages) if modified_texts: inputs["texts"] = modified_texts + rewritten_messages: Final = self._structured_messages_with_modifications( + structured_messages, modified_messages + ) + if rewritten_messages is not None: + inputs["structured_messages"] = rewritten_messages return inputs + def _is_sent_to_protect(self, message: Mapping[str, object]) -> bool: + return self.check_tool_results or message.get("role") in _PROTECT_ROLES + + def _structured_messages_with_modifications( + self, + structured_messages: Sequence[AllMessageValues], + modified_messages: Sequence[Mapping[str, object]], + ) -> list[AllMessageValues] | None: + sent_indices: Final = tuple( + index for index, message in enumerate(structured_messages) if self._is_sent_to_protect(message) + ) + if not sent_indices or len(sent_indices) != len(modified_messages): + return None + rewritten: Final = tuple( + message_with_slot_texts(structured_messages[index], self._extract_texts_from_messages((modified,))) + for index, modified in zip(sent_indices, modified_messages) + ) + replacements: Final = MappingProxyType( + {index: message for index, message in zip(sent_indices, rewritten) if message is not None} + ) + if len(replacements) != len(sent_indices): + return None + return [ # mutable-ok: guardrail inputs take a list + replacements.get(index, message) for index, message in enumerate(structured_messages) + ] + async def _apply_guardrail_on_response( self, inputs: GenericGuardrailAPIInputs, @@ -678,14 +712,13 @@ class PromptSecurityGuardrail(CustomGuardrail): This allows checking tool results for indirect prompt injection when enabled. """ - supported_roles: Final = ["system", "user", "assistant"] filtered_messages: Final = [] transformed_count = 0 filtered_count = 0 for message in messages: role = message.get("role", "") - if role in supported_roles: + if role in _PROTECT_ROLES: filtered_messages.append(message) else: if self.check_tool_results: diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 5a29a96829f..4e2291fdec1 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1893,6 +1893,49 @@ class TestScanOnlyToolResults: assert data["messages"][4]["content"] == "and then?" +class ToolDroppingTextGuardrail(CustomGuardrail): + """Answers one text per non-tool message it saw, the way a guardrail that + filters tool rows out before scanning does, and hands back only texts.""" + + def __init__(self): + super().__init__(guardrail_name="tool-dropping-redactor") + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + kept = [m for m in inputs.get("structured_messages") or [] if m.get("role") != "tool"] + return {**inputs, "texts": [str(m.get("content")).replace("POISON", "[BLOCKED]") for m in kept]} + + +class TestPerMessageTextWriteBack: + """Texts that no longer pair one-to-one with what the handler extracted must be + rejected by name instead of sliding onto the wrong messages.""" + + @pytest.mark.asyncio + async def test_fewer_texts_than_extracted_over_a_tool_message_is_rejected(self): + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + handler = OpenAIChatCompletionsHandler() + original_messages = [ + {"role": "system", "content": "SYSTEM-PROMPT"}, + {"role": "user", "content": "fetch the page"}, + {"role": "assistant", "content": "fetching"}, + {"role": "tool", "tool_call_id": "call_1", "content": "page says POISON here"}, + {"role": "user", "content": "and then?"}, + ] + data = {"messages": json.loads(json.dumps(original_messages))} + + with pytest.raises(UnappliableRequestRewrite) as excinfo: + await handler.process_input_messages(data=data, guardrail_to_apply=ToolDroppingTextGuardrail()) + + assert excinfo.value.guardrail_name == "tool-dropping-redactor" + assert data["messages"] == original_messages, "a rejected rewrite must leave the request untouched" + + class TestBuildBlockSseChunks: """build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE chunks""" diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index a4f0a77a9b6..b4b467773da 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -2338,6 +2338,109 @@ def _parallel_tool_call_input() -> list: ] +SSN = "123-45-6789" +REDACTED_SSN = "" + + +def _slot_texts(message: dict) -> list[str]: + content = message.get("content") + if isinstance(content, str): + return [content] + if isinstance(content, list): + return [part["text"] for part in content if isinstance(part, dict) and isinstance(part.get("text"), str)] + return [] + + +class PerMessageRedactionGuardrail(CustomGuardrail): + """Guardrail that answers one redacted text per message it was shown and hands + back only texts, the way Prompt Security in modify mode and a generic guardrail + API server that scans per message do.""" + + def __init__(self, extra_texts: int = 0): + super().__init__(guardrail_name="per-message-redactor") + self.extra_texts = extra_texts + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = inputs.get("structured_messages") or [] + texts = [text.replace(SSN, REDACTED_SSN) for message in messages for text in _slot_texts(message)] + return {**inputs, "texts": texts + ["junk"] * self.extra_texts} + + +class TestPerMessageTextWriteBack: + """A guardrail that rewrites one text per message it saw must land on the + instructions and the input items those messages came from, not be rejected.""" + + @pytest.mark.asyncio + async def test_instructions_plus_tool_replay_gets_each_rewrite_in_place(self): + handler = OpenAIResponsesHandler() + function_call_item = { + "type": "function_call", + "call_id": "call_1", + "name": "lookup_customer", + "arguments": '{"query": "' + SSN + '"}', + } + data = { + "model": "gpt-5.6", + "instructions": "Never repeat the SSN " + SSN + " back.", + "input": [ + {"role": "user", "content": "Look up " + SSN + " for me."}, + function_call_item, + {"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'}, + ], + } + + result = await handler.process_input_messages(data, PerMessageRedactionGuardrail()) + + assert result["instructions"] == "Never repeat the SSN " + REDACTED_SSN + " back." + assert [item.get("type", item.get("role")) for item in result["input"]] == [ + "user", + "function_call", + "function_call_output", + ] + assert _slot_texts(result["input"][0]) == ["Look up " + REDACTED_SSN + " for me."] + assert result["input"][1] == function_call_item + assert result["input"][2]["output"] == '{"ssn": "' + REDACTED_SSN + '"}' + assert result["input"][2]["call_id"] == "call_1" + + @pytest.mark.asyncio + async def test_string_input_with_instructions_keeps_the_two_apart(self): + handler = OpenAIResponsesHandler() + data = { + "model": "gpt-5.6", + "instructions": "Redact " + SSN + " everywhere.", + "input": "My SSN is " + SSN + ".", + } + + result = await handler.process_input_messages(data, PerMessageRedactionGuardrail()) + + assert result["instructions"] == "Redact " + REDACTED_SSN + " everywhere." + assert [_slot_texts(item) for item in result["input"]] == [["My SSN is " + REDACTED_SSN + "."]] + + @pytest.mark.asyncio + async def test_count_matching_neither_texts_nor_messages_is_still_rejected(self): + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + handler = OpenAIResponsesHandler() + original_input = [ + {"role": "user", "content": "Look up " + SSN + " for me."}, + {"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'}, + ] + data = {"model": "gpt-5.6", "instructions": "Be terse.", "input": copy.deepcopy(original_input)} + + with pytest.raises(UnappliableRequestRewrite) as excinfo: + await handler.process_input_messages(data, PerMessageRedactionGuardrail(extra_texts=1)) + + assert excinfo.value.guardrail_name == "per-message-redactor" + assert data["input"] == original_input + assert data["instructions"] == "Be terse." + + class TestProvenancePatching: """The O(n) provenance pass must keep patching rewritten rows in place for the shapes real agent loops produce, and fall back safely everywhere else.""" diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py index e650f796f29..9e83098eb04 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -174,6 +174,95 @@ async def test_apply_guardrail_modify_request(monkeypatch: pytest.MonkeyPatch): assert result["texts"] == ["User prompt with PII: SSN [REDACTED]"] +def _modify_response(modified_messages: list) -> Response: + mock_response = Response( + json={"result": {"prompt": {"action": "modify", "modified_messages": modified_messages}}}, + status_code=200, + request=Request(method="POST", url="https://test.prompt.security/api/protect"), + ) + mock_response.raise_for_status = lambda: None + return mock_response + + +def _tool_replay_messages() -> list: + return [ + {"role": "system", "content": "Never echo an SSN like 123-45-6789."}, + { + "role": "user", + "content": [ + {"type": "text", "text": "Look up 123-45-6789"}, + {"type": "image_url", "image_url": {"url": "https://example.com/id-card.png"}}, + ], + }, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "123-45-6789"}'}, + {"role": "user", "content": "Summarize what you found."}, + ] + + +@pytest.mark.asyncio +async def test_modify_returns_structured_messages_with_tool_rows_kept(monkeypatch: pytest.MonkeyPatch): + """A per-message modify verdict comes back as structured_messages so the + endpoint handler can write it back by message, with the rows Prompt Security + never saw (tool results) and the non-text parts (images) left in place.""" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") + guardrail = PromptSecurityGuardrail(guardrail_name="test-guard", event_hook="pre_call", default_on=True) + messages = _tool_replay_messages() + inputs = {"texts": ["Look up 123-45-6789", "Summarize what you found."], "structured_messages": messages} + modified_messages = [ + {"role": "system", "content": "Never echo an SSN like [REDACTED]."}, + {"role": "user", "content": [{"type": "text", "text": "Look up [REDACTED]"}]}, + {"role": "assistant", "content": None}, + {"role": "user", "content": "Summarize what you found."}, + ] + + with patch.object(guardrail.async_handler, "post", return_value=_modify_response(modified_messages)): + result = await guardrail.apply_guardrail(inputs=inputs, request_data={"messages": messages}, input_type="request") + + assert result["structured_messages"] == [ + {"role": "system", "content": "Never echo an SSN like [REDACTED]."}, + { + "role": "user", + "content": [ + {"type": "text", "text": "Look up [REDACTED]"}, + {"type": "image_url", "image_url": {"url": "https://example.com/id-card.png"}}, + ], + }, + messages[2], + messages[3], + {"role": "user", "content": "Summarize what you found."}, + ] + assert result["structured_messages"] is not messages + assert result["texts"] == [ + "Never echo an SSN like [REDACTED].", + "Look up [REDACTED]", + "Summarize what you found.", + ] + + +@pytest.mark.asyncio +async def test_modify_with_unexpected_message_count_keeps_texts_only(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") + guardrail = PromptSecurityGuardrail(guardrail_name="test-guard", event_hook="pre_call", default_on=True) + messages = _tool_replay_messages() + inputs = {"texts": ["Look up 123-45-6789", "Summarize what you found."], "structured_messages": messages} + modified_messages = [{"role": "user", "content": "Look up [REDACTED]"}] + + with patch.object(guardrail.async_handler, "post", return_value=_modify_response(modified_messages)): + result = await guardrail.apply_guardrail(inputs=inputs, request_data={"messages": messages}, input_type="request") + + assert result["structured_messages"] is messages + assert result["texts"] == ["Look up [REDACTED]"] + + @pytest.mark.asyncio async def test_apply_guardrail_allow_request(monkeypatch: pytest.MonkeyPatch): """Test that apply_guardrail allows safe prompts""" From 09314f239c34f35a86ce4155b8c4d445944a6038 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 00:54:29 -0700 Subject: [PATCH 2/9] fix(guardrails): hand per-message rewrites back as structured_messages A guardrail that rewrites text per chat message now returns the rewritten rows as structured_messages instead of only texts, so the Responses and chat handlers write the rewrite back through the structured path. The generic guardrail API response accepts an optional structured_messages list, Prompt Security modify builds one from modified_messages, and rows a server echoes back exactly as shown are restored to the original row objects because the request model drops undeclared keys. Texts-only per-message answers keep the named rejection on both endpoints. --- .../base_llm/guardrail_translation/utils.py | 42 ++--- .../guardrail_translation/handler.py | 27 +--- .../generic_guardrail_api.py | 30 +++- .../prompt_security/prompt_security.py | 8 +- .../guardrail_hooks/generic_guardrail_api.py | 15 +- ...test_openai_responses_guardrail_handler.py | 148 +++++++++--------- .../test_generic_guardrail_api.py | 105 +++++++++++++ .../test_prompt_security_guardrails.py | 14 +- 8 files changed, 244 insertions(+), 145 deletions(-) diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 383e668e45c..1172c93959b 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -2,7 +2,6 @@ from __future__ import annotations import json from collections.abc import Callable, Iterator, Mapping, Sequence -from itertools import accumulate from types import MappingProxyType from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles @@ -390,7 +389,7 @@ def _part_with_text(part: object, text: str) -> object: return {**part, "text": text} # mutable-ok: content parts stay JSON-plain dicts -def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> list[object]: +def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> Sequence[object]: text_part_indices: Final = tuple( index for index, part in enumerate(content) if _content_part_text(part) is not None ) @@ -401,39 +400,18 @@ def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> ] -def _message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues: +def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues | None: + """Swap one rewritten text into each text slot of a chat row, in order. + + A slot is a string ``content`` or one list part carrying a string ``text``; + images and other parts ride along untouched. Returns None unless the counts + line up exactly, so a rewrite never lands on the wrong slot. + """ + if message_text_slot_count(message) != len(texts): + return None content: Final = message.get("content") if not isinstance(content, (str, list)): return message rewritten_content: Final = texts[0] if isinstance(content, str) else _content_with_slot_texts(content, texts) rewritten: Final = {**message, "content": rewritten_content} # mutable-ok: chat rows stay JSON-plain dicts return cast("AllMessageValues", rewritten) # cast-ok: the same row with only its text slots swapped - - -def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues | None: - if message_text_slot_count(message) != len(texts): - return None - return _message_with_slot_texts(message, texts) - - -def messages_with_slot_texts( - messages: Sequence[AllMessageValues], - texts: Sequence[str], -) -> list[AllMessageValues] | None: - """Spread one flat list of rewritten texts over the messages' text slots, in order. - - A slot is a string ``content`` or one list part carrying a string ``text``; - images and other parts ride along untouched. A guardrail that answers one - text per message it saw produces exactly this shape, which stops matching - the endpoint's own per-text extraction as soon as the request carries - instructions or tool items. Returns None unless the counts line up exactly, - so a rewrite never lands on the wrong slot. - """ - slot_counts: Final = tuple(message_text_slot_count(message) for message in messages) - if sum(slot_counts) != len(texts): - return None - offsets: Final = tuple(accumulate(slot_counts, initial=0)) - return [ # mutable-ok: guardrail rows travel as a list - _message_with_slot_texts(message, texts[start:end]) - for message, start, end in zip(messages, offsets, offsets[1:]) - ] diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 61903a54cd3..2fe11d9f7bd 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -53,7 +53,6 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import ( ) from litellm.llms.base_llm.guardrail_translation.utils import ( blocked_responses_stream_usage, - messages_with_slot_texts, stream_item_field, stream_item_fingerprint, stream_item_items, @@ -396,20 +395,6 @@ def _patched_request_fields( ) -def _guardrailed_structured_messages( - structured_messages: Sequence[AllMessageValues] | None, - sent_text_count: int, - guardrailed_inputs: GenericGuardrailAPIInputs, -) -> Sequence[AllMessageValues] | None: - returned: Final = guardrailed_inputs.get("structured_messages") - if returned is not None and returned is not structured_messages: - return returned - rewritten_texts: Final = guardrailed_inputs.get("texts") - if not structured_messages or rewritten_texts is None or len(rewritten_texts) == sent_text_count: - return None - return messages_with_slot_texts(structured_messages, rewritten_texts) - - def _patch_or_convert_request_fields( raw_input: object, instructions: object, @@ -488,8 +473,7 @@ class OpenAIResponsesHandler(BaseTranslation): form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(original_tools) ) extracted: Final = self._extract_guardrail_inputs(data, input_data, flattened_tool_groups) - sent_texts: Final = extracted.inputs.get("texts") - if not sent_texts: + if not extracted.inputs.get("texts"): return data if structured_messages: extracted.inputs["structured_messages"] = structured_messages @@ -502,9 +486,7 @@ class OpenAIResponsesHandler(BaseTranslation): self._apply_guardrailed_tools_to_data( data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools") ) - written_back: Final = self._written_back_request_fields( - data, structured_messages, len(sent_texts), guardrailed_inputs - ) + written_back: Final = self._written_back_request_fields(data, structured_messages, guardrailed_inputs) if written_back is not None: data["input"] = list(written_back.input) # mutable-ok: JSON body if written_back.instructions is None: @@ -571,11 +553,10 @@ class OpenAIResponsesHandler(BaseTranslation): def _written_back_request_fields( data: Mapping[str, object], structured_messages: Sequence[AllMessageValues] | None, - sent_text_count: int, guardrailed_inputs: GenericGuardrailAPIInputs, ) -> _RequestFields | None: - guardrailed: Final = _guardrailed_structured_messages(structured_messages, sent_text_count, guardrailed_inputs) - if guardrailed is None: + guardrailed: Final = guardrailed_inputs.get("structured_messages") + if guardrailed is None or guardrailed is structured_messages: return None return _patch_or_convert_request_fields( data.get("input"), diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index d8296003ae9..16159d32a7f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -7,7 +7,7 @@ import fnmatch import os -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal, Optional import httpx @@ -24,7 +24,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.llms.openai import ChatCompletionToolParam +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( GenericGuardrailAPIMetadata, GenericGuardrailAPIRequest, @@ -150,6 +150,22 @@ def _extract_inbound_headers( return None +def _rows_with_unchanged_originals( + original_rows: Sequence[AllMessageValues] | None, + shown_rows: Sequence[AllMessageValues] | None, + returned_rows: Sequence[AllMessageValues], +) -> tuple[AllMessageValues, ...]: + """The request model drops row keys its message types do not declare, so a + row the server echoes back verbatim is restored to the original row object; + only rows the server actually changed reach the endpoint write-back.""" + if original_rows is None or shown_rows is None or len(returned_rows) != len(original_rows): + return tuple(returned_rows) + return tuple( + original if returned == shown else returned + for original, shown, returned in zip(original_rows, shown_rows, returned_rows) + ) + + class GenericGuardrailAPI(CustomGuardrail): """ Generic Guardrail API integration for LiteLLM. @@ -322,6 +338,8 @@ class GenericGuardrailAPI(CustomGuardrail): texts: list, images: list[str] | None, tools: list[ChatCompletionToolParam] | None, + structured_messages: Sequence[AllMessageValues] | None, + shown_messages: Sequence[AllMessageValues] | None, guardrail_response: GenericGuardrailAPIResponse, ) -> GenericGuardrailAPIInputs: # Action is NONE or no modifications needed @@ -336,6 +354,12 @@ class GenericGuardrailAPI(CustomGuardrail): return_inputs["tools"] = guardrail_response.tools elif tools: return_inputs["tools"] = tools + if guardrail_response.structured_messages: + return_inputs["structured_messages"] = list( # mutable-ok: guardrail inputs take a list + _rows_with_unchanged_originals( + structured_messages, shown_messages, guardrail_response.structured_messages + ) + ) if guardrail_response.stream_holdback_chars is not None: return_inputs["stream_holdback_chars"] = guardrail_response.stream_holdback_chars return return_inputs @@ -473,6 +497,8 @@ class GenericGuardrailAPI(CustomGuardrail): texts=texts, images=images, tools=tools, + structured_messages=structured_messages, + shown_messages=guardrail_request.structured_messages, guardrail_response=guardrail_response, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 72c87a793a5..41d3b202344 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -287,7 +287,7 @@ class PromptSecurityGuardrail(CustomGuardrail): structured_messages, modified_messages ) if rewritten_messages is not None: - inputs["structured_messages"] = rewritten_messages + inputs["structured_messages"] = list(rewritten_messages) # mutable-ok: guardrail inputs take a list return inputs @@ -298,7 +298,7 @@ class PromptSecurityGuardrail(CustomGuardrail): self, structured_messages: Sequence[AllMessageValues], modified_messages: Sequence[Mapping[str, object]], - ) -> list[AllMessageValues] | None: + ) -> tuple[AllMessageValues, ...] | None: sent_indices: Final = tuple( index for index, message in enumerate(structured_messages) if self._is_sent_to_protect(message) ) @@ -313,9 +313,7 @@ class PromptSecurityGuardrail(CustomGuardrail): ) if len(replacements) != len(sent_indices): return None - return [ # mutable-ok: guardrail inputs take a list - replacements.get(index, message) for index, message in enumerate(structured_messages) - ] + return tuple(replacements.get(index, message) for index, message in enumerate(structured_messages)) async def _apply_guardrail_on_response( self, diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index 4a868c48352..44e2cc2404f 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -1,4 +1,5 @@ -from typing import Any, Final, Literal +from collections.abc import Mapping, Sequence +from typing import Any, Final, Literal, cast # noqa: TID251 # JSON chat rows have no typed constructor across roles from pydantic import BaseModel, ConfigDict, Field from typing_extensions import TypedDict @@ -158,12 +159,21 @@ def coerce_stream_holdback_value(value: Any) -> int: return 0 +def structured_messages_from_response(value: object) -> Sequence[AllMessageValues] | None: + if not isinstance(value, list): + return None + if not all(isinstance(message, Mapping) and isinstance(message.get("role"), str) for message in value): + return None + return cast("Sequence[AllMessageValues]", value) # cast-ok: JSON rows checked for a role, the same trust texts get + + class GenericGuardrailAPIResponse: """Response model for the Generic Guardrail API""" texts: list[str] | None images: list[str] | None tools: list[GuardrailToolParam] | None + structured_messages: Sequence[AllMessageValues] | None action: str blocked_reason: str | None stream_holdback_chars: list[int] | None @@ -176,12 +186,14 @@ class GenericGuardrailAPIResponse: images: list[str] | None = None, tools: list[GuardrailToolParam] | None = None, stream_holdback_chars: list[int] | None = None, + structured_messages: Sequence[AllMessageValues] | None = None, ) -> None: self.action = action self.blocked_reason = blocked_reason self.texts = texts self.images = images self.tools = tools + self.structured_messages = structured_messages # Number of trailing chars, indexed the same as ``texts``, that the # framework must withhold from streaming emission until the next # processing round (word-boundary safety for text transformations). @@ -200,4 +212,5 @@ class GenericGuardrailAPIResponse: images=data.get("images"), tools=data.get("tools"), stream_holdback_chars=stream_holdback_chars, + structured_messages=structured_messages_from_response(data.get("structured_messages")), ) diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index b4b467773da..394f134e99c 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -8,7 +8,7 @@ with guardrail transformations. import copy from collections.abc import Callable from typing import Any, List, Literal, Optional, Tuple -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import logging @@ -31,6 +31,7 @@ from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, ) from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools +from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GenericGuardrailAPI from litellm.types.llms.openai import ChatCompletionToolCallChunk from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, @@ -2342,103 +2343,94 @@ SSN = "123-45-6789" REDACTED_SSN = "" -def _slot_texts(message: dict) -> list[str]: - content = message.get("content") - if isinstance(content, str): - return [content] - if isinstance(content, list): - return [part["text"] for part in content if isinstance(part, dict) and isinstance(part.get("text"), str)] - return [] +def _redacted(value: object) -> object: + if isinstance(value, str): + return value.replace(SSN, REDACTED_SSN) + if isinstance(value, list): + return [{**part, "text": _redacted(part["text"])} if "text" in part else part for part in value] + return value -class PerMessageRedactionGuardrail(CustomGuardrail): - """Guardrail that answers one redacted text per message it was shown and hands - back only texts, the way Prompt Security in modify mode and a generic guardrail - API server that scans per message do.""" +def _per_message_guardrail_server(structured_messages_in_answer: bool) -> Callable[..., MagicMock]: + """Answers one redacted text per chat row it was shown, the way a guardrail + that scans per message does, and optionally the rewritten rows themselves.""" - def __init__(self, extra_texts: int = 0): - super().__init__(guardrail_name="per-message-redactor") - self.extra_texts = extra_texts + def post(url: str, json: dict, headers: dict) -> MagicMock: + rows = json["structured_messages"] + answer: dict = { + "action": "GUARDRAIL_INTERVENED", + "texts": [_redacted(row["content"]) if isinstance(row.get("content"), str) else "" for row in rows], + } + if structured_messages_in_answer: + answer["structured_messages"] = [{**row, "content": _redacted(row.get("content"))} for row in rows] + response = MagicMock() + response.json.return_value = answer + response.raise_for_status = MagicMock() + return response - async def apply_guardrail( - self, - inputs: GenericGuardrailAPIInputs, - request_data: dict, - input_type: Literal["request", "response"], - logging_obj: Optional[Any] = None, - ) -> GenericGuardrailAPIInputs: - messages = inputs.get("structured_messages") or [] - texts = [text.replace(SSN, REDACTED_SSN) for message in messages for text in _slot_texts(message)] - return {**inputs, "texts": texts + ["junk"] * self.extra_texts} + return post -class TestPerMessageTextWriteBack: - """A guardrail that rewrites one text per message it saw must land on the - instructions and the input items those messages came from, not be rejected.""" +def _per_message_redactor() -> GenericGuardrailAPI: + return GenericGuardrailAPI( + api_base="https://guardrail.test", + guardrail_name="per-message-redactor", + event_hook="pre_call", + default_on=True, + ) + + +def _tool_replay_request() -> dict: + return { + "model": "gpt-5.6", + "instructions": "Never repeat the SSN " + SSN + " back.", + "input": [ + {"role": "user", "content": "Look up " + SSN + " for me."}, + {"type": "function_call", "call_id": "call_1", "name": "lookup_customer", "arguments": '{"id": "42"}'}, + {"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'}, + ], + } + + +class TestPerMessageRewriteWriteBack: + """A guardrail that rewrites per chat row hands the rows back as + structured_messages, and the handler lands them on the instructions and the + input items they came from; the same rewrite handed back as texts alone has + no item to land on and is rejected by name instead of sent unrewritten.""" @pytest.mark.asyncio - async def test_instructions_plus_tool_replay_gets_each_rewrite_in_place(self): - handler = OpenAIResponsesHandler() - function_call_item = { - "type": "function_call", - "call_id": "call_1", - "name": "lookup_customer", - "arguments": '{"query": "' + SSN + '"}', - } - data = { - "model": "gpt-5.6", - "instructions": "Never repeat the SSN " + SSN + " back.", - "input": [ - {"role": "user", "content": "Look up " + SSN + " for me."}, - function_call_item, - {"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'}, - ], - } + async def test_structured_rows_land_on_instructions_and_tool_output(self): + guardrail = _per_message_redactor() + data = _tool_replay_request() + function_call_item = data["input"][1] - result = await handler.process_input_messages(data, PerMessageRedactionGuardrail()) + with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(True)): + result = await OpenAIResponsesHandler().process_input_messages(data, guardrail) assert result["instructions"] == "Never repeat the SSN " + REDACTED_SSN + " back." - assert [item.get("type", item.get("role")) for item in result["input"]] == [ - "user", - "function_call", - "function_call_output", - ] - assert _slot_texts(result["input"][0]) == ["Look up " + REDACTED_SSN + " for me."] + assert _texts(result["input"][0]) == ["Look up " + REDACTED_SSN + " for me."] assert result["input"][1] == function_call_item - assert result["input"][2]["output"] == '{"ssn": "' + REDACTED_SSN + '"}' - assert result["input"][2]["call_id"] == "call_1" - - @pytest.mark.asyncio - async def test_string_input_with_instructions_keeps_the_two_apart(self): - handler = OpenAIResponsesHandler() - data = { - "model": "gpt-5.6", - "instructions": "Redact " + SSN + " everywhere.", - "input": "My SSN is " + SSN + ".", + assert result["input"][2] == { + "type": "function_call_output", + "call_id": "call_1", + "output": '{"ssn": "' + REDACTED_SSN + '"}', } - result = await handler.process_input_messages(data, PerMessageRedactionGuardrail()) - - assert result["instructions"] == "Redact " + REDACTED_SSN + " everywhere." - assert [_slot_texts(item) for item in result["input"]] == [["My SSN is " + REDACTED_SSN + "."]] - @pytest.mark.asyncio - async def test_count_matching_neither_texts_nor_messages_is_still_rejected(self): + async def test_texts_only_per_message_answer_is_rejected_by_name(self): from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite - handler = OpenAIResponsesHandler() - original_input = [ - {"role": "user", "content": "Look up " + SSN + " for me."}, - {"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'}, - ] - data = {"model": "gpt-5.6", "instructions": "Be terse.", "input": copy.deepcopy(original_input)} + guardrail = _per_message_redactor() + data = _tool_replay_request() + original = copy.deepcopy(data) - with pytest.raises(UnappliableRequestRewrite) as excinfo: - await handler.process_input_messages(data, PerMessageRedactionGuardrail(extra_texts=1)) + with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(False)): + with pytest.raises(UnappliableRequestRewrite) as excinfo: + await OpenAIResponsesHandler().process_input_messages(data, guardrail) assert excinfo.value.guardrail_name == "per-message-redactor" - assert data["input"] == original_input - assert data["instructions"] == "Be terse." + assert data["input"] == original["input"] + assert data["instructions"] == original["instructions"] class TestProvenancePatching: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 83cc9ae8bb9..cc9942e0e40 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -582,6 +582,111 @@ class TestGuardrailActions: assert result_images is None +class TestStructuredMessagesInResponse: + """A guardrail server that rewrites per chat row answers with the rewritten + rows as structured_messages, which the endpoint handlers write back by row.""" + + @pytest.mark.asyncio + async def test_returned_rows_are_handed_back_as_structured_messages( + self, generic_guardrail, mock_request_data_input + ): + rewritten_rows = [ + {"role": "system", "content": "Never repeat an SSN."}, + {"role": "user", "content": "Look up [REDACTED] for me."}, + {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "[REDACTED]"}'}, + ] + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "texts": ["Never repeat an SSN.", "Look up [REDACTED] for me.", '{"ssn": "[REDACTED]"}'], + "structured_messages": rewritten_rows, + } + mock_response.raise_for_status = MagicMock() + + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response): + guardrailed_inputs = await generic_guardrail.apply_guardrail( + inputs={"texts": ["Look up 123-45-6789 for me."]}, + request_data=mock_request_data_input, + input_type="request", + ) + + assert guardrailed_inputs["structured_messages"] == rewritten_rows + assert guardrailed_inputs["texts"] == mock_response.json.return_value["texts"] + + @pytest.mark.asyncio + async def test_rows_echoed_back_as_shown_keep_their_original_keys( + self, generic_guardrail, mock_request_data_input + ): + tool_call_row = { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}, "index": 0} + ], + } + original_rows = [ + {"role": "user", "content": "Look up 123-45-6789 for me.", "name": "pat"}, + tool_call_row, + {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "123-45-6789"}'}, + ] + + def echo_with_tool_output_redacted(url, json, headers): + shown_rows = json["structured_messages"] + assert "index" not in shown_rows[1]["tool_calls"][0] + assert "name" not in shown_rows[0] + answer = MagicMock() + answer.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "texts": ["Look up 123-45-6789 for me."], + "structured_messages": [ + shown_rows[0], + shown_rows[1], + {**shown_rows[2], "content": '{"ssn": "[REDACTED]"}'}, + ], + } + answer.raise_for_status = MagicMock() + return answer + + with patch.object(generic_guardrail.async_handler, "post", side_effect=echo_with_tool_output_redacted): + guardrailed_inputs = await generic_guardrail.apply_guardrail( + inputs={"texts": ["Look up 123-45-6789 for me."], "structured_messages": original_rows}, + request_data=mock_request_data_input, + input_type="request", + ) + + returned_rows = guardrailed_inputs["structured_messages"] + assert returned_rows[0] is original_rows[0] + assert returned_rows[1] is tool_call_row + assert returned_rows[2] == {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "[REDACTED]"}'} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "structured_messages", + [[], [{"content": "a row with no role"}], "not a list"], + ids=["empty", "no_role", "not_a_list"], + ) + async def test_rows_that_are_not_chat_messages_are_ignored( + self, generic_guardrail, mock_request_data_input, structured_messages + ): + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "texts": ["[REDACTED]"], + "structured_messages": structured_messages, + } + mock_response.raise_for_status = MagicMock() + + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response): + guardrailed_inputs = await generic_guardrail.apply_guardrail( + inputs={"texts": ["Look up 123-45-6789 for me."]}, + request_data=mock_request_data_input, + input_type="request", + ) + + assert "structured_messages" not in guardrailed_inputs + assert guardrailed_inputs["texts"] == ["[REDACTED]"] + + class TestImageSupport: """Test image handling in guardrail requests""" diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py index 9e83098eb04..70083e50f01 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -1,5 +1,6 @@ import asyncio import base64 +from collections.abc import Mapping, Sequence from unittest.mock import AsyncMock, patch import pytest @@ -12,6 +13,7 @@ from litellm.proxy.guardrails.guardrail_hooks.prompt_security.prompt_security im PromptSecurityGuardrailMissingSecrets, ) from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 +from litellm.types.llms.openai import AllMessageValues def test_prompt_security_guard_config(monkeypatch: pytest.MonkeyPatch): @@ -174,7 +176,7 @@ async def test_apply_guardrail_modify_request(monkeypatch: pytest.MonkeyPatch): assert result["texts"] == ["User prompt with PII: SSN [REDACTED]"] -def _modify_response(modified_messages: list) -> Response: +def _modify_response(modified_messages: Sequence[Mapping[str, object]]) -> Response: mock_response = Response( json={"result": {"prompt": {"action": "modify", "modified_messages": modified_messages}}}, status_code=200, @@ -184,7 +186,7 @@ def _modify_response(modified_messages: list) -> Response: return mock_response -def _tool_replay_messages() -> list: +def _tool_replay_messages() -> list[AllMessageValues]: return [ {"role": "system", "content": "Never echo an SSN like 123-45-6789."}, { @@ -224,7 +226,9 @@ async def test_modify_returns_structured_messages_with_tool_rows_kept(monkeypatc ] with patch.object(guardrail.async_handler, "post", return_value=_modify_response(modified_messages)): - result = await guardrail.apply_guardrail(inputs=inputs, request_data={"messages": messages}, input_type="request") + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={"messages": messages}, input_type="request" + ) assert result["structured_messages"] == [ {"role": "system", "content": "Never echo an SSN like [REDACTED]."}, @@ -257,7 +261,9 @@ async def test_modify_with_unexpected_message_count_keeps_texts_only(monkeypatch modified_messages = [{"role": "user", "content": "Look up [REDACTED]"}] with patch.object(guardrail.async_handler, "post", return_value=_modify_response(modified_messages)): - result = await guardrail.apply_guardrail(inputs=inputs, request_data={"messages": messages}, input_type="request") + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={"messages": messages}, input_type="request" + ) assert result["structured_messages"] is messages assert result["texts"] == ["Look up [REDACTED]"] From 93f3911b32e04c124027e8c5a12961fdcb0fca90 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:22:49 -0700 Subject: [PATCH 3/9] fix(guardrails): drop the types import CodeQL reads as a package cycle --- litellm/llms/base_llm/guardrail_translation/utils.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 1172c93959b..34d648cf184 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -2,7 +2,6 @@ from __future__ import annotations import json from collections.abc import Callable, Iterator, Mapping, Sequence -from types import MappingProxyType from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles from pydantic import BaseModel @@ -390,13 +389,10 @@ def _part_with_text(part: object, text: str) -> object: def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> Sequence[object]: - text_part_indices: Final = tuple( - index for index, part in enumerate(content) if _content_part_text(part) is not None - ) - replacement_by_index: Final = MappingProxyType(dict(zip(text_part_indices, texts))) + remaining_texts: Final = iter(texts) return [ # mutable-ok: message content stays a JSON list - _part_with_text(part, replacement_by_index[index]) if index in replacement_by_index else part - for index, part in enumerate(content) + _part_with_text(part, next(remaining_texts)) if _content_part_text(part) is not None else part + for part in content ] From 1760fe628436fac7fb292fae9ecdf6016a51fc55 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:22:49 -0700 Subject: [PATCH 4/9] refactor(guardrails): return a fresh inputs mapping from the Prompt Security modify branch --- .../prompt_security/prompt_security.py | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 41d3b202344..5d533e98ff6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -37,6 +37,18 @@ class PromptSecurityGuardrailMissingSecrets(Exception): pass +def _inputs_with_modifications( + inputs: GenericGuardrailAPIInputs, + modified_texts: list[str], + rewritten_messages: Sequence[AllMessageValues] | None, +) -> GenericGuardrailAPIInputs: + texts_patch: Final[GenericGuardrailAPIInputs] = {"texts": modified_texts} if modified_texts else {} + messages_patch: Final[GenericGuardrailAPIInputs] = ( + {"structured_messages": list(rewritten_messages)} if rewritten_messages is not None else {} + ) + return {**inputs, **texts_patch, **messages_patch} + + class _ProtectVerdict(TypedDict, total=False): """One side (``prompt`` or ``response``) of an ``/api/protect`` verdict.""" @@ -280,14 +292,11 @@ class PromptSecurityGuardrail(CustomGuardrail): ) elif action == "modify": modified_messages: Final = result.get("modified_messages", []) - modified_texts: Final = self._extract_texts_from_messages(modified_messages) - if modified_texts: - inputs["texts"] = modified_texts - rewritten_messages: Final = self._structured_messages_with_modifications( - structured_messages, modified_messages + return _inputs_with_modifications( + inputs, + self._extract_texts_from_messages(modified_messages), + self._structured_messages_with_modifications(structured_messages, modified_messages), ) - if rewritten_messages is not None: - inputs["structured_messages"] = list(rewritten_messages) # mutable-ok: guardrail inputs take a list return inputs From 23a98cb85134342e20ccb7387875c9b0037bc7d2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:22:49 -0700 Subject: [PATCH 5/9] chore(ui): regenerate schema.d.ts after merging main --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..839aa52fa84 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16781,7 +16781,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16886,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) From 37447c98f77116a54e3c902a64a64fae618c7996 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:28:43 -0700 Subject: [PATCH 6/9] fix(guardrails): reject per-message texts that cannot land on a string input or a Messages request --- .../chat/guardrail_translation/handler.py | 3 ++ .../base_llm/guardrail_translation/utils.py | 6 +++ .../chat/guardrail_translation/handler.py | 5 +- .../guardrail_translation/handler.py | 7 +-- .../test_anthropic_guardrail_handler.py | 52 +++++++++++++++++++ ...test_openai_responses_guardrail_handler.py | 35 +++++++++++++ 6 files changed, 102 insertions(+), 6 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 9d50345d70d..656e9978eff 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -44,6 +44,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( scoped_structured_message_indices, stream_item_field, stream_item_fingerprint, + unappliable_request_rewrite, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, @@ -570,6 +571,8 @@ class AnthropicMessagesHandler(BaseTranslation): preserve_system_messages=has_midturn_system_message, ) else: + if guardrailed_texts and len(guardrailed_texts) != len(scanned): + raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) # Step 3: Map guardrail responses back to original message structure await self._apply_guardrail_responses_to_input( messages=messages, diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 34d648cf184..a80e20c5404 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -411,3 +411,9 @@ def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> rewritten_content: Final = texts[0] if isinstance(content, str) else _content_with_slot_texts(content, texts) rewritten: Final = {**message, "content": rewritten_content} # mutable-ok: chat rows stay JSON-plain dicts return cast("AllMessageValues", rewritten) # cast-ok: the same row with only its text slots swapped + + +def unappliable_request_rewrite(guardrail_name: str | None) -> Exception: + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + return UnappliableRequestRewrite(guardrail_name or "unknown") diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 56fda636e9a..ee68f8f6546 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -42,6 +42,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( stream_item_field, stream_item_fingerprint, stream_item_items, + unappliable_request_rewrite, ) from litellm.main import stream_chunk_builder from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -197,9 +198,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Step 3: Map guardrail responses back to original message structure if guardrailed_texts and texts_to_check: if len(guardrailed_texts) != len(text_task_mappings): - from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite - - raise UnappliableRequestRewrite(guardrail_to_apply.guardrail_name or "unknown") + raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) await self._apply_guardrail_responses_to_input_texts( messages=messages, responses=guardrailed_texts, diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 2fe11d9f7bd..27ff55f120c 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -56,6 +56,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( stream_item_field, stream_item_fingerprint, stream_item_items, + unappliable_request_rewrite, ) from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools from litellm.responses.litellm_completion_transformation.transformation import ( @@ -495,13 +496,13 @@ class OpenAIResponsesHandler(BaseTranslation): data["instructions"] = written_back.instructions # rebind-ok: data is an out-param elif isinstance(input_data, str): guardrailed_texts: Final = guardrailed_inputs.get("texts") or () + if len(guardrailed_texts) > 1: + raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data # rebind-ok: data is an out-param else: rewritten_texts: Final = guardrailed_inputs.get("texts") or () if len(rewritten_texts) != len(extracted.task_mappings): - from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite - - raise UnappliableRequestRewrite(guardrail_to_apply.guardrail_name or "unknown") + raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) await self._apply_guardrail_responses_to_input( messages=input_data, responses=rewritten_texts, diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 9fe56f4dc65..d6fd30638cc 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -2272,6 +2272,58 @@ class TestAnthropicMessagesHandlerStreamingScanKey: assert ended_key != open_key +class PerRowTextGuardrail(CustomGuardrail): + """Answers one redacted text per chat row it was shown, the way a guardrail + that scans per message does, and hands back only texts.""" + + def __init__(self): + super().__init__(guardrail_name="per-row-redactor") + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + rows = inputs.get("structured_messages") or [] + return {**inputs, "texts": [str(row.get("content")).replace("123-45-6789", "") for row in rows]} + + +class TestPerMessageTextWriteBack: + """Texts that no longer pair one-to-one with what the handler extracted must be + rejected by name instead of sliding onto the wrong messages.""" + + @pytest.mark.asyncio + async def test_one_text_per_row_over_a_system_prompt_is_rejected_by_name(self): + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + data = { + "model": "claude-sonnet-4-5", + "system": "Reply with exactly the SSN you were given.", + "messages": [{"role": "user", "content": "My SSN is 123-45-6789."}], + } + original = json.loads(json.dumps(data)) + + with pytest.raises(UnappliableRequestRewrite) as excinfo: + await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=PerRowTextGuardrail()) + + assert excinfo.value.guardrail_name == "per-row-redactor" + assert data["system"] == original["system"], "a rejected rewrite must leave the request untouched" + assert data["messages"] == original["messages"], "a rejected rewrite must leave the request untouched" + + @pytest.mark.asyncio + async def test_one_text_per_row_without_a_system_prompt_is_applied(self): + data = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "My SSN is 123-45-6789."}], + } + + await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=PerRowTextGuardrail()) + + assert data["messages"] == [{"role": "user", "content": "My SSN is ."}] + + class TestAnthropicMessagesHandlerPostCallHookResponse: def test_openai_shaped_stream_assembly_reaches_the_hook_as_a_messages_response(self): from litellm.types.utils import Choices, Message, ModelResponse, Usage diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 394f134e99c..ac719da169c 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -2392,6 +2392,14 @@ def _tool_replay_request() -> dict: } +def _string_input_request() -> dict: + return { + "model": "gpt-5.6", + "instructions": "Never repeat the SSN " + SSN + " back.", + "input": "My SSN is " + SSN + ".", + } + + class TestPerMessageRewriteWriteBack: """A guardrail that rewrites per chat row hands the rows back as structured_messages, and the handler lands them on the instructions and the @@ -2432,6 +2440,33 @@ class TestPerMessageRewriteWriteBack: assert data["input"] == original["input"] assert data["instructions"] == original["instructions"] + @pytest.mark.asyncio + async def test_structured_rows_land_on_instructions_and_string_input(self): + guardrail = _per_message_redactor() + data = _string_input_request() + + with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(True)): + result = await OpenAIResponsesHandler().process_input_messages(data, guardrail) + + assert result["instructions"] == "Never repeat the SSN " + REDACTED_SSN + " back." + assert [_texts(item) for item in result["input"]] == [["My SSN is " + REDACTED_SSN + "."]] + + @pytest.mark.asyncio + async def test_texts_only_per_message_answer_over_a_string_input_is_rejected_by_name(self): + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + guardrail = _per_message_redactor() + data = _string_input_request() + original = copy.deepcopy(data) + + with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(False)): + with pytest.raises(UnappliableRequestRewrite) as excinfo: + await OpenAIResponsesHandler().process_input_messages(data, guardrail) + + assert excinfo.value.guardrail_name == "per-message-redactor" + assert data["input"] == original["input"] + assert data["instructions"] == original["instructions"] + class TestProvenancePatching: """The O(n) provenance pass must keep patching rewritten rows in place for the From 2923c4ac5520642bc3c3e728f9b2974850a43f7e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 03:02:15 -0700 Subject: [PATCH 7/9] fix(guardrails): read the rewrite from texts when a guardrail echoes every row back unchanged --- .../generic_guardrail_api.py | 25 ++++++++------ .../test_generic_guardrail_api.py | 34 +++++++++++++++++++ 2 files changed, 49 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 16159d32a7f..3d1a173635e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -150,16 +150,20 @@ def _extract_inbound_headers( return None -def _rows_with_unchanged_originals( +def _structured_rows_to_write_back( original_rows: Sequence[AllMessageValues] | None, shown_rows: Sequence[AllMessageValues] | None, returned_rows: Sequence[AllMessageValues], -) -> tuple[AllMessageValues, ...]: +) -> tuple[AllMessageValues, ...] | None: """The request model drops row keys its message types do not declare, so a - row the server echoes back verbatim is restored to the original row object; - only rows the server actually changed reach the endpoint write-back.""" + row the server echoes back verbatim is restored to the original row object. + A server that echoes every row back unchanged has not rewritten anything + per row, so its answer is read from texts, as it was before rows could be + returned at all.""" if original_rows is None or shown_rows is None or len(returned_rows) != len(original_rows): return tuple(returned_rows) + if all(returned == shown for shown, returned in zip(shown_rows, returned_rows)): + return None return tuple( original if returned == shown else returned for original, shown, returned in zip(original_rows, shown_rows, returned_rows) @@ -354,12 +358,13 @@ class GenericGuardrailAPI(CustomGuardrail): return_inputs["tools"] = guardrail_response.tools elif tools: return_inputs["tools"] = tools - if guardrail_response.structured_messages: - return_inputs["structured_messages"] = list( # mutable-ok: guardrail inputs take a list - _rows_with_unchanged_originals( - structured_messages, shown_messages, guardrail_response.structured_messages - ) - ) + rows_to_write_back: Final = ( + _structured_rows_to_write_back(structured_messages, shown_messages, guardrail_response.structured_messages) + if guardrail_response.structured_messages + else None + ) + if rows_to_write_back is not None: + return_inputs["structured_messages"] = list(rows_to_write_back) # mutable-ok: guardrail inputs take a list if guardrail_response.stream_holdback_chars is not None: return_inputs["stream_holdback_chars"] = guardrail_response.stream_holdback_chars return return_inputs diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index cc9942e0e40..a5e79f84ef1 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -659,6 +659,40 @@ class TestStructuredMessagesInResponse: assert returned_rows[1] is tool_call_row assert returned_rows[2] == {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "[REDACTED]"}'} + @pytest.mark.asyncio + async def test_rows_all_echoed_back_as_shown_leave_the_rewrite_to_texts( + self, generic_guardrail, mock_request_data_input + ): + """A server written against the texts contract that echoes the request rows back + untouched while rewriting texts still gets its texts rewrite applied.""" + original_rows = [ + {"role": "system", "content": "Never repeat an SSN."}, + {"role": "user", "content": "Look up 123-45-6789 for me."}, + ] + + def echo_rows_and_rewrite_texts(url, json, headers): + answer = MagicMock() + answer.json.return_value = { + "action": "NONE", + "texts": [text.replace("123-45-6789", "[REDACTED]") for text in json["texts"]], + "structured_messages": json["structured_messages"], + } + answer.raise_for_status = MagicMock() + return answer + + with patch.object(generic_guardrail.async_handler, "post", side_effect=echo_rows_and_rewrite_texts): + guardrailed_inputs = await generic_guardrail.apply_guardrail( + inputs={ + "texts": ["Never repeat an SSN.", "Look up 123-45-6789 for me."], + "structured_messages": original_rows, + }, + request_data=mock_request_data_input, + input_type="request", + ) + + assert "structured_messages" not in guardrailed_inputs + assert guardrailed_inputs["texts"] == ["Never repeat an SSN.", "Look up [REDACTED] for me."] + @pytest.mark.asyncio @pytest.mark.parametrize( "structured_messages", From 16c326537f5aa06c597fa198a5cb9e7a01fd9327 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 18:39:54 -0700 Subject: [PATCH 8/9] fix(guardrails): define UnappliableRequestRewrite in the shared guardrail translation utils The three guardrail translation handlers imported the exception from the proxy policy engine through a function-local import, which CodeQL flagged as a cyclic import. The exception and its helper now live next to the handlers in the shared guardrail translation utils, and the tests import it from there. The Prompt Security modify-mode helper is also restructured into early-return TypedDict displays so the LIT002 budget stays at its limit --- .../base_llm/guardrail_translation/utils.py | 11 ++++++++-- .../prompt_security/prompt_security.py | 21 ++++++++++++++----- .../proxy/policy_engine/pipeline_executor.py | 9 -------- .../test_anthropic_guardrail_handler.py | 2 +- .../test_openai_guardrail_handler.py | 2 +- ...test_openai_responses_guardrail_handler.py | 4 ++-- .../guardrail_hooks/test_crowdstrike_aidr.py | 2 +- 7 files changed, 30 insertions(+), 21 deletions(-) diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index a80e20c5404..c47a56b5ea9 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -413,7 +413,14 @@ def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> return cast("AllMessageValues", rewritten) # cast-ok: the same row with only its text slots swapped -def unappliable_request_rewrite(guardrail_name: str | None) -> Exception: - from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite +class UnappliableRequestRewrite(Exception): + def __init__(self, guardrail_name: str) -> None: + super().__init__( + f"Guardrail '{guardrail_name}' rewrote the request in a way this endpoint cannot apply, " + "so the request was rejected rather than sent unrewritten" + ) + self.guardrail_name: Final = guardrail_name + +def unappliable_request_rewrite(guardrail_name: str | None) -> UnappliableRequestRewrite: return UnappliableRequestRewrite(guardrail_name or "unknown") diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 4581b8b863f..3f29b3e751c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -38,16 +38,27 @@ class PromptSecurityGuardrailMissingSecrets(Exception): pass +def _inputs_with_structured_messages( + inputs: GenericGuardrailAPIInputs, rewritten_messages: Sequence[AllMessageValues] | None +) -> GenericGuardrailAPIInputs: + if rewritten_messages is None: + return inputs + patched: Final[GenericGuardrailAPIInputs] = { + **inputs, + "structured_messages": list(rewritten_messages), # mutable-ok: the TypedDict field is declared as a list + } + return patched + + def _inputs_with_modifications( inputs: GenericGuardrailAPIInputs, modified_texts: list[str], rewritten_messages: Sequence[AllMessageValues] | None, ) -> GenericGuardrailAPIInputs: - texts_patch: Final[GenericGuardrailAPIInputs] = {"texts": modified_texts} if modified_texts else {} - messages_patch: Final[GenericGuardrailAPIInputs] = ( - {"structured_messages": list(rewritten_messages)} if rewritten_messages is not None else {} - ) - return {**inputs, **texts_patch, **messages_patch} + if not modified_texts: + return _inputs_with_structured_messages(inputs, rewritten_messages) + with_texts: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": modified_texts} + return _inputs_with_structured_messages(with_texts, rewritten_messages) class _ProtectVerdict(TypedDict, total=False): diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index ad45781d5d2..ed193c7f434 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -58,15 +58,6 @@ class UndeliverableStreamRewrite(Exception): self.guardrail_name: Final = guardrail_name -class UnappliableRequestRewrite(Exception): - def __init__(self, guardrail_name: str) -> None: - super().__init__( - f"Guardrail '{guardrail_name}' rewrote the request in a way this endpoint cannot apply, " - "so the request was rejected rather than sent unrewritten" - ) - self.guardrail_name: Final = guardrail_name - - def _tool_call_shape(tool_call: object) -> tuple[object, object]: plain: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call function: Final = plain.get("function") if isinstance(plain, Mapping) else None diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index d6fd30638cc..0882b329c49 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -2296,7 +2296,7 @@ class TestPerMessageTextWriteBack: @pytest.mark.asyncio async def test_one_text_per_row_over_a_system_prompt_is_rejected_by_name(self): - from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite data = { "model": "claude-sonnet-4-5", diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 4e2291fdec1..8ee0e982aa0 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1917,7 +1917,7 @@ class TestPerMessageTextWriteBack: @pytest.mark.asyncio async def test_fewer_texts_than_extracted_over_a_tool_message_is_rejected(self): - from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite handler = OpenAIChatCompletionsHandler() original_messages = [ diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index ac719da169c..48d86384633 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -2426,7 +2426,7 @@ class TestPerMessageRewriteWriteBack: @pytest.mark.asyncio async def test_texts_only_per_message_answer_is_rejected_by_name(self): - from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite guardrail = _per_message_redactor() data = _tool_replay_request() @@ -2453,7 +2453,7 @@ class TestPerMessageRewriteWriteBack: @pytest.mark.asyncio async def test_texts_only_per_message_answer_over_a_string_input_is_rejected_by_name(self): - from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite guardrail = _per_message_redactor() data = _string_input_request() diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py index a9ca13a463d..9849ad7ec88 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py @@ -1820,7 +1820,7 @@ async def test_unalignable_rewrite_is_rejected_never_sent_unredacted( Skipping the write-back would hand the model the unredacted text, so a guardrail could be bypassed by adding ``instructions`` or a tool call. """ - from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite data: dict[str, object] = {"model": "gpt-4o", "input": responses_input} if instructions is not None: From e01d97ea08dd27606672187f029adac680a49b3c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:31:22 -0700 Subject: [PATCH 9/9] fix(guardrails): read Prompt Security modified rows with the slot count's own predicate A chat row whose content carried an empty text part counted two slots in the chat completions handler while Prompt Security read one text out of the modified row, so the structured rewrite was dropped and the request got the named rejection. One shared helper now lists a row's slot texts and both the slot count and the modified-row reader use it. --- .../base_llm/guardrail_translation/utils.py | 12 ++++++---- .../prompt_security/prompt_security.py | 16 ++----------- .../test_prompt_security_guardrails.py | 24 +++++++++++++++++++ 3 files changed, 34 insertions(+), 18 deletions(-) diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index c47a56b5ea9..51d43436fc9 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -373,13 +373,17 @@ def _content_part_text(part: object) -> str | None: return text if isinstance(text, str) else None -def message_text_slot_count(message: AllMessageValues) -> int: +def message_slot_texts(message: Mapping[str, object]) -> tuple[str, ...]: content: Final = message.get("content") if isinstance(content, str): - return 1 + return (content,) if isinstance(content, list): - return sum(1 for part in content if _content_part_text(part) is not None) - return 0 + return tuple(text for part in content if (text := _content_part_text(part)) is not None) + return () + + +def message_text_slot_count(message: AllMessageValues) -> int: + return len(message_slot_texts(message)) def _part_with_text(part: object, text: str) -> object: diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 3f29b3e751c..7e43566f224 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -15,7 +15,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) -from litellm.llms.base_llm.guardrail_translation.utils import message_with_slot_texts +from litellm.llms.base_llm.guardrail_translation.utils import message_slot_texts, message_with_slot_texts from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -399,19 +399,7 @@ class PromptSecurityGuardrail(CustomGuardrail): return inputs def _extract_texts_from_messages(self, messages: Sequence[Mapping[str, object]]) -> list[str]: - """Extract text content from messages.""" - texts: Final = [] - for message in messages: - content = message.get("content") - if isinstance(content, str): - texts.append(content) - elif isinstance(content, list): - for item in content: - if isinstance(item, dict) and item.get("type") == "text": - text = item.get("text") - if text: - texts.append(text) - return texts + return [text for message in messages for text in message_slot_texts(message)] async def _process_standalone_images(self, images: list[str], user_api_key_alias: str | None) -> None: """Process standalone images from inputs (data URLs).""" diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py index 785895b6190..3218632a8d2 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -269,6 +269,30 @@ async def test_modify_with_unexpected_message_count_keeps_texts_only(monkeypatch assert result["texts"] == ["Look up [REDACTED]"] +@pytest.mark.asyncio +async def test_modify_keeps_empty_text_parts_as_slots(monkeypatch: pytest.MonkeyPatch): + """The chat handler counts an empty text part as a slot, so a modify verdict + that echoes the empty part still lines up with the row and its texts.""" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") + guardrail = PromptSecurityGuardrail(guardrail_name="test-guard", event_hook="pre_call", default_on=True) + messages: list[AllMessageValues] = [ + {"role": "user", "content": [{"type": "text", "text": "Look up 123-45-6789"}, {"type": "text", "text": ""}]} + ] + inputs = {"texts": ["Look up 123-45-6789", ""], "structured_messages": messages} + modified_messages = [ + {"role": "user", "content": [{"type": "text", "text": "Look up [REDACTED]"}, {"type": "text", "text": ""}]} + ] + + with patch.object(guardrail.async_handler, "post", return_value=_modify_response(modified_messages)): + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={"messages": messages}, input_type="request" + ) + + assert result["structured_messages"] == modified_messages + assert result["texts"] == ["Look up [REDACTED]", ""] + + @pytest.mark.asyncio async def test_apply_guardrail_allow_request(monkeypatch: pytest.MonkeyPatch): """Test that apply_guardrail allows safe prompts"""