scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(guardrails): scan replayed Claude server tool calls in Lasso, Lakera and Qualifire request checks

Browse source
This commit is contained in:
mateo-berri 2026-11-05 23:33:17 -08:00
parent b70c14803c
commit 418409e9d7
7 changed files with 238 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py
View file

@ -22,6 +22,7 @@ from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
@ -163,11 +164,11 @@ class lakeraAI_Moderation(CustomGuardrail):
for message in data["messages"]:
role = message.get("role")
if role in stringified_roles:
if "tool_calls" in message:
tool_call_messages = [
*tool_call_messages,
*message["tool_calls"],
]
tool_call_messages = [
*tool_call_messages,
*anthropic_server_tool_call_params(message),
*message.get("tool_calls", ()),
]
if role == Role.SYSTEM.value: # we need this for later
system_message = message
continue
@ -185,10 +186,11 @@ class lakeraAI_Moderation(CustomGuardrail):
if system_message is not None:
if not litellm.add_function_to_prompt:
content = system_message.get("content")
function_input: Final = []
for tool_call in tool_call_messages:
if "function" in tool_call:
function_input.append(tool_call["function"]["arguments"])
function_input: Final[tuple[str, ...]] = tuple(
tool_call["function"]["arguments"]
for tool_call in tool_call_messages
if "function" in tool_call
)
if len(function_input) > 0:
content += " Function Input: " + " ".join(function_input)

12
litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py
View file

@ -45,6 +45,7 @@ from litellm.llms.anthropic.common_utils import (
anthropic_server_tool_calls,
with_server_tool_call_arguments,
write_anthropic_server_tool_calls,
write_request_server_tool_calls,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
@ -571,6 +572,15 @@ class LassoGuardrail(CustomGuardrail):
elif role == "assistant" and not content and orig_msg.get("tool_calls"):
msg["tool_calls"] = self._update_tool_calls_from_masked(orig_msg["tool_calls"], masked_tool_use)
if role == "assistant":
write_request_server_tool_calls(
msg,
tuple(
self._masked_server_tool_call(server_tool_call, masked_tool_use)
for server_tool_call in anthropic_server_tool_call_params(orig_msg)
),
)
result.append(msg)
return result
@ -711,7 +721,7 @@ class LassoGuardrail(CustomGuardrail):
expanded.append({"role": role, "content": content})
if role == "assistant":
for call in iter_items(msg.get("tool_calls") or []):
for call in (*anthropic_server_tool_call_params(msg), *iter_items(msg.get("tool_calls") or [])):
call_id, name, input_data = self._extract_tool_call_fields(call)
if not call_id or not name:
verbose_proxy_logger.warning(

5
litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py
View file

@ -18,6 +18,7 @@ from litellm.integrations.custom_guardrail import (
log_guardrail_information,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
@ -180,8 +181,8 @@ class QualifireGuardrail(CustomGuardrail):
api_message["tool_call_id"] = tool_call_id
# Handle tool calls if present
tool_calls = msg.get("tool_calls")
if tool_calls and isinstance(tool_calls, list):
tool_calls = [*anthropic_server_tool_call_params(msg), *(msg.get("tool_calls") or [])]
if tool_calls:
api_tool_calls = []
for tc in tool_calls:
if isinstance(tc, dict):

13
tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py
View file

@ -6,10 +6,11 @@ with guardrail transformations, including tool calls.
"""
import json
from collections.abc import Mapping
from collections.abc import Iterable, Mapping
from typing import Any, Final, Literal, Optional
import pytest
from openai.types.responses import ResponseFunctionWebSearch
from litellm.integrations.custom_guardrail import CustomGuardrail
@ -149,14 +150,14 @@ def _server_tool_call(call_id: str, name: str, arguments: Mapping[str, object])
return {"id": call_id, "type": "function", "function": {"name": name, "arguments": json.dumps(arguments)}}
def _tool_call_ids_names_and_arguments(tool_calls) -> list[tuple[str, str, str]]:
def _tool_call_ids_names_and_arguments(tool_calls: Iterable[Mapping[str, Any]]) -> list[tuple[str, str, str]]:
return [
(tool_call["id"], tool_call["function"]["name"], tool_call["function"]["arguments"])
for tool_call in tool_calls
]
def _hosted_web_call_inputs(web_search_calls) -> list[tuple[str, str, Optional[str]]]:
def _hosted_web_call_inputs(web_search_calls: Iterable[ResponseFunctionWebSearch]) -> list[tuple[str, str, Optional[str]]]:
return [
(call.id, call.status, call.action.query if call.action.type == "search" else call.action.url)
for call in web_search_calls
@ -177,9 +178,9 @@ class QueryRedactingGuardrail(CustomGuardrail):
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
def redacted(tool_call):
arguments = json.loads(tool_call["function"]["arguments"])
redacted_arguments = {
def redacted(tool_call: Mapping[str, Any]) -> Mapping[str, Any]:
arguments: Final[dict[str, object]] = json.loads(tool_call["function"]["arguments"])
redacted_arguments: Final = {
key: "[REDACTED]" if key in ("query", "url") else value for key, value in arguments.items()
}
if redacted_arguments == arguments:

57
tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py Normal file
View file

@ -0,0 +1,57 @@
import json
from typing import Final
import httpx
import pytest
import respx
import litellm
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import lakeraAI_Moderation
@pytest.mark.asyncio
async def test_prompt_check_appends_replayed_server_tool_call_arguments_to_system_message(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(litellm, "add_function_to_prompt", False)
route: Final = respx_mock.post("https://api.lakera.ai/v1/prompt_injection").mock(
return_value=httpx.Response(200, json={"results": [{"flagged": False}]})
)
moderation: Final = lakeraAI_Moderation(api_key="test-key", moderation_check="pre_call")
await moderation._check(
data={
"messages": [
{"role": "system", "content": "You are a shopping assistant."},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_01",
"type": "function",
"function": {"name": "lookup_order", "arguments": '{"order_id": "A-17"}'},
}
],
"provider_specific_fields": {
"server_tool_calls": [
{
"id": "srvtoolu_01",
"type": "function",
"function": {"name": "web_fetch", "arguments": '{"url": "https://example.com"}'},
}
]
},
},
]
},
user_api_key_dict=UserAPIKeyAuth(),
call_type="completion",
)
assert json.loads(route.calls.last.request.content)["input"][0] == {
"role": "system",
"content": 'You are a shopping assistant. Function Input: {"url": "https://example.com"} {"order_id": "A-17"}',
}

110
tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py
View file

@ -1691,3 +1691,113 @@ async def test_post_call_masking_rewrites_server_tool_call_arguments(
}
]
}
def _replayed_web_fetch_request(url: str) -> dict[str, object]:
return {
"messages": [
{"role": "user", "content": "open the order page"},
{
"role": "assistant",
"content": None,
"provider_specific_fields": {
"server_tool_calls": [
{
"id": "srvtoolu_01",
"type": "function",
"function": {"name": "web_fetch", "arguments": json.dumps({"url": url})},
}
]
},
},
{"role": "user", "content": "summarize it"},
]
}
@pytest.mark.asyncio
async def test_pre_call_scan_sends_replayed_server_tool_calls_as_tool_use_entries(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
lasso_guardrail: Final = LassoGuardrail(
lasso_api_key="test-key",
user_id="test-user",
conversation_id="test-conversation",
guardrail_name="test-guard",
event_hook="pre_call",
default_on=True,
)
route: Final = respx_mock.post("https://server.lasso.security/gateway/v3/classify").mock(
return_value=httpx.Response(200, json={"violations_detected": False, "deputies": {}, "findings": {}})
)
await lasso_guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(),
cache=DualCache(),
data=_replayed_web_fetch_request("https://example.com/?email=jane@example.com"),
call_type="completion",
)
assert json.loads(route.calls.last.request.content)["messages"] == [
{"role": "user", "content": "open the order page"},
{
"role": "model",
"content": {
"type": "tool_use",
"id": "srvtoolu_01",
"name": "web_fetch",
"input": {"url": "https://example.com/?email=jane@example.com"},
},
},
{"role": "user", "content": "summarize it"},
]
@pytest.mark.asyncio
async def test_pre_call_masking_rewrites_replayed_server_tool_call_arguments(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
lasso_guardrail: Final = LassoGuardrail(
lasso_api_key="test-key",
user_id="test-user",
conversation_id="test-conversation",
mask=True,
guardrail_name="test-guard",
event_hook="pre_call",
default_on=True,
)
respx_mock.post("https://server.lasso.security/gateway/v3/classifix").mock(
return_value=httpx.Response(
200,
json={
"violations_detected": True,
"deputies": {"pattern-detection": True},
"findings": {"pattern-detection": [{"action": "AUTO_MASKING", "severity": "HIGH"}]},
"messages": [
{"role": "user", "content": "open the order page"},
{
"role": "model",
"content": {
"type": "tool_use",
"id": "srvtoolu_01",
"name": "web_fetch",
"input": {"url": "https://example.com/?email=<EMAIL_ADDRESS>"},
},
},
{"role": "user", "content": "summarize it"},
],
},
)
)
result: Final = await lasso_guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(),
cache=DualCache(),
data=_replayed_web_fetch_request("https://example.com/?email=jane@example.com"),
call_type="completion",
)
assert isinstance(result, dict)
assert result["messages"] == _replayed_web_fetch_request("https://example.com/?email=<EMAIL_ADDRESS>")["messages"]

39
tests/unit/proxy/guardrails/guardrail_hooks/test_qualifire.py
View file

@ -2,6 +2,7 @@
Unit tests for Qualifire guardrail integration.
"""
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -735,3 +736,41 @@ class TestQualifireGuardrailRegistry:
assert "qualifire" in guardrail_class_registry
assert guardrail_class_registry["qualifire"] == QualifireGuardrail
def test_convert_messages_sends_replayed_server_tool_calls_with_client_tool_calls() -> None:
from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import (
QualifireGuardrail,
)
guardrail: Final = QualifireGuardrail(api_key="test_key", guardrail_name="test_guardrail")
result: Final = guardrail._convert_messages_to_api_format(
[
{
"role": "assistant",
"content": "",
"tool_calls": [
{
"id": "call_123",
"type": "function",
"function": {"name": "get_weather", "arguments": '{"location": "NYC"}'},
}
],
"provider_specific_fields": {
"server_tool_calls": [
{
"id": "srvtoolu_01",
"type": "function",
"function": {"name": "web_fetch", "arguments": '{"url": "https://example.com"}'},
}
]
},
},
]
)
assert result[0]["tool_calls"] == [
{"id": "srvtoolu_01", "name": "web_fetch", "arguments": {"url": "https://example.com"}},
{"id": "call_123", "name": "get_weather", "arguments": {"location": "NYC"}},
]