mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
feat(guardrails): give response scans the full conversation
Post-call guardrail scans only ever received the extracted response texts, so a guardrail judging a reply had no way to see what was asked. Guardrails wanting that context had to cache request state between the two hooks, which does not survive streaming, retries, or more than one proxy worker. Response scans now carry the conversation on inputs["structured_messages"]: the request turns as the request scan saw them, with the same skip_system, skip_tool and scan_only_tool_results scoping, and the model's reply appended. The request's tool definitions ride along on inputs["tools"]. Chat Completions, Anthropic Messages and Responses are all covered, streaming included. Guardrails that already read structured_messages as the thing to scan keep their previous response-scan input, so none of them starts scanning the user's prompt in place of the model's answer.
This commit is contained in:
Youlian Simidjiyski
parent
aae36f4bd4
commit
d338fb4fb4
19 changed files with 1280 additions and 39 deletions
|
|
@ -111,6 +111,17 @@ class ExtractedInput:
|
|||
EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=())
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ScopedRequestView:
|
||||
"""The request as guardrail scans see it, with the operator scoping applied."""
|
||||
|
||||
translated_request: ChatCompletionRequest
|
||||
full_messages: tuple[AllMessageValues, ...]
|
||||
scoped_indices: tuple[int, ...]
|
||||
hoisted_system_message: AllMessageValues | None
|
||||
has_midturn_system_message: bool
|
||||
|
||||
|
||||
class AnthropicMessagesHandler(BaseTranslation):
|
||||
"""Process Anthropic messages with guardrails.
|
||||
|
||||
|
|
@ -314,6 +325,78 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
)
|
||||
return result if result else None
|
||||
|
||||
def _scoped_request_view(
|
||||
self,
|
||||
data: dict, # mutable-ok: API request payload
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
) -> "_ScopedRequestView":
|
||||
"""
|
||||
Translate the request the way guardrail scans see it: exclude the
|
||||
trusted top-level prompt, hoist it back unless skip_system, keep
|
||||
in-sequence system entries in scope. Single source of truth for the
|
||||
request and response scans, which must not disagree.
|
||||
"""
|
||||
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 data.items() if key != "system"},
|
||||
"model": data.get("model") or "",
|
||||
}
|
||||
translated_request: Final = self._translate_to_openai(translation_source)
|
||||
translated_messages: Final = tuple(translated_request.get("messages") or ())
|
||||
has_midturn_system_message: Final = any(
|
||||
str(message.get("role") or "").lower() == "system" for message in translated_messages
|
||||
)
|
||||
hoisted_system_message: Final = None if skip_system else self._hoisted_top_level_system_message(data)
|
||||
full_messages: Final = (
|
||||
(hoisted_system_message, *translated_messages)
|
||||
if hoisted_system_message is not None
|
||||
else translated_messages
|
||||
)
|
||||
# skip_system already excluded the trusted top-level prompt (it is simply
|
||||
# not hoisted); in-sequence system entries are untrusted and stay in scope.
|
||||
scoped_indices: Final = scoped_structured_message_indices(
|
||||
full_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),
|
||||
)
|
||||
return _ScopedRequestView(
|
||||
translated_request=translated_request,
|
||||
full_messages=full_messages,
|
||||
scoped_indices=scoped_indices,
|
||||
hoisted_system_message=hoisted_system_message,
|
||||
has_midturn_system_message=has_midturn_system_message,
|
||||
)
|
||||
|
||||
def scoped_request_conversation(
|
||||
self,
|
||||
request_data: dict, # mutable-ok: API request payload
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
) -> tuple[AllMessageValues, ...] | None:
|
||||
if request_data.get("messages") is None:
|
||||
return None
|
||||
view: Final = self._scoped_request_view(request_data, guardrail_to_apply)
|
||||
return tuple(view.full_messages[index] for index in view.scoped_indices) or None
|
||||
|
||||
def request_tools_for_guardrail(
|
||||
self,
|
||||
request_data: dict, # mutable-ok: API request payload
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
) -> tuple[ChatCompletionToolParam, ...] | None:
|
||||
if effective_scan_only_tool_results_for_guardrail(guardrail_to_apply):
|
||||
return None
|
||||
if not request_data.get("tools"):
|
||||
return None
|
||||
probe: Final = self._translate_to_openai(
|
||||
{ # mutable-ok: API message payload
|
||||
"model": request_data.get("model") or "",
|
||||
"messages": [], # mutable-ok: API message payload
|
||||
"tools": request_data["tools"],
|
||||
}
|
||||
)
|
||||
tools: Final = probe.get("tools")
|
||||
return tuple(tools) if tools else None
|
||||
|
||||
async def process_input_messages(
|
||||
self,
|
||||
data: dict,
|
||||
|
|
@ -334,33 +417,17 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
# Exclude only the trusted top-level prompt. In-sequence system entries are untrusted
|
||||
# and must stay aligned with texts_to_check for positional masking. When the top-level
|
||||
# prompt is included, the pre-existing count mismatch disables positional masking.
|
||||
translation_source: Final = { # mutable-ok: API message payload
|
||||
key: value for key, value in data.items() if key != "system"
|
||||
}
|
||||
chat_completion_compatible_request: Final = self._translate_to_openai(translation_source)
|
||||
view: Final = self._scoped_request_view(data, guardrail_to_apply)
|
||||
full_structured_messages: Final = view.full_messages
|
||||
has_midturn_system_message: Final = view.has_midturn_system_message
|
||||
hoisted_system_message: Final = view.hoisted_system_message
|
||||
scoped_message_indices: Final = view.scoped_indices
|
||||
structured_messages: Final = [ # mutable-ok: GenericGuardrailAPIInputs takes list
|
||||
full_structured_messages[index] for index in scoped_message_indices
|
||||
]
|
||||
|
||||
full_structured_messages: Final = cast(
|
||||
list[AllMessageValues],
|
||||
chat_completion_compatible_request.get("messages", []),
|
||||
)
|
||||
has_midturn_system_message: Final = any(
|
||||
str(message.get("role") or "").lower() == "system" for message in full_structured_messages
|
||||
)
|
||||
hoisted_system_message: Final = None if skip_system else self._hoisted_top_level_system_message(data)
|
||||
if hoisted_system_message is not None:
|
||||
full_structured_messages.insert(0, hoisted_system_message)
|
||||
# skip_system already excluded the trusted top-level prompt (it is simply not hoisted);
|
||||
# in-sequence system entries are untrusted and always stay in scope.
|
||||
scoped_message_indices: Final = scoped_structured_message_indices(
|
||||
full_structured_messages,
|
||||
scan_only_tool_results=scan_only_tool_results,
|
||||
skip_system=False,
|
||||
skip_tool=skip_tool,
|
||||
)
|
||||
structured_messages: Final = [full_structured_messages[index] for index in scoped_message_indices]
|
||||
|
||||
tools_to_check: Final[list[ChatCompletionToolParam]] = (
|
||||
[] if scan_only_tool_results else chat_completion_compatible_request.get("tools", [])
|
||||
tools_to_check: Final[list[ChatCompletionToolParam]] = ( # mutable-ok: GenericGuardrailAPIInputs takes list
|
||||
[] if scan_only_tool_results else list(view.translated_request.get("tools") or ())
|
||||
)
|
||||
|
||||
# Step 1: Extract all text content and images
|
||||
|
|
@ -896,6 +963,13 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
response,
|
||||
)
|
||||
|
||||
self.attach_response_scan_context(
|
||||
inputs,
|
||||
request_data,
|
||||
guardrail_to_apply,
|
||||
self.assistant_turn_from_extraction(texts_to_check, tool_calls_to_check),
|
||||
)
|
||||
|
||||
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=request_data,
|
||||
|
|
@ -962,6 +1036,23 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
user_api_key_dict,
|
||||
key="response",
|
||||
)
|
||||
stream_tool_calls: Final = tuple(
|
||||
self.assistant_tool_call(
|
||||
tool_call_id=tool_call.id,
|
||||
name=tool_call.function.name,
|
||||
arguments=tool_call.function.arguments,
|
||||
)
|
||||
for tool_call in tool_calls_list or ()
|
||||
)
|
||||
self.attach_response_scan_context(
|
||||
guardrail_inputs,
|
||||
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_calls,
|
||||
),
|
||||
)
|
||||
_guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=guardrail_inputs,
|
||||
request_data=prepared_request_data,
|
||||
|
|
|
|||
|
|
@ -1,7 +1,20 @@
|
|||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any, Final, Optional
|
||||
|
||||
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,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
ChatCompletionAssistantMessage,
|
||||
ChatCompletionAssistantToolCall,
|
||||
ChatCompletionToolCallFunctionChunk,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
|
|
@ -9,7 +22,12 @@ 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,
|
||||
ChatCompletionToolCallChunk,
|
||||
ChatCompletionToolParam,
|
||||
)
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
|
|
@ -157,6 +175,140 @@ class BaseTranslation(ABC):
|
|||
"""
|
||||
return None
|
||||
|
||||
def scoped_request_conversation(
|
||||
self,
|
||||
request_data: dict, # mutable-ok: API request payload
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
) -> tuple["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),
|
||||
)
|
||||
return tuple(structured_messages[index] for index in scoped_indices) or None
|
||||
|
||||
def response_scan_conversation(
|
||||
self,
|
||||
request_data: dict | None, # mutable-ok: API request payload
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
response_turns: Sequence["AllMessageValues"],
|
||||
) -> tuple["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) or when nothing was
|
||||
extracted from the response; guardrails then fall back to scanning the
|
||||
extracted texts and tool calls, which must not be shadowed by a
|
||||
conversation that lacks a response turn.
|
||||
"""
|
||||
if request_data is None or not response_turns:
|
||||
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 attach_response_scan_context(
|
||||
self,
|
||||
inputs: "GenericGuardrailAPIInputs",
|
||||
request_data: dict | None, # mutable-ok: API request payload
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
response_turns: Sequence["AllMessageValues"],
|
||||
) -> None:
|
||||
"""
|
||||
Put the response-scan conversation and the request's tools on ``inputs``
|
||||
when the request context allows building them; no-op otherwise.
|
||||
"""
|
||||
structured_conversation: Final = self.response_scan_conversation(
|
||||
request_data, guardrail_to_apply, response_turns
|
||||
)
|
||||
if structured_conversation is None or request_data is None:
|
||||
return
|
||||
inputs["structured_messages"] = list(structured_conversation) # rebind-ok: out-param; field type is list
|
||||
response_scan_tools: Final = self.request_tools_for_guardrail(request_data, guardrail_to_apply)
|
||||
if response_scan_tools:
|
||||
inputs["tools"] = list(response_scan_tools) # rebind-ok: out-parameter; the field type is a list
|
||||
|
||||
def request_tools_for_guardrail(
|
||||
self,
|
||||
request_data: dict, # mutable-ok: API request payload
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
) -> tuple["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_tool_call(
|
||||
tool_call_id: str | None,
|
||||
name: str | None,
|
||||
arguments: str,
|
||||
) -> ChatCompletionAssistantToolCall:
|
||||
"""One assistant-message tool call in the OpenAI wire shape."""
|
||||
return ChatCompletionAssistantToolCall(
|
||||
id=tool_call_id,
|
||||
type="function",
|
||||
function=ChatCompletionToolCallFunctionChunk(name=name, arguments=arguments),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def assistant_turn(
|
||||
content: str | None,
|
||||
tool_calls: Sequence[ChatCompletionAssistantToolCall],
|
||||
) -> ChatCompletionAssistantMessage | None:
|
||||
"""One OpenAI-shape assistant turn, or None when it would carry nothing."""
|
||||
if content is None and not tool_calls:
|
||||
return None
|
||||
if tool_calls:
|
||||
turn_with_tools: Final[ChatCompletionAssistantMessage] = {
|
||||
"role": "assistant",
|
||||
"content": content,
|
||||
"tool_calls": list(tool_calls), # mutable-ok: the message field type is a list
|
||||
}
|
||||
return turn_with_tools
|
||||
turn: Final[ChatCompletionAssistantMessage] = {"role": "assistant", "content": content}
|
||||
return turn
|
||||
|
||||
@staticmethod
|
||||
def assistant_turn_from_extraction(
|
||||
texts: Sequence[str],
|
||||
tool_calls: Sequence["ChatCompletionAssistantToolCall | ChatCompletionToolCallChunk"] | None = None,
|
||||
) -> tuple["ChatCompletionAssistantMessage", ...]:
|
||||
"""
|
||||
One OpenAI-shape assistant turn built from the texts and tool calls a
|
||||
handler's response extraction collected; empty when there is nothing.
|
||||
Tool calls are normalized to the assistant-message shape, dropping
|
||||
extraction-only fields such as ``index``.
|
||||
"""
|
||||
tool_call_items: Final = tuple(
|
||||
BaseTranslation.assistant_tool_call(
|
||||
tool_call_id=item.get("id"),
|
||||
name=item["function"].get("name"),
|
||||
arguments=item["function"].get("arguments") or "",
|
||||
)
|
||||
for item in tool_calls or ()
|
||||
)
|
||||
if not texts and not tool_call_items:
|
||||
return ()
|
||||
turn: Final = BaseTranslation.assistant_turn("\n".join(texts), tool_call_items)
|
||||
return (turn,) if turn is not None else ()
|
||||
|
||||
def extract_request_tool_names(self, data: dict) -> list[str]:
|
||||
"""
|
||||
Extract tool names from the request body for allowlist/policy checks.
|
||||
|
|
|
|||
|
|
@ -33,7 +33,11 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
|
|||
scoped_structured_message_indices,
|
||||
)
|
||||
from litellm.main import stream_chunk_builder
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
ChatCompletionAssistantMessage,
|
||||
ChatCompletionToolParam,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
coerce_stream_holdback_value,
|
||||
)
|
||||
|
|
@ -396,6 +400,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
if hasattr(response, "model") and response.model:
|
||||
inputs["model"] = response.model
|
||||
|
||||
self.attach_response_scan_context(
|
||||
inputs, request_data, guardrail_to_apply, self._build_response_turns(response)
|
||||
)
|
||||
|
||||
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=request_data,
|
||||
|
|
@ -566,6 +574,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
# Include model information from the first response if available
|
||||
if responses_so_far and hasattr(responses_so_far[0], "model") and responses_so_far[0].model:
|
||||
inputs["model"] = responses_so_far[0].model
|
||||
self.attach_response_scan_context(
|
||||
inputs, request_data, guardrail_to_apply, self.assistant_turn_from_extraction(texts_to_check)
|
||||
)
|
||||
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=request_data,
|
||||
|
|
@ -660,6 +671,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
|
||||
if responses_so_far and getattr(responses_so_far[0], "model", None):
|
||||
inputs["model"] = responses_so_far[0].model
|
||||
self.attach_response_scan_context(
|
||||
inputs,
|
||||
request_data,
|
||||
guardrail_to_apply,
|
||||
tuple(turn for index in indices if (turn := self.assistant_turn(raw_by_index[index], ())) is not None),
|
||||
)
|
||||
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=request_data,
|
||||
|
|
@ -832,6 +849,32 @@ 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, # mutable-ok: API request payload
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
) -> tuple[ChatCompletionToolParam, ...] | None:
|
||||
if effective_scan_only_tool_results_for_guardrail(guardrail_to_apply):
|
||||
return None
|
||||
tools: Final = request_data.get("tools")
|
||||
return tuple(tools) if tools else None
|
||||
|
||||
def _build_response_turns(self, response: "ModelResponse") -> tuple[ChatCompletionAssistantMessage, ...]:
|
||||
"""Assistant turns for the response-scan conversation, one per choice."""
|
||||
return tuple(turn for choice in response.choices if (turn := self._choice_assistant_turn(choice)) is not None)
|
||||
|
||||
def _choice_assistant_turn(self, choice: Choices) -> ChatCompletionAssistantMessage | None:
|
||||
tool_calls: Final = tuple(
|
||||
self.assistant_tool_call(
|
||||
tool_call_id=converted.get("id"),
|
||||
name=(converted.get("function") or {}).get("name"),
|
||||
arguments=(converted.get("function") or {}).get("arguments") or "",
|
||||
)
|
||||
for tool_call in choice.message.tool_calls or ()
|
||||
if (converted := self._convert_tool_call_to_dict(tool_call)) is not None and converted.get("function")
|
||||
)
|
||||
return self.assistant_turn(choice.message.content, tool_calls)
|
||||
|
||||
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,32 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools)
|
||||
tools_to_check.extend(cast(list[ChatCompletionToolParam], transformed_tools))
|
||||
|
||||
def scoped_request_conversation(
|
||||
self,
|
||||
request_data: dict, # mutable-ok: API request payload
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
) -> tuple[AllMessageValues, ...] | None:
|
||||
"""
|
||||
This surface's request scan sends its structured messages without the
|
||||
operator scoping flags (the input path applies none), so the response
|
||||
scan mirrors that; request and response scans must not disagree about
|
||||
the conversation.
|
||||
"""
|
||||
structured_messages: Final = self.get_structured_messages(request_data)
|
||||
return tuple(structured_messages) if structured_messages else None
|
||||
|
||||
def request_tools_for_guardrail(
|
||||
self,
|
||||
request_data: dict, # mutable-ok: API request payload
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
) -> tuple[ChatCompletionToolParam, ...] | None:
|
||||
raw_tools: Final = request_data.get("tools")
|
||||
if not raw_tools:
|
||||
return None
|
||||
tools_to_check: Final[list[ChatCompletionToolParam]] = [] # mutable-ok: _extract_and_transform_tools appends
|
||||
self._extract_and_transform_tools(raw_tools, tools_to_check)
|
||||
return tuple(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 +481,13 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
if response_model:
|
||||
inputs["model"] = response_model
|
||||
|
||||
self.attach_response_scan_context(
|
||||
inputs,
|
||||
request_data,
|
||||
guardrail_to_apply,
|
||||
self.assistant_turn_from_extraction(texts_to_check, tool_calls_to_check),
|
||||
)
|
||||
|
||||
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=request_data,
|
||||
|
|
@ -545,6 +578,13 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
if response_model:
|
||||
inputs["model"] = response_model
|
||||
|
||||
self.attach_response_scan_context(
|
||||
inputs,
|
||||
request_data,
|
||||
guardrail_to_apply,
|
||||
self.assistant_turn_from_extraction(texts_to_check, tool_calls_to_check),
|
||||
)
|
||||
|
||||
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
|
||||
|
|
|
|||
|
|
@ -211,7 +211,9 @@ class HiddenlayerGuardrail(CustomGuardrail):
|
|||
hl_request_metadata["requester_id"] = headers.get("hl-requester-id") or "LiteLLM"
|
||||
project_id: Final = headers.get("hl-project-id")
|
||||
|
||||
if scan_params := inputs.get("structured_messages"):
|
||||
# Response scans keep the texts path: their structured_messages carry the
|
||||
# whole conversation, whose last turn is not necessarily the scan target.
|
||||
if input_type == "request" and (scan_params := inputs.get("structured_messages")):
|
||||
last_msg: Final = scan_params[-1]
|
||||
result: _HiddenlayerResponse = await self._call_hiddenlayer(
|
||||
project_id,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -424,8 +424,10 @@ class QualifireGuardrail(CustomGuardrail):
|
|||
# Get dynamic params from request body (allows runtime overrides)
|
||||
dynamic_params: Final = self.get_guardrail_dynamic_request_body_params(request_data=request_data)
|
||||
|
||||
# Extract messages from structured_messages or request_data
|
||||
messages: list[AllMessageValues] | None = inputs.get("structured_messages")
|
||||
# Extract messages from structured_messages or request_data. Response
|
||||
# scans keep the request_data path: their structured_messages carry the
|
||||
# model's answer too, which would land in `messages` on top of `output`.
|
||||
messages: list[AllMessageValues] | None = inputs.get("structured_messages") if input_type == "request" else None
|
||||
if not messages:
|
||||
messages = request_data.get("messages")
|
||||
|
||||
|
|
|
|||
|
|
@ -380,10 +380,15 @@ class StraikerGuardrail(CustomGuardrail):
|
|||
call_id: Final = getattr(logging_obj, "litellm_call_id", None) if logging_obj else None
|
||||
event_id: Final = f"{call_id or 'litellm'}:{input_type}"
|
||||
|
||||
# On response scans this content object is the envelope's ``response``,
|
||||
# while ``structured_messages`` carries the whole conversation; it is
|
||||
# reported under ``request`` instead so neither side misattributes turns.
|
||||
content: Final = StraikerWebhookContent(
|
||||
texts=list(inputs.get("texts") or []),
|
||||
images=list(inputs.get("images") or []),
|
||||
structured_messages=_opaque_dict_list(inputs.get("structured_messages")),
|
||||
structured_messages=(
|
||||
_opaque_dict_list(inputs.get("structured_messages")) if input_type == "request" else None
|
||||
),
|
||||
tools=_opaque_dict_list(inputs.get("tools")),
|
||||
tool_calls=_opaque_dict_list(inputs.get("tool_calls")),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1818,3 +1818,204 @@ 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[tuple[str, dict]] = []
|
||||
|
||||
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_skip_system_scoping_matches_request_scan(self):
|
||||
handler = AnthropicMessagesHandler()
|
||||
guardrail = MockInputsRecordingGuardrail()
|
||||
guardrail.skip_system_message_in_guardrail = True
|
||||
request_data = self._request_data(
|
||||
system="trusted top-level system prompt",
|
||||
messages=[
|
||||
{"role": "user", "content": "safe text"},
|
||||
{"role": "system", "content": "prohibited correction"},
|
||||
{"role": "user", "content": "continue"},
|
||||
],
|
||||
)
|
||||
response = {
|
||||
"id": "msg_1",
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [{"type": "text", "text": "Done."}],
|
||||
}
|
||||
|
||||
await handler.process_input_messages(data=dict(request_data), guardrail_to_apply=guardrail)
|
||||
await handler.process_output_response(
|
||||
response=response,
|
||||
guardrail_to_apply=guardrail,
|
||||
request_data=dict(request_data),
|
||||
)
|
||||
|
||||
(request_type, request_inputs), (response_type, response_inputs) = guardrail.calls
|
||||
assert (request_type, response_type) == ("request", "response")
|
||||
request_conversation = request_inputs["structured_messages"]
|
||||
response_conversation = response_inputs["structured_messages"]
|
||||
assert [message["role"] for message in request_conversation] == ["user", "system", "user"]
|
||||
assert response_conversation[: len(request_conversation)] == request_conversation
|
||||
assert response_conversation[len(request_conversation) :] == [{"role": "assistant", "content": "Done."}]
|
||||
assert "trusted top-level system prompt" not in str(response_conversation)
|
||||
|
||||
@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 "index" not in assistant_turn["tool_calls"][0], (
|
||||
"extraction-only fields must not leak into the assistant-message tool call shape"
|
||||
)
|
||||
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."]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_scan_survives_request_data_without_model(self):
|
||||
"""Buffered streaming hands the response scan a request payload with no
|
||||
``model``; the conversation must still be built rather than blowing up
|
||||
inside the Anthropic-to-OpenAI adapter."""
|
||||
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={"messages": [{"role": "user", "content": "what's the weather?"}]},
|
||||
)
|
||||
|
||||
_, inputs = guardrail.calls[-1]
|
||||
assert inputs["structured_messages"] == [
|
||||
{"role": "user", "content": "what's the weather?"},
|
||||
{"role": "assistant", "content": "Sunny today."},
|
||||
]
|
||||
|
|
|
|||
|
|
@ -1559,3 +1559,300 @@ 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[tuple[str, dict]] = []
|
||||
|
||||
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"},
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_midstream_scan_receives_the_partial_assistant_turn(self):
|
||||
"""Before the stream ends there is no assembled response to fall back on,
|
||||
so the in-flight scan must still see the conversation plus what the model
|
||||
has produced so far."""
|
||||
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=None, delta=Delta(content="the"))],
|
||||
),
|
||||
]
|
||||
|
||||
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 the"},
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_transform_scan_receives_conversation_per_choice(self):
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
||||
|
||||
handler = OpenAIChatCompletionsHandler()
|
||||
guardrail = RecordingGuardrail()
|
||||
request_data = {"messages": [{"role": "user", "content": "pick one"}]}
|
||||
chunks = [
|
||||
ModelResponseStream(
|
||||
id="chatcmpl-1",
|
||||
created=1234567890,
|
||||
model="gpt-4",
|
||||
object="chat.completion.chunk",
|
||||
choices=[
|
||||
StreamingChoices(index=1, finish_reason=None, delta=Delta(content="two", role="assistant")),
|
||||
StreamingChoices(index=0, finish_reason=None, delta=Delta(content="one", role="assistant")),
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
await handler.process_output_streaming_response(
|
||||
responses_so_far=chunks,
|
||||
guardrail_to_apply=guardrail,
|
||||
request_data=request_data,
|
||||
stream_transform_sink=StreamTransformSink(),
|
||||
)
|
||||
|
||||
_, inputs = guardrail.calls[-1]
|
||||
assert inputs["structured_messages"] == [
|
||||
{"role": "user", "content": "pick one"},
|
||||
{"role": "assistant", "content": "one"},
|
||||
{"role": "assistant", "content": "two"},
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_response_turns_yield_no_conversation(self):
|
||||
"""A conversation without a response turn would shadow the tool_calls/texts
|
||||
fallback in conversation-preferring guardrails, so none is built."""
|
||||
handler = OpenAIChatCompletionsHandler()
|
||||
guardrail = RecordingGuardrail()
|
||||
request_data = {"messages": [{"role": "user", "content": "hi"}]}
|
||||
|
||||
assert handler.response_scan_conversation(request_data, guardrail, []) is None
|
||||
|
|
|
|||
|
|
@ -1229,3 +1229,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[tuple[str, dict]] = []
|
||||
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -987,3 +987,50 @@ 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. The fake
|
||||
moderation endpoint flags only the harmful model answer: moderating the
|
||||
(benign) user prompt instead would let the flagged output through."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}):
|
||||
guardrail = OpenAIModerationGuardrail(
|
||||
guardrail_name="test-openai-moderation",
|
||||
)
|
||||
|
||||
async def moderate(input_text: str) -> OpenAIModerationResponse:
|
||||
flagged = "harmful answer" in input_text
|
||||
return OpenAIModerationResponse(
|
||||
id="modr-123",
|
||||
model="omni-moderation-latest",
|
||||
results=[
|
||||
OpenAIModerationResult(
|
||||
flagged=flagged,
|
||||
categories={"violence": flagged},
|
||||
category_scores={"violence": 0.99 if flagged else 0.0},
|
||||
category_applied_input_types={"violence": []},
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
with patch.object(guardrail, "async_make_request", side_effect=moderate):
|
||||
inputs = GenericGuardrailAPIInputs(
|
||||
texts=["a harmful answer"],
|
||||
structured_messages=[
|
||||
{"role": "user", "content": "a benign question"},
|
||||
{"role": "assistant", "content": "a harmful answer"},
|
||||
],
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException):
|
||||
await guardrail.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data={},
|
||||
input_type="response",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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"}]
|
||||
|
|
@ -1088,3 +1088,37 @@ class TestHiddenlayerGuardrailV2:
|
|||
config_model = HiddenlayerGuardrailV2.get_config_model()
|
||||
assert config_model is not None
|
||||
assert config_model.__name__ == "HiddenlayerGuardrailConfigModel"
|
||||
|
||||
|
||||
class TestHiddenlayerResponseScanConversation:
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_scan_ignores_structured_messages(self, monkeypatch):
|
||||
"""Response scans now carry the conversation in structured_messages; the
|
||||
v1 guardrail must keep scanning texts instead of str()-coercing the last
|
||||
conversation turn (a tool-call-only turn has content None, and scanning
|
||||
the literal string 'None' would produce junk verdicts)."""
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
guardrail = HiddenlayerGuardrail(
|
||||
guardrail_name="hiddenlayer", event_hook="post_call", default_on=True
|
||||
)
|
||||
|
||||
async def scan(project_id, metadata, payload, input_type):
|
||||
if payload["messages"][-1]["content"] == "None":
|
||||
return {"evaluation": {"action": "BLOCK"}, "analysis": []}
|
||||
return {}
|
||||
|
||||
inputs = GenericGuardrailAPIInputs(
|
||||
structured_messages=[
|
||||
{"role": "user", "content": "look this up"},
|
||||
{"role": "assistant", "content": None},
|
||||
]
|
||||
)
|
||||
|
||||
with patch.object(guardrail, "_call_hiddenlayer", side_effect=scan):
|
||||
result = await guardrail.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data={},
|
||||
input_type="response",
|
||||
)
|
||||
|
||||
assert result is inputs
|
||||
|
|
|
|||
|
|
@ -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 *********"]
|
||||
|
|
|
|||
|
|
@ -679,3 +679,49 @@ class TestQualifireGuardrailRegistry:
|
|||
|
||||
assert "qualifire" in guardrail_class_registry
|
||||
assert guardrail_class_registry["qualifire"] == QualifireGuardrail
|
||||
|
||||
|
||||
class TestQualifireResponseScanConversation:
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_scan_keeps_request_messages_out_of_conversation(self):
|
||||
"""Response scans now receive the whole conversation in structured_messages;
|
||||
Qualifire must keep sending the request messages plus a separate output,
|
||||
not a conversation that already embeds the model's answer."""
|
||||
from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import (
|
||||
QualifireGuardrail,
|
||||
)
|
||||
|
||||
guardrail = QualifireGuardrail(
|
||||
api_key="test_key",
|
||||
prompt_injections=True,
|
||||
guardrail_name="test_guardrail",
|
||||
)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"score": 100,
|
||||
"status": "completed",
|
||||
"evaluationResults": [],
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
guardrail.async_handler.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
request_messages = [{"role": "user", "content": "what's my balance?"}]
|
||||
inputs = {
|
||||
"texts": ["Your balance is $5"],
|
||||
"structured_messages": request_messages
|
||||
+ [{"role": "assistant", "content": "Your balance is $5"}],
|
||||
}
|
||||
|
||||
result = await guardrail.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data={"messages": request_messages},
|
||||
input_type="response",
|
||||
)
|
||||
|
||||
assert result is inputs
|
||||
payload = guardrail.async_handler.post.call_args[1]["json"]
|
||||
assert all(m.get("role") != "assistant" for m in payload["messages"]), (
|
||||
"the model's answer belongs in `output`, not in the conversation"
|
||||
)
|
||||
assert payload["output"] == "Your balance is $5"
|
||||
|
|
|
|||
|
|
@ -595,6 +595,37 @@ async def test_non_streamed_response_intervention_redacts():
|
|||
assert out["texts"] == ["[redacted]"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_envelope_keeps_conversation_out_of_response_content():
|
||||
g = _make_guardrail()
|
||||
g.async_handler.post.return_value = _mock_response("NONE")
|
||||
response = ModelResponse(
|
||||
choices=[Choices(finish_reason="stop", index=0, message=Message(content="answer", role="assistant"))],
|
||||
model="gpt-4o-mini",
|
||||
)
|
||||
conversation = [
|
||||
{"role": "user", "content": "question"},
|
||||
{"role": "assistant", "content": "answer"},
|
||||
]
|
||||
request_data = {
|
||||
"model": "gpt-4o-mini",
|
||||
"messages": [{"role": "user", "content": "question"}],
|
||||
"response": response,
|
||||
}
|
||||
|
||||
await g.apply_guardrail(
|
||||
inputs={"texts": ["answer"], "structured_messages": conversation, "model": "gpt-4o-mini"},
|
||||
request_data=request_data,
|
||||
input_type="response",
|
||||
logging_obj=_logging_obj(),
|
||||
)
|
||||
|
||||
payload = _posted_payload(g)
|
||||
assert payload["response"]["texts"] == ["answer"]
|
||||
assert "structured_messages" not in payload["response"]
|
||||
assert payload["request"]["structured_messages"] == [{"role": "user", "content": "question"}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_guardrail_intervened_without_texts_blocks():
|
||||
g = _make_guardrail()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue