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: