ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(guardrails): optional skip tool message in unified guardrail inputs

Browse Source 
Mirrors the system-message skip in PR #25481 for tool-role messages.
Adds a global litellm.skip_tool_message_in_guardrail flag and a
per-guardrail litellm_params.skip_tool_message_in_guardrail override,
applied in the OpenAI and Anthropic chat translation handlers.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
shivam 2026-05-07 18:39:26 -07:00
parent ee8d8c4137
commit f58f8927f2
No known key found for this signature in database
7 changed files with 192 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
litellm/__init__.py
View File

@ -206,6 +206,7 @@ add_user_information_to_llm_headers: Optional[bool] = (
)
store_audit_logs = False # Enterprise feature, allow users to see audit logs
skip_system_message_in_guardrail: bool = False
skip_tool_message_in_guardrail: bool = False
### end of callbacks #############
email: Optional[str] = (

12
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View File

@ -23,7 +23,9 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_skip_system_message_for_guardrail,
effective_skip_tool_message_for_guardrail,
openai_messages_without_system,
openai_messages_without_tool,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
@ -108,6 +110,7 @@ class AnthropicMessagesHandler(BaseTranslation):
return data
skip_system = effective_skip_system_message_for_guardrail(guardrail_to_apply)
skip_tool = effective_skip_tool_message_for_guardrail(guardrail_to_apply)
chat_completion_compatible_request = self._translate_to_openai(data)
@ -117,6 +120,8 @@ class AnthropicMessagesHandler(BaseTranslation):
)
if skip_system:
structured_messages = openai_messages_without_system(structured_messages)
if skip_tool:
structured_messages = openai_messages_without_tool(structured_messages)
texts_to_check: List[str] = []
images_to_check: List[str] = []
@ -134,6 +139,7 @@ class AnthropicMessagesHandler(BaseTranslation):
images_to_check=images_to_check,
task_mappings=task_mappings,
skip_system_message=skip_system,
skip_tool_message=skip_tool,
)
# Step 2: Apply guardrail to all texts in batch
@ -198,13 +204,17 @@ class AnthropicMessagesHandler(BaseTranslation):
images_to_check: List[str],
task_mappings: List[Tuple[int, Optional[int]]],
skip_system_message: bool = False,
skip_tool_message: bool = False,
) -> None:
"""
Extract text content and images from a message.
Override this method to customize text/image extraction logic.
"""
if skip_system_message and str(message.get("role") or "").lower() == "system":
role = str(message.get("role") or "").lower()
if skip_system_message and role == "system":
return
if skip_tool_message and role == "tool":
return
content = message.get("content", None)

15
litellm/llms/base_llm/guardrail_translation/utils.py
View File

@ -14,7 +14,22 @@ def effective_skip_system_message_for_guardrail(guardrail_to_apply: Any) -> bool
return bool(getattr(litellm, "skip_system_message_in_guardrail", False))
def effective_skip_tool_message_for_guardrail(guardrail_to_apply: Any) -> bool:
per = getattr(guardrail_to_apply, "skip_tool_message_in_guardrail", None)
if per is not None:
return bool(per)
import litellm
return bool(getattr(litellm, "skip_tool_message_in_guardrail", False))
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"]
def openai_messages_without_tool(
messages: List[AllMessageValues],
) -> List[AllMessageValues]:
return [m for m in messages if str((m or {}).get("role") or "").lower() != "tool"]

24
litellm/llms/openai/chat/guardrail_translation/handler.py
View File

@ -21,7 +21,9 @@ from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_skip_system_message_for_guardrail,
effective_skip_tool_message_for_guardrail,
openai_messages_without_system,
openai_messages_without_tool,
)
from litellm.main import stream_chunk_builder
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
@ -73,6 +75,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
return data
skip_system = effective_skip_system_message_for_guardrail(guardrail_to_apply)
skip_tool = effective_skip_tool_message_for_guardrail(guardrail_to_apply)
texts_to_check: List[str] = []
images_to_check: List[str] = []
@ -91,6 +94,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
text_task_mappings=text_task_mappings,
tool_call_task_mappings=tool_call_task_mappings,
skip_system_message=skip_system,
skip_tool_message=skip_tool,
)
# Step 2: Apply guardrail to all texts and tool calls in batch
@ -102,11 +106,15 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
inputs["tool_calls"] = tool_calls_to_check # type: ignore
structured_messages = self.get_structured_messages(data)
if structured_messages:
inputs["structured_messages"] = (
openai_messages_without_system(structured_messages)
if skip_system
else 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
# Pass tools (function definitions) to the guardrail
tools = data.get("tools")
if tools:
@ -176,13 +184,17 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
text_task_mappings: List[Tuple[int, Optional[int]]],
tool_call_task_mappings: List[Tuple[int, int]],
skip_system_message: bool = False,
skip_tool_message: bool = False,
) -> None:
"""
Extract text content, images, and tool calls from a message.
Override this method to customize text/image/tool call extraction logic.
"""
if skip_system_message and str(message.get("role") or "").lower() == "system":
role = str(message.get("role") or "").lower()
if skip_system_message and role == "system":
return
if skip_tool_message and role == "tool":
return
content = message.get("content", None)

5
litellm/proxy/guardrails/guardrail_registry.py
View File

@ -482,6 +482,11 @@ class InMemoryGuardrailHandler:
"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),
)
parsed_guardrail = Guardrail(
guardrail_id=guardrail.get("guardrail_id"),

10
litellm/types/guardrails.py
View File

@ -633,6 +633,16 @@ class BaseLitellmParams(
),
)
skip_tool_message_in_guardrail: Optional[bool] = Field(
default=None,
description=(
"When True, unified guardrails skip tool-role messages when building "
"evaluation inputs (texts and structured_messages). When False, tool "
"messages are included even if litellm_settings sets a global skip. When "
"None, use the global litellm.skip_tool_message_in_guardrail setting."
),
)
# Lakera specific params
category_thresholds: Optional[LakeraCategoryThresholds] = Field(
default=None,

132
tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py
View File

@ -8,7 +8,9 @@ from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_skip_system_message_for_guardrail,
effective_skip_tool_message_for_guardrail,
openai_messages_without_system,
openai_messages_without_tool,
)
from litellm.llms.openai.chat.guardrail_translation.handler import (
OpenAIChatCompletionsHandler,
@ -180,6 +182,136 @@ class TestUnifiedLLMGuardrails:
}
assert "system" in roles
class TestSkipToolMessageForChatCompletions:
def test_openai_messages_without_tool(self):
msgs = [
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {"name": "f", "arguments": "{}"},
}
],
},
{"role": "tool", "content": "tool result", "tool_call_id": "call_1"},
]
out = openai_messages_without_tool(msgs)
assert len(out) == 2
assert all(m["role"] != "tool" for m in out)
assert msgs[2]["content"] == "tool result"
def test_effective_skip_tool_respects_per_guardrail_over_global(
self, monkeypatch
):
monkeypatch.setattr(
litellm, "skip_tool_message_in_guardrail", True, raising=False
)
class G:
skip_tool_message_in_guardrail = False
assert effective_skip_tool_message_for_guardrail(G()) is False
class G2:
skip_tool_message_in_guardrail = None
assert effective_skip_tool_message_for_guardrail(G2()) is True
@pytest.mark.asyncio
async def test_openai_handler_skips_tool_in_guardrail_inputs(self, monkeypatch):
monkeypatch.setattr(
litellm, "skip_tool_message_in_guardrail", True, raising=False
)
captured = {}
class MockGuardrail:
skip_tool_message_in_guardrail = None
async def apply_guardrail(
self, inputs, request_data, input_type, logging_obj=None
):
captured["inputs"] = inputs
return inputs
data = {
"messages": [
{"role": "user", "content": "hello"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {"name": "f", "arguments": "{}"},
}
],
},
{
"role": "tool",
"content": "secret tool result",
"tool_call_id": "call_1",
},
],
"model": "gpt-4o",
}
handler = OpenAIChatCompletionsHandler()
await handler.process_input_messages(
data=data,
guardrail_to_apply=MockGuardrail(),
litellm_logging_obj=None,
)
assert "secret tool result" not in captured["inputs"]["texts"]
sm = captured["inputs"].get("structured_messages") or []
assert all(m.get("role") != "tool" for m in sm)
assert data["messages"][2]["content"] == "secret tool result"
@pytest.mark.asyncio
async def test_openai_handler_per_guardrail_skip_tool_false_overrides_global(
self, monkeypatch
):
monkeypatch.setattr(
litellm, "skip_tool_message_in_guardrail", True, raising=False
)
captured = {}
class MockGuardrail:
skip_tool_message_in_guardrail = False
async def apply_guardrail(
self, inputs, request_data, input_type, logging_obj=None
):
captured["inputs"] = inputs
return inputs
data = {
"messages": [
{"role": "user", "content": "u"},
{"role": "tool", "content": "tr", "tool_call_id": "call_1"},
],
}
await OpenAIChatCompletionsHandler().process_input_messages(
data=data,
guardrail_to_apply=MockGuardrail(),
litellm_logging_obj=None,
)
assert "tr" in captured["inputs"]["texts"]
roles = {
m.get("role")
for m in (captured["inputs"].get("structured_messages") or [])
}
assert "tool" in roles
class TestAsyncPreCallHook:
@pytest.mark.asyncio
async def test_uses_mcp_event_type(self):