From 7966666e6e5c3cb36831ede4d1336ca0bc226722 Mon Sep 17 00:00:00 2001 From: Youlian Simidjiyski Date: Sat, 22 Aug 2026 08:23:07 -0400 Subject: [PATCH] feat(guardrails): provide conversation context on response scans Response scans previously received only flat texts and tool_calls, so conversation-context guardrails had to cache the request-scan view and replay it at response time. The chat, Anthropic, and Responses API translation handlers now populate structured_messages (the scoped request conversation with the model's assistant turns appended) and tools on response scans, mirroring the request scan's operator scoping. Guardrails that treated structured_messages as request-scan-only are gated so their response behavior is unchanged: openai moderation keeps moderating the output texts, PromptGuard keeps its texts redact path, and Akto keeps ingesting the actual request body. --- .../chat/guardrail_translation/handler.py | 76 +++++++ .../guardrail_translation/base_translation.py | 88 +++++++- .../chat/guardrail_translation/handler.py | 45 ++++ .../guardrail_translation/handler.py | 34 +++ .../guardrails/guardrail_hooks/akto/akto.py | 15 +- .../guardrail_hooks/openai/moderations.py | 6 +- .../promptguard/promptguard.py | 8 +- .../test_anthropic_guardrail_handler.py | 138 ++++++++++++ .../test_openai_guardrail_handler.py | 211 ++++++++++++++++++ ...test_openai_responses_guardrail_handler.py | 161 +++++++++++++ .../openai/test_moderations.py | 46 ++++ .../guardrails/guardrail_hooks/test_akto.py | 31 +++ .../guardrail_hooks/test_promptguard.py | 43 ++++ 13 files changed, 893 insertions(+), 9 deletions(-) create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_akto.py diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 721a6653597..0db559e9794 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -314,6 +314,53 @@ class AnthropicMessagesHandler(BaseTranslation): ) return result if result else None + def scoped_request_conversation( + self, + request_data: dict, + guardrail_to_apply: "CustomGuardrail", + ) -> list[AllMessageValues] | None: + """ + Mirror the request scan's scoping: translate without the trusted + top-level prompt, hoist it back unless skip_system, and keep + in-sequence system entries in scope. + """ + if request_data.get("messages") is None: + return None + skip_system: Final = effective_skip_system_message_for_guardrail(guardrail_to_apply) + translation_source: Final = { # mutable-ok: API message payload + key: value for key, value in request_data.items() if key != "system" + } + full_structured_messages: Final = cast( + list[AllMessageValues], + self._translate_to_openai(translation_source).get("messages", []), + ) + hoisted_system_message: Final = None if skip_system else self._hoisted_top_level_system_message(request_data) + if hoisted_system_message is not None: + full_structured_messages.insert(0, hoisted_system_message) + scoped_indices: Final = scoped_structured_message_indices( + full_structured_messages, + scan_only_tool_results=effective_scan_only_tool_results_for_guardrail(guardrail_to_apply), + skip_system=False, + skip_tool=effective_skip_tool_message_for_guardrail(guardrail_to_apply), + ) + scoped: Final = [full_structured_messages[index] for index in scoped_indices] + return scoped or None + + def request_tools_for_guardrail( + self, + request_data: dict, + guardrail_to_apply: "CustomGuardrail", + ) -> list[ChatCompletionToolParam] | None: + if effective_scan_only_tool_results_for_guardrail(guardrail_to_apply): + return None + if not request_data.get("tools"): + return None + translation_source: Final = { # mutable-ok: API message payload + key: value for key, value in request_data.items() if key != "system" + } + tools: Final = self._translate_to_openai(translation_source).get("tools", []) + return list(tools) if tools else None + async def process_input_messages( self, data: dict, @@ -896,6 +943,17 @@ class AnthropicMessagesHandler(BaseTranslation): response, ) + structured_conversation: Final = self.response_scan_conversation( + request_data, + guardrail_to_apply, + self.assistant_turn_from_extraction(texts_to_check, tool_calls_to_check), + ) + if structured_conversation: + inputs["structured_messages"] = structured_conversation + response_scan_tools: Final = self.request_tools_for_guardrail(request_data, guardrail_to_apply) + if response_scan_tools: + inputs["tools"] = response_scan_tools + guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=request_data, @@ -962,6 +1020,24 @@ class AnthropicMessagesHandler(BaseTranslation): user_api_key_dict, key="response", ) + stream_tool_call_dicts: Final = tuple( + tool_call.model_dump() for tool_call in tool_calls_list or () + ) + structured_conversation: Final = self.response_scan_conversation( + prepared_request_data, + guardrail_to_apply, + self.assistant_turn_from_extraction( + [string_so_far] if isinstance(string_so_far, str) and string_so_far else [], + stream_tool_call_dicts, + ), + ) + if structured_conversation: + guardrail_inputs["structured_messages"] = structured_conversation + response_scan_tools: Final = self.request_tools_for_guardrail( + prepared_request_data, guardrail_to_apply + ) + if response_scan_tools: + guardrail_inputs["tools"] = response_scan_tools _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=guardrail_inputs, request_data=prepared_request_data, diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index ba96ab3dc99..da37cbf1a93 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -1,6 +1,14 @@ from abc import ABC, abstractmethod +from collections.abc import Mapping, Sequence from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Final, Optional +from typing import TYPE_CHECKING, Any, Final, Optional, cast + +from litellm.llms.base_llm.guardrail_translation.utils import ( + effective_scan_only_tool_results_for_guardrail, + effective_skip_system_message_for_guardrail, + effective_skip_tool_message_for_guardrail, + scoped_structured_message_indices, +) if TYPE_CHECKING: from litellm.integrations.custom_guardrail import ( @@ -9,7 +17,7 @@ if TYPE_CHECKING: ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import UserAPIKeyAuth - from litellm.types.llms.openai import AllMessageValues + from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @dataclass(slots=True) @@ -157,6 +165,82 @@ class BaseTranslation(ABC): """ return None + def scoped_request_conversation( + self, + request_data: dict, + guardrail_to_apply: "CustomGuardrail", + ) -> list["AllMessageValues"] | None: + """ + The request conversation as the guardrail's request scan saw it: the + handler's structured messages with the operator scoping flags applied. + + Override when the request scan scopes differently (e.g. Anthropic's + top-level system prompt hoisting). + """ + structured_messages: Final = self.get_structured_messages(request_data) + if not structured_messages: + return None + scoped_indices: Final = scoped_structured_message_indices( + structured_messages, + scan_only_tool_results=effective_scan_only_tool_results_for_guardrail(guardrail_to_apply), + skip_system=effective_skip_system_message_for_guardrail(guardrail_to_apply), + skip_tool=effective_skip_tool_message_for_guardrail(guardrail_to_apply), + ) + scoped: Final = [structured_messages[index] for index in scoped_indices] + return scoped or None + + def response_scan_conversation( + self, + request_data: dict | None, + guardrail_to_apply: "CustomGuardrail", + response_turns: Sequence["AllMessageValues"], + ) -> list["AllMessageValues"] | None: + """ + Full conversation for a response scan: the scoped request conversation + with the model's response turns appended. + + Returns None when the request context is unavailable (SDK/direct-call + path fabricates request_data without messages); guardrails then fall + back to scanning the extracted texts. + """ + if request_data is None: + return None + request_conversation: Final = self.scoped_request_conversation(request_data, guardrail_to_apply) + if request_conversation is None: + return None + return [*request_conversation, *response_turns] + + def request_tools_for_guardrail( + self, + request_data: dict, + guardrail_to_apply: "CustomGuardrail", + ) -> list["ChatCompletionToolParam"] | None: + """ + The request's tool definitions in the shape the request scan sends them. + + Override in tool-capable handlers; default returns None. + """ + return None + + @staticmethod + def assistant_turn_from_extraction( + texts: Sequence[str], + tool_calls: Sequence[Mapping[str, object]] | None = None, + ) -> list["AllMessageValues"]: + """ + One OpenAI-shape assistant turn built from the texts and tool calls a + handler's response extraction collected; empty when there is nothing. + """ + tool_call_items: Final = tuple(tool_calls or ()) + if not texts and not tool_call_items: + return [] + turn: Final = { + "role": "assistant", + "content": "\n".join(texts), + **({"tool_calls": list(tool_call_items)} if tool_call_items else {}), + } + return [cast("AllMessageValues", turn)] + def extract_request_tool_names(self, data: dict) -> list[str]: """ Extract tool names from the request body for allowlist/policy checks. diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index e411dc497fc..5352be2a0ff 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -396,6 +396,15 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if hasattr(response, "model") and response.model: inputs["model"] = response.model + structured_conversation: Final = self.response_scan_conversation( + request_data, guardrail_to_apply, self._build_response_turns(response) + ) + if structured_conversation: + inputs["structured_messages"] = structured_conversation + response_scan_tools: Final = self.request_tools_for_guardrail(request_data, guardrail_to_apply) + if response_scan_tools: + inputs["tools"] = response_scan_tools + guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=request_data, @@ -832,6 +841,42 @@ class OpenAIChatCompletionsHandler(BaseTranslation): tool_calls_to_check.append(tool_call_dict) tool_call_task_mappings.append((choice_idx, int(tool_call_idx))) + def request_tools_for_guardrail( + self, + request_data: dict, + guardrail_to_apply: "CustomGuardrail", + ) -> list[ChatCompletionToolParam] | None: + if effective_scan_only_tool_results_for_guardrail(guardrail_to_apply): + return None + tools: Final = request_data.get("tools") + return cast(list[ChatCompletionToolParam], tools) if tools else None + + def _build_response_turns(self, response: "ModelResponse") -> list[AllMessageValues]: + """Assistant turns for the response-scan conversation, one per choice.""" + return [ + turn + for choice in response.choices + if isinstance(choice, litellm.Choices) and (turn := self._choice_assistant_turn(choice)) is not None + ] + + def _choice_assistant_turn(self, choice: Choices) -> AllMessageValues | None: + tool_call_dicts: Final = tuple( + converted + for tool_call in (choice.message.tool_calls or []) + if (converted := self._convert_tool_call_to_dict(tool_call)) is not None + ) + content: Final = choice.message.content + if content is None and not tool_call_dicts: + return None + return cast( + AllMessageValues, + { + "role": "assistant", + "content": content, + **({"tool_calls": list(tool_call_dicts)} if tool_call_dicts else {}), + }, + ) + def _convert_tool_call_to_dict(self, tool_call: dict[str, Any] | Any) -> dict[str, Any] | None: """ Convert a tool call object to dictionary format. diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 7c5d8ac99ad..2d747d43968 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -248,6 +248,18 @@ class OpenAIResponsesHandler(BaseTranslation): ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools) tools_to_check.extend(cast(list[ChatCompletionToolParam], transformed_tools)) + def request_tools_for_guardrail( + self, + request_data: dict, + guardrail_to_apply: "CustomGuardrail", + ) -> list[ChatCompletionToolParam] | None: + raw_tools: Final = request_data.get("tools") + if not raw_tools: + return None + tools_to_check: Final[list[ChatCompletionToolParam]] = [] + self._extract_and_transform_tools(raw_tools, tools_to_check) + return tools_to_check or None + def _remap_tools_to_responses_api_format(self, guardrailed_tools: list[Any]) -> list[dict[str, object]]: """ Remap guardrail-returned tools (Chat Completion format) back to @@ -455,6 +467,17 @@ class OpenAIResponsesHandler(BaseTranslation): if response_model: inputs["model"] = response_model + structured_conversation: Final = self.response_scan_conversation( + request_data, + guardrail_to_apply, + self.assistant_turn_from_extraction(texts_to_check, tool_calls_to_check), + ) + if structured_conversation: + inputs["structured_messages"] = structured_conversation + response_scan_tools: Final = self.request_tools_for_guardrail(request_data, guardrail_to_apply) + if response_scan_tools: + inputs["tools"] = response_scan_tools + guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=request_data, @@ -545,6 +568,17 @@ class OpenAIResponsesHandler(BaseTranslation): if response_model: inputs["model"] = response_model + structured_conversation: Final = self.response_scan_conversation( + request_data, + guardrail_to_apply, + self.assistant_turn_from_extraction(texts_to_check, tool_calls_to_check), + ) + if structured_conversation: + inputs["structured_messages"] = structured_conversation + response_scan_tools: Final = self.request_tools_for_guardrail(request_data, guardrail_to_apply) + if response_scan_tools: + inputs["tools"] = response_scan_tools + guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=request_data, diff --git a/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py b/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py index 3716d00774f..92332d896f7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py +++ b/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py @@ -162,12 +162,19 @@ class AktoGuardrail(CustomGuardrail): def build_request_body( inputs: GenericGuardrailAPIInputs, request_data: dict | None = None, + *, + prefer_structured_messages: bool = True, ) -> dict[str, Any]: - """Build the LLM request body from guardrail inputs (messages, model, tools).""" + """Build the LLM request body from guardrail inputs (messages, model, tools). + + ``prefer_structured_messages`` is False on response scans, where + ``structured_messages`` carries the response turns too and would + misrepresent the request in the ingested payload. + """ model: Final = inputs.get("model", "") or "" body: Final[dict[str, Any]] = {"model": model} - structured: Final = inputs.get("structured_messages") + structured: Final = inputs.get("structured_messages") if prefer_structured_messages else None if structured: body["messages"] = structured elif request_data is not None and request_data.get("messages"): @@ -232,7 +239,9 @@ class AktoGuardrail(CustomGuardrail): """ request_path: Final = self.extract_request_path(request_data) request_headers: Final = self.build_request_headers(request_data) - request_body: Final = self.build_request_body(inputs, request_data) + request_body: Final = self.build_request_body( + inputs, request_data, prefer_structured_messages=not include_response + ) tag: Final = self.build_tag_metadata(request_data) response_payload = json.dumps({}) # Empty body wrapper when no response yet diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py index 10683550f85..fdd53148a88 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py @@ -195,8 +195,10 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): # Extract text to moderate from inputs text_to_moderate: str | None = None - # Prefer structured_messages if available (has role context) - if structured_messages := inputs.get("structured_messages"): + # Prefer structured_messages if available (has role context). Response + # scans moderate the model output via texts; the conversation would + # point moderation back at the user prompt. + if input_type == "request" and (structured_messages := inputs.get("structured_messages")): text_to_moderate = self.get_user_prompt(structured_messages) # Fall back to texts diff --git a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py index c25f704567e..2b77220490d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py @@ -105,7 +105,11 @@ class PromptGuardGuardrail(CustomGuardrail): structured_messages: Final = inputs.get("structured_messages", []) model: Final = inputs.get("model") - if structured_messages: + # Response scans keep the texts path: the redact write-back extracts + # user-role texts, which would clobber the response texts if the + # conversation were sent instead. + use_structured: Final = input_type == "request" and bool(structured_messages) + if use_structured: messages = list(structured_messages) elif texts: messages = [{"role": "user", "content": text} for text in texts] @@ -175,7 +179,7 @@ class PromptGuardGuardrail(CustomGuardrail): if decision == "redact": redacted: Final = result.get("redacted_messages") if redacted: - if structured_messages: + if use_structured: inputs["structured_messages"] = redacted if "texts" in inputs: extracted: Final = self._extract_texts_from_messages( 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 b219dcba491..ff18f6b93bf 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 @@ -1823,3 +1823,141 @@ class TestAnthropicMessagesScanOnlyToolResults: assert guardrail.captured_inputs is not None assert guardrail.captured_inputs.get("images") == ["TOOL_IMG"] + + +class MockInputsRecordingGuardrail(CustomGuardrail): + """Records the inputs of every apply_guardrail call without modifying anything.""" + + def __init__(self): + super().__init__(guardrail_name="inputs-recording") + self.calls: list = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + import copy + + self.calls.append((input_type, copy.deepcopy(dict(inputs)))) + return inputs + + +class TestAnthropicResponseScanConversation: + """Response scans receive the request conversation (with the top-level system + prompt hoisted, mirroring the request scan) plus the model's assistant turn, + and the request's tools translated to OpenAI shape.""" + + def _request_data(self, **overrides) -> dict: + data = { + "model": "claude-sonnet-4-5", + "system": "you are a helpful assistant", + "messages": [{"role": "user", "content": "what's the weather?"}], + **overrides, + } + return {key: value for key, value in data.items() if value is not None} + + @pytest.mark.asyncio + async def test_response_scan_receives_hoisted_system_and_assistant_turn(self): + handler = AnthropicMessagesHandler() + guardrail = MockInputsRecordingGuardrail() + response = { + "id": "msg_1", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "Sunny today."}], + } + + await handler.process_output_response( + response=response, + guardrail_to_apply=guardrail, + request_data=self._request_data(), + ) + + input_type, inputs = guardrail.calls[-1] + assert input_type == "response" + conversation = inputs["structured_messages"] + assert conversation[0]["role"] == "system" + assert "helpful assistant" in str(conversation[0]["content"]) + assert conversation[1] == {"role": "user", "content": "what's the weather?"} + assert conversation[-1] == {"role": "assistant", "content": "Sunny today."} + + @pytest.mark.asyncio + async def test_skip_system_excludes_hoisted_prompt_from_response_scan(self): + handler = AnthropicMessagesHandler() + guardrail = MockInputsRecordingGuardrail() + guardrail.skip_system_message_in_guardrail = True + response = { + "id": "msg_1", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "Sunny today."}], + } + + await handler.process_output_response( + response=response, + guardrail_to_apply=guardrail, + request_data=self._request_data(), + ) + + _, inputs = guardrail.calls[-1] + conversation = inputs["structured_messages"] + assert all(message["role"] != "system" for message in conversation) + assert conversation[0] == {"role": "user", "content": "what's the weather?"} + assert conversation[-1] == {"role": "assistant", "content": "Sunny today."} + + @pytest.mark.asyncio + async def test_response_scan_includes_translated_tools_and_tool_call_turn(self): + handler = AnthropicMessagesHandler() + guardrail = MockInputsRecordingGuardrail() + request_data = self._request_data( + system=None, + tools=[ + { + "name": "get_weather", + "description": "look up weather", + "input_schema": {"type": "object", "properties": {}}, + } + ], + ) + response = { + "id": "msg_1", + "model": "claude-sonnet-4-5", + "content": [ + {"type": "text", "text": "Checking."}, + {"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "Paris"}}, + ], + } + + await handler.process_output_response( + response=response, + guardrail_to_apply=guardrail, + request_data=request_data, + ) + + _, inputs = guardrail.calls[-1] + assistant_turn = inputs["structured_messages"][-1] + assert assistant_turn["role"] == "assistant" + assert assistant_turn["content"] == "Checking." + assert assistant_turn["tool_calls"][0]["function"]["name"] == "get_weather" + assert inputs["tools"][0]["function"]["name"] == "get_weather" + + @pytest.mark.asyncio + async def test_response_scan_without_request_messages_sends_no_conversation(self): + handler = AnthropicMessagesHandler() + guardrail = MockInputsRecordingGuardrail() + response = { + "id": "msg_1", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "Hello."}], + } + + await handler.process_output_response( + response=response, + guardrail_to_apply=guardrail, + request_data=None, + ) + + _, inputs = guardrail.calls[-1] + assert "structured_messages" not in inputs + assert inputs["texts"] == ["Hello."] 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 2e75f29b1c5..8676cb4392e 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 @@ -1564,3 +1564,214 @@ class TestScanOnlyToolResults: assert data["messages"][3]["content"] == "page says [BLOCKED] here" assert data["messages"][3]["tool_call_id"] == "call_1" assert data["messages"][4]["content"] == "and then?" + + +class RecordingGuardrail(CustomGuardrail): + """Captures the inputs of every apply_guardrail call without modifying anything.""" + + def __init__(self): + super().__init__(guardrail_name="recording") + self.calls: list = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + import copy + + self.calls.append((input_type, copy.deepcopy(dict(inputs)))) + return inputs + + +class TestResponseScanConversation: + """Response scans receive the request conversation with the model's assistant + turns appended (inputs["structured_messages"]) plus the request's tool + definitions (inputs["tools"]), so conversation-context guardrails see the + same dialogue on both scans without caching state between hooks.""" + + def _response(self, *choices: Choices) -> ModelResponse: + return ModelResponse( + id="chatcmpl-1", + created=1234567890, + model="gpt-4", + object="chat.completion", + choices=list(choices), + ) + + @pytest.mark.asyncio + async def test_response_scan_receives_conversation_and_tools(self): + handler = OpenAIChatCompletionsHandler() + guardrail = RecordingGuardrail() + request_data = { + "messages": [ + {"role": "system", "content": "be helpful"}, + {"role": "user", "content": "what's the weather?"}, + ], + "tools": [{"type": "function", "function": {"name": "get_weather", "parameters": {}}}], + } + response = self._response( + Choices( + finish_reason="tool_calls", + index=0, + message=Message( + content="Checking now.", + role="assistant", + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_1", + type="function", + function=Function(name="get_weather", arguments='{"city": "Paris"}'), + ) + ], + ), + ) + ) + + await handler.process_output_response(response, guardrail, request_data=request_data) + + input_type, inputs = guardrail.calls[-1] + assert input_type == "response" + assert inputs["structured_messages"] == [ + {"role": "system", "content": "be helpful"}, + {"role": "user", "content": "what's the weather?"}, + { + "role": "assistant", + "content": "Checking now.", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + }, + ] + assert inputs["tools"] == request_data["tools"] + + @pytest.mark.asyncio + async def test_response_scan_scoping_matches_request_scan(self): + handler = OpenAIChatCompletionsHandler() + guardrail = RecordingGuardrail() + guardrail.skip_system_message_in_guardrail = True + guardrail.skip_tool_message_in_guardrail = True + request_data = { + "messages": [ + {"role": "system", "content": "SYSTEM-PROMPT"}, + {"role": "user", "content": "run the tool"}, + {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"}, + {"role": "user", "content": "summarize"}, + ] + } + + await handler.process_input_messages(data=dict(request_data), guardrail_to_apply=guardrail) + response = self._response( + Choices(finish_reason="stop", index=0, message=Message(content="Done.", role="assistant")) + ) + await handler.process_output_response(response, guardrail, request_data=dict(request_data)) + + (_, request_inputs), (_, response_inputs) = guardrail.calls + request_conversation = request_inputs["structured_messages"] + response_conversation = response_inputs["structured_messages"] + assert response_conversation[: len(request_conversation)] == request_conversation + assert response_conversation[len(request_conversation) :] == [{"role": "assistant", "content": "Done."}] + assert all(m["role"] not in ("system", "tool") for m in response_conversation) + + @pytest.mark.asyncio + async def test_scan_only_tool_results_scopes_response_conversation_and_tools(self): + handler = OpenAIChatCompletionsHandler() + guardrail = RecordingGuardrail() + guardrail.scan_only_tool_results = True + request_data = { + "messages": [ + {"role": "user", "content": "fetch the page"}, + {"role": "tool", "tool_call_id": "call_1", "content": "page content"}, + ], + "tools": [{"type": "function", "function": {"name": "fetch", "parameters": {}}}], + } + response = self._response( + Choices(finish_reason="stop", index=0, message=Message(content="Summary.", role="assistant")) + ) + + await handler.process_output_response(response, guardrail, request_data=request_data) + + _, inputs = guardrail.calls[-1] + assert inputs["structured_messages"] == [ + {"role": "tool", "tool_call_id": "call_1", "content": "page content"}, + {"role": "assistant", "content": "Summary."}, + ] + assert "tools" not in inputs + + @pytest.mark.asyncio + async def test_sdk_path_without_request_data_sends_no_conversation(self): + handler = OpenAIChatCompletionsHandler() + guardrail = RecordingGuardrail() + response = self._response( + Choices(finish_reason="stop", index=0, message=Message(content="Hello.", role="assistant")) + ) + + await handler.process_output_response(response, guardrail, request_data=None) + + _, inputs = guardrail.calls[-1] + assert "structured_messages" not in inputs + assert "tools" not in inputs + assert inputs["texts"] == ["Hello."] + + @pytest.mark.asyncio + async def test_multi_choice_response_appends_one_turn_per_choice(self): + handler = OpenAIChatCompletionsHandler() + guardrail = RecordingGuardrail() + request_data = {"messages": [{"role": "user", "content": "pick one"}]} + response = self._response( + Choices(finish_reason="stop", index=0, message=Message(content="candidate one", role="assistant")), + Choices(finish_reason="stop", index=1, message=Message(content="candidate two", role="assistant")), + ) + + await handler.process_output_response(response, guardrail, request_data=request_data) + + _, inputs = guardrail.calls[-1] + assert inputs["structured_messages"] == [ + {"role": "user", "content": "pick one"}, + {"role": "assistant", "content": "candidate one"}, + {"role": "assistant", "content": "candidate two"}, + ] + + @pytest.mark.asyncio + async def test_end_of_stream_scan_receives_conversation(self): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + handler = OpenAIChatCompletionsHandler() + guardrail = RecordingGuardrail() + request_data = {"messages": [{"role": "user", "content": "say hi"}]} + chunks = [ + ModelResponseStream( + id="chatcmpl-1", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices(index=0, finish_reason=None, delta=Delta(content="hi ", role="assistant")) + ], + ), + ModelResponseStream( + id="chatcmpl-1", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[StreamingChoices(index=0, finish_reason="stop", delta=Delta(content="there"))], + ), + ] + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + request_data=request_data, + ) + + _, inputs = guardrail.calls[-1] + assert inputs["structured_messages"] == [ + {"role": "user", "content": "say hi"}, + {"role": "assistant", "content": "hi there"}, + ] 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 4c45eaac7b9..f721eaf908e 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 @@ -1234,3 +1234,164 @@ class TestOpenAIResponsesHandlerToolInjection: names = [t.get("name") for t in result["tools"]] assert "get_weather" in names assert "injected_tool" in names + + +class MockInputsRecordingGuardrail(CustomGuardrail): + """Records the inputs of every apply_guardrail call without modifying anything.""" + + def __init__(self): + super().__init__(guardrail_name="inputs-recording") + self.calls: list = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + import copy + + self.calls.append((input_type, copy.deepcopy(dict(inputs)))) + return inputs + + +class TestResponsesResponseScanConversation: + """Response scans receive the request conversation (input plus instructions, + translated to chat shape) with the model's assistant turn appended, and the + request's tools translated to chat-completion shape.""" + + def _response(self, output: list) -> ResponsesAPIResponse: + return ResponsesAPIResponse( + id="resp_1", + created_at=1234567890, + model="gpt-4", + object="response", + status="completed", + output=output, + ) + + @pytest.mark.asyncio + async def test_response_scan_receives_conversation_and_tools(self): + handler = OpenAIResponsesHandler() + guardrail = MockInputsRecordingGuardrail() + request_data = { + "model": "gpt-4", + "instructions": "be helpful", + "input": [{"role": "user", "content": "what's the weather?"}], + "tools": [ + { + "type": "function", + "name": "get_weather", + "parameters": {"type": "object", "properties": {}}, + } + ], + } + response = self._response( + [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Sunny today."}], + } + ] + ) + + await handler.process_output_response(response, guardrail, request_data=request_data) + + input_type, inputs = guardrail.calls[-1] + assert input_type == "response" + conversation = inputs["structured_messages"] + expected_request_conversation = handler.get_structured_messages(request_data) + assert conversation[:-1] == expected_request_conversation + assert conversation[0]["role"] == "system" + assert "be helpful" in str(conversation[0]["content"]) + assert conversation[-1] == {"role": "assistant", "content": "Sunny today."} + assert inputs["tools"][0]["function"]["name"] == "get_weather" + + @pytest.mark.asyncio + async def test_response_scan_appends_tool_call_turn(self): + handler = OpenAIResponsesHandler() + guardrail = MockInputsRecordingGuardrail() + request_data = { + "model": "gpt-4", + "input": [{"role": "user", "content": "look it up"}], + } + response = self._response( + [ + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + "status": "completed", + } + ] + ) + + await handler.process_output_response(response, guardrail, request_data=request_data) + + _, inputs = guardrail.calls[-1] + assistant_turn = inputs["structured_messages"][-1] + assert assistant_turn["role"] == "assistant" + assert assistant_turn["tool_calls"][0]["function"]["name"] == "get_weather" + + @pytest.mark.asyncio + async def test_response_scan_without_request_data_sends_no_conversation(self): + handler = OpenAIResponsesHandler() + guardrail = MockInputsRecordingGuardrail() + response = self._response( + [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Hello."}], + } + ] + ) + + await handler.process_output_response(response, guardrail, request_data=None) + + _, inputs = guardrail.calls[-1] + assert "structured_messages" not in inputs + assert inputs["texts"] == ["Hello."] + + @pytest.mark.asyncio + async def test_streaming_completed_event_receives_conversation(self): + handler = OpenAIResponsesHandler() + guardrail = MockInputsRecordingGuardrail() + request_data = { + "model": "gpt-4", + "input": [{"role": "user", "content": "say hi"}], + } + final_chunk = { + "type": "response.completed", + "response": { + "model": "gpt-4", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "hi there"}], + } + ], + }, + } + + await handler.process_output_streaming_response( + responses_so_far=[final_chunk], + guardrail_to_apply=guardrail, + request_data=request_data, + ) + + _, inputs = guardrail.calls[-1] + conversation = inputs["structured_messages"] + assert conversation[-1] == {"role": "assistant", "content": "hi there"} + assert conversation[:-1] == handler.get_structured_messages(request_data) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py index 729dcb54309..3187f1d322a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py @@ -989,3 +989,49 @@ async def test_openai_moderation_initialize_guardrail_forwards_streaming_flags() assert guardrail.streaming_sampling_rate == 2 finally: litellm.logging_callback_manager._reset_all_callbacks() + + +@pytest.mark.asyncio +async def test_openai_moderation_response_scan_moderates_output_not_user_prompt(): + """On response scans the conversation is available in structured_messages, + but moderation must still target the model output carried in texts.""" + from unittest.mock import AsyncMock + + from litellm.types.utils import GenericGuardrailAPIInputs + + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + guardrail = OpenAIModerationGuardrail( + guardrail_name="test-openai-moderation", + ) + + mock_response = OpenAIModerationResponse( + id="modr-123", + model="omni-moderation-latest", + results=[ + OpenAIModerationResult( + flagged=False, + categories={}, + category_scores={}, + category_applied_input_types={}, + ) + ], + ) + + with patch.object( + guardrail, "async_make_request", new_callable=AsyncMock, return_value=mock_response + ) as mock_request: + inputs = GenericGuardrailAPIInputs( + texts=["the model's answer"], + structured_messages=[ + {"role": "user", "content": "the user's question"}, + {"role": "assistant", "content": "the model's answer"}, + ], + ) + + await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="response", + ) + + assert mock_request.call_args.kwargs["input_text"] == "the model's answer" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_akto.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_akto.py new file mode 100644 index 00000000000..5aa9fa0dee3 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_akto.py @@ -0,0 +1,31 @@ +"""Tests for the Akto guardrail's ingest payload construction.""" + +from litellm.proxy.guardrails.guardrail_hooks.akto.akto import AktoGuardrail + + +class TestBuildRequestBody: + def test_request_scan_uses_structured_messages(self): + body = AktoGuardrail.build_request_body( + inputs={ + "texts": ["hi"], + "structured_messages": [{"role": "user", "content": "hi"}], + }, + request_data={"messages": [{"role": "user", "content": "raw"}]}, + ) + assert body["messages"] == [{"role": "user", "content": "hi"}] + + def test_response_scan_ingests_request_messages_not_conversation(self): + """On response scans structured_messages carries the response turns too; + the ingested request body must stay the actual request.""" + body = AktoGuardrail.build_request_body( + inputs={ + "texts": ["the reply"], + "structured_messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "the reply"}, + ], + }, + request_data={"messages": [{"role": "user", "content": "hi"}]}, + prefer_structured_messages=False, + ) + assert body["messages"] == [{"role": "user", "content": "hi"}] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_promptguard.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_promptguard.py index efd14379ddd..0ef532f0dbc 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_promptguard.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_promptguard.py @@ -815,3 +815,46 @@ class TestPromptGuardInitializer: from litellm.types.guardrails import SupportedGuardrailIntegrations assert SupportedGuardrailIntegrations.PROMPTGUARD.value == "promptguard" + + +class TestPromptGuardResponseScanIgnoresConversation: + @pytest.mark.asyncio + async def test_response_scan_sends_texts_and_redacts_them( + self, promptguard_guardrail, mock_request_data + ): + """A response scan now carries the conversation in structured_messages, + but PromptGuard must keep scanning the response texts: sending the + conversation would make the user-role redact write-back clobber the + response texts with request content.""" + resp = _make_response( + { + "decision": "redact", + "event_id": "evt-007", + "confidence": 0.99, + "threat_type": "pii_detected", + "redacted_messages": [ + {"role": "user", "content": "Your SSN is *********"} + ], + "threats": [], + "latency_ms": 50.0, + } + ) + with patch.object( + promptguard_guardrail.async_handler, "post", return_value=resp + ) as mock_post: + result = await promptguard_guardrail.apply_guardrail( + inputs={ + "texts": ["Your SSN is 123-45-6789"], + "structured_messages": [ + {"role": "user", "content": "what's my SSN?"}, + {"role": "assistant", "content": "Your SSN is 123-45-6789"}, + ], + }, + request_data=mock_request_data, + input_type="response", + ) + sent_messages = mock_post.call_args.kwargs["json"]["messages"] + assert sent_messages == [ + {"role": "user", "content": "Your SSN is 123-45-6789"} + ] + assert result["texts"] == ["Your SSN is *********"]