mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
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:
parent
822e3557be
commit
7966666e6e
13 changed files with 893 additions and 9 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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."]
|
||||
|
|
|
|||
|
|
@ -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"},
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"}]
|
||||
|
|
@ -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 *********"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue