Skip to content
Merged
41 changes: 32 additions & 9 deletions litellm/llms/anthropic/chat/guardrail_translation/handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,10 +26,11 @@
)
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
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,
openai_messages_without_system,
openai_messages_without_tool,
merge_guardrailed_scoped_messages,
scoped_structured_message_indices,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
Expand Down Expand Up @@ -326,19 +327,25 @@ async def process_input_messages(

skip_system: Final = effective_skip_system_message_for_guardrail(guardrail_to_apply)
skip_tool: Final = effective_skip_tool_message_for_guardrail(guardrail_to_apply)
scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply)

chat_completion_compatible_request: Final = self._translate_to_openai(data)

structured_messages = cast(
full_structured_messages: Final = cast(
list[AllMessageValues],
chat_completion_compatible_request.get("messages", []),
)
if skip_system:
structured_messages = openai_messages_without_system(structured_messages)
if skip_tool:
structured_messages = openai_messages_without_tool(structured_messages)
scoped_message_indices: Final = scoped_structured_message_indices(
full_structured_messages,
scan_only_tool_results=scan_only_tool_results,
skip_system=skip_system,
skip_tool=skip_tool,
)
structured_messages: Final = [full_structured_messages[index] for index in scoped_message_indices]

tools_to_check: Final[list[ChatCompletionToolParam]] = chat_completion_compatible_request.get("tools", [])
tools_to_check: Final[list[ChatCompletionToolParam]] = (
[] if scan_only_tool_results else chat_completion_compatible_request.get("tools", [])
)

# Step 1: Extract all text content and images
extracted: Final = tuple(
Expand All @@ -347,6 +354,7 @@ async def process_input_messages(
msg_idx=msg_idx,
skip_system_message=skip_system,
skip_tool_message=skip_tool,
scan_only_tool_results=scan_only_tool_results,
)
for msg_idx, message in enumerate(messages)
)
Expand Down Expand Up @@ -395,7 +403,14 @@ async def process_input_messages(
guardrailed_structured_messages is not None
and guardrailed_structured_messages is not original_structured_messages
):
self._write_back_structured_messages(data, guardrailed_structured_messages)
self._write_back_structured_messages(
data,
merge_guardrailed_scoped_messages(
full_messages=full_structured_messages,
scoped_indices=scoped_message_indices,
guardrailed_scoped=guardrailed_structured_messages,
),
)
else:
# Step 3: Map guardrail responses back to original message structure
await self._apply_guardrail_responses_to_input(
Expand Down Expand Up @@ -461,6 +476,7 @@ def _extract_input_text_and_images(
msg_idx: int,
skip_system_message: bool = False,
skip_tool_message: bool = False,
scan_only_tool_results: bool = False,
) -> ExtractedInput:
"""
Extract text content and images from a message.
Expand All @@ -471,6 +487,8 @@ def _extract_input_text_and_images(

content: Final = message.get("content", None)
if isinstance(content, str):
if scan_only_tool_results:
return EMPTY_EXTRACTED_INPUT
return ExtractedInput(scanned=(ScannedText(content, MessageContentTarget(msg_idx)),), images=())
if not isinstance(content, list):
return EMPTY_EXTRACTED_INPUT
Expand All @@ -481,6 +499,7 @@ def _extract_input_text_and_images(
msg_idx=msg_idx,
content_idx=content_idx,
skip_tool_message=skip_tool_message,
scan_only_tool_results=scan_only_tool_results,
)
for content_idx, content_item in enumerate(content)
if isinstance(content_item, dict)
Expand All @@ -497,12 +516,16 @@ def _extract_content_block(
msg_idx: int,
content_idx: int,
skip_tool_message: bool,
scan_only_tool_results: bool = False,
) -> ExtractedInput:
if content_item.get("type") == "tool_result":
if skip_tool_message:
return EMPTY_EXTRACTED_INPUT
return cls._extract_tool_result(content_item=content_item, msg_idx=msg_idx, content_idx=content_idx)

if scan_only_tool_results:
return EMPTY_EXTRACTED_INPUT

text_str: Final = content_item.get("text", None)
return ExtractedInput(
scanned=(
Expand Down
85 changes: 80 additions & 5 deletions litellm/llms/base_llm/guardrail_translation/utils.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

import json
from collections.abc import Iterator, Sequence
from typing import Any, Final

from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage
Expand Down Expand Up @@ -113,13 +114,87 @@ def effective_skip_tool_message_for_guardrail(guardrail_to_apply: Any) -> bool:
return bool(getattr(litellm, "skip_tool_message_in_guardrail", False))


def _message_role(message: AllMessageValues) -> str:
return str((message or {}).get("role") or "").lower()


def openai_messages_without_system(
messages: list[AllMessageValues],
) -> list[AllMessageValues]:
return [m for m in messages if str((m or {}).get("role") or "").lower() != "system"]
messages: Sequence[AllMessageValues],
) -> tuple[AllMessageValues, ...]:
return tuple(m for m in messages if _message_role(m) != "system")


def openai_messages_without_tool(
messages: list[AllMessageValues],
messages: Sequence[AllMessageValues],
) -> tuple[AllMessageValues, ...]:
return tuple(m for m in messages if _message_role(m) != "tool")


def effective_scan_only_tool_results_for_guardrail(guardrail_to_apply: Any) -> bool:
return getattr(guardrail_to_apply, "scan_only_tool_results", None) is True


def role_out_of_guardrail_scope(
role: str,
*,
skip_system_message: bool,
skip_tool_message: bool,
scan_only_tool_results: bool = False,
) -> bool:
if skip_system_message and role == "system":
return True
if skip_tool_message and role == "tool":
return True
return scan_only_tool_results and role != "tool"
Comment thread
veria-ai[bot] marked this conversation as resolved.
Outdated


def scoped_structured_message_indices(
messages: Sequence[AllMessageValues],
*,
scan_only_tool_results: bool,
skip_system: bool,
skip_tool: bool,
) -> tuple[int, ...]:
return tuple(
index
for index, message in enumerate(messages)
if not role_out_of_guardrail_scope(
_message_role(message),
skip_system_message=skip_system,
skip_tool_message=skip_tool,
scan_only_tool_results=scan_only_tool_results,
)
)


def merge_guardrailed_scoped_messages(
full_messages: Sequence[AllMessageValues],
scoped_indices: Sequence[int],
guardrailed_scoped: Sequence[AllMessageValues],
) -> list[AllMessageValues]:
return [m for m in messages if str((m or {}).get("role") or "").lower() != "tool"]
"""Substitute guardrail-returned messages back into the full conversation.

Guardrails only ever see the scoped subset of messages, so a replacement
list they hand back describes that subset, not the whole request. Writing
it over ``data["messages"]`` wholesale would silently drop every
out-of-scope message (system prompt, prior turns). Instead, swap each
returned message into the position its scoped original came from; extra
returned messages land after the last scoped position, and scoped
originals without a counterpart are treated as removed by the guardrail.
When nothing was filtered out this degenerates to the returned list
itself, preserving wholesale-replacement behavior for unscoped guardrails.
"""
replacements: Final = dict(zip(scoped_indices, guardrailed_scoped))
removed: Final = frozenset(scoped_indices[len(guardrailed_scoped) :])
appended: Final = tuple(guardrailed_scoped[len(scoped_indices) :])
last_scoped_index: Final = scoped_indices[-1] if scoped_indices else None

def _merged() -> Iterator[AllMessageValues]:
for index, message in enumerate(full_messages):
if index in removed:
continue
yield replacements.get(index, message)
if index == last_scoped_index:
yield from appended

return list(_merged())
41 changes: 27 additions & 14 deletions litellm/llms/openai/chat/guardrail_translation/handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,12 @@
StreamTransformSink,
)
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,
openai_messages_without_system,
openai_messages_without_tool,
merge_guardrailed_scoped_messages,
role_out_of_guardrail_scope,
scoped_structured_message_indices,
)
from litellm.main import stream_chunk_builder
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
Expand Down Expand Up @@ -82,6 +84,7 @@ async def process_input_messages(

skip_system: Final = effective_skip_system_message_for_guardrail(guardrail_to_apply)
skip_tool: Final = effective_skip_tool_message_for_guardrail(guardrail_to_apply)
scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply)

texts_to_check: Final[list[str]] = []
images_to_check: Final[list[str]] = []
Expand All @@ -101,6 +104,7 @@ async def process_input_messages(
tool_call_task_mappings=tool_call_task_mappings,
skip_system_message=skip_system,
skip_tool_message=skip_tool,
scan_only_tool_results=scan_only_tool_results,
)

# Step 2: Apply guardrail to all texts and tool calls in batch
Expand All @@ -110,16 +114,18 @@ async def process_input_messages(
inputs["images"] = images_to_check
if tool_calls_to_check:
inputs["tool_calls"] = tool_calls_to_check
structured_messages = self.get_structured_messages(data)
structured_messages: Final = self.get_structured_messages(data)
scoped_message_indices: Final = scoped_structured_message_indices(
structured_messages or [],
scan_only_tool_results=scan_only_tool_results,
skip_system=skip_system,
skip_tool=skip_tool,
)
if structured_messages:
if skip_system:
structured_messages = openai_messages_without_system(structured_messages)
if skip_tool:
structured_messages = openai_messages_without_tool(structured_messages)
inputs["structured_messages"] = structured_messages
inputs["structured_messages"] = [structured_messages[index] for index in scoped_message_indices]
Comment thread
veria-ai[bot] marked this conversation as resolved.
# Pass tools (function definitions) to the guardrail
tools: Final = data.get("tools")
if tools:
if tools and not scan_only_tool_results:
inputs["tools"] = tools
Comment thread
greptile-apps[bot] marked this conversation as resolved.
# Include model information if available
model: Final = data.get("model")
Expand All @@ -145,7 +151,11 @@ async def process_input_messages(
guardrailed_structured_messages is not None
and guardrailed_structured_messages is not original_structured_messages
):
data["messages"] = guardrailed_structured_messages
data["messages"] = merge_guardrailed_scoped_messages(
full_messages=structured_messages or [],
scoped_indices=scoped_message_indices,
guardrailed_scoped=guardrailed_structured_messages,
)
else:
# Step 3: Map guardrail responses back to original message structure
if guardrailed_texts and texts_to_check:
Expand Down Expand Up @@ -194,16 +204,19 @@ def _extract_inputs(
tool_call_task_mappings: list[tuple[int, int]],
skip_system_message: bool = False,
skip_tool_message: bool = False,
scan_only_tool_results: bool = False,
) -> None:
"""
Extract text content, images, and tool calls from a message.

Override this method to customize text/image/tool call extraction logic.
"""
role: Final = str(message.get("role") or "").lower()
if skip_system_message and role == "system":
return
if skip_tool_message and role == "tool":
if role_out_of_guardrail_scope(
str(message.get("role") or "").lower(),
skip_system_message=skip_system_message,
skip_tool_message=skip_tool_message,
scan_only_tool_results=scan_only_tool_results,
):
return

content: Final = message.get("content", None)
Expand Down
12 changes: 4 additions & 8 deletions litellm/proxy/guardrails/guardrail_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -487,16 +487,12 @@ def initialize_guardrail(
raise ValueError(f"Unsupported guardrail: {guardrail_type}")

if custom_guardrail_callback is not None:
setattr(
custom_guardrail_callback,
for scoping_param in (
"skip_system_message_in_guardrail",
getattr(litellm_params, "skip_system_message_in_guardrail", None),
)
setattr(
custom_guardrail_callback,
"skip_tool_message_in_guardrail",
getattr(litellm_params, "skip_tool_message_in_guardrail", None),
)
"scan_only_tool_results",
):
setattr(custom_guardrail_callback, scoping_param, getattr(litellm_params, scoping_param, None))
configured_run_in_parallel: Final = getattr(litellm_params, "run_in_parallel", None)
if configured_run_in_parallel is not None:
custom_guardrail_callback.run_in_parallel = bool(configured_run_in_parallel)
Expand Down
10 changes: 10 additions & 0 deletions litellm/types/guardrails.py
Original file line number Diff line number Diff line change
Expand Up @@ -757,6 +757,16 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up
),
)

scan_only_tool_results: Optional[bool] = Field(
default=None,
description=(
"When True, unified guardrails only evaluate tool results, the untrusted data an "
"agent feeds back into the model, and skip system, user, and assistant content. "
"Intended for agent harnesses whose own prompt scaffolding is trusted but often "
"trips prompt-attack detectors."
),
)

# Lakera specific params
category_thresholds: Optional[LakeraCategoryThresholds] = Field(
default=None,
Expand Down
Loading
Loading