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.
This commit is contained in:
Youlian Simidjiyski 2026-08-22 08:23:07 -04:00
parent 822e3557be
commit 7966666e6e
13 changed files with 893 additions and 9 deletions

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -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"},
]

View file

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

View file

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

View file

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

View file

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