scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

Merge branch 'litellm_lakera_v2_gemini_contents' into litellm_passthrough_applied_guardrails_header

Browse source
This commit is contained in:
yucheng 2026-11-02 18:05:07 -08:00
commit 4afb8292d5
8 changed files with 203 additions and 101 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -15,6 +15,7 @@ from typing import Final, Literal
import httpx
from fastapi import HTTPException
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
@ -22,6 +23,7 @@ from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
@ -45,6 +47,20 @@ INPUT_POSITIONING_MAP: Final = {
Role.ASSISTANT.value: 2,
}
_CONTENT_PARTS: Final = TypeAdapter(list[object])
def _input_item_text(item: object) -> str:
"""The text of one entry of a list ``input``: a string, or a Responses message's string or parts ``content``."""
if isinstance(item, str):
return item
content: Final = (as_str_mapping(item) or {}).get("content")
if isinstance(content, str):
return content
parts: Final = _CONTENT_PARTS.validate_python(content) if isinstance(content, list) else ()
part_texts: Final = ((as_str_mapping(part) or {}).get("text") for part in parts)
return "\n".join(text for text in part_texts if isinstance(text, str))
class lakeraAI_Moderation(CustomGuardrail):
@classmethod
@ -219,7 +235,7 @@ class lakeraAI_Moderation(CustomGuardrail):
}
)
elif "input" in data and isinstance(data["input"], list):
text = "\n".join(item for item in data["input"] if isinstance(item, str))
text = "\n".join(filter(None, map(_input_item_text, data["input"])))
if not text:
verbose_proxy_logger.debug("Skipping lakera prompt injection, input has no text")
return
@ -230,6 +246,10 @@ class lakeraAI_Moderation(CustomGuardrail):
}
)
if not _json_data and call_type == "pass_through_endpoint":
verbose_proxy_logger.debug("Skipping lakera prompt injection, pass-through body has no messages or input")
return
verbose_proxy_logger.debug("Lakera AI Request Args %s", _json_data)
# https://platform.lakera.ai/account/api-keys

22
litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py
View file

@ -298,18 +298,16 @@ def _opted_out_of_global_guardrail(guardrail: "LakeraAIGuardrail", user_api_key_
"""True if the key's or team's admin metadata opts out of this ``default_on`` guardrail. A pass-through body has no
admin metadata for ``should_run_guardrail`` to read, so it is read from the key here."""
admin_fields: Final = user_api_key_dict.model_dump(include={"metadata", "team_metadata"})
return (
guardrail.default_on is True
and guardrail.should_run_guardrail( # pyright: ignore[reportUnknownMemberType] # takes an untyped data dict
data={
"metadata": {
"user_api_key_metadata": admin_fields.get("metadata"),
"user_api_key_team_metadata": admin_fields.get("team_metadata"),
}
},
event_type=GuardrailEventHooks.pre_call,
)
is not True
admin_data: Final = {
"metadata": {
"user_api_key_metadata": admin_fields.get("metadata"),
"user_api_key_team_metadata": admin_fields.get("team_metadata"),
}
}
return guardrail.default_on is True and (
guardrail.get_disable_global_guardrail(admin_data) # pyright: ignore[reportUnknownMemberType] # takes an untyped data dict
is True
or guardrail.guardrail_name in guardrail.get_opted_out_global_guardrails_from_metadata(admin_data) # pyright: ignore[reportUnknownMemberType] # takes an untyped data dict
)

52
tests/integration/authorization/test_passthrough_auth_guardrails.py
View file

@ -22,6 +22,7 @@ GUARDRAIL: Final = "passthrough-auth-guardrail"
GLOBAL_POST_CALL: Final = "passthrough-global-post-call"
KEY_POST_CALL: Final = "passthrough-key-post-call"
GENERIC_POST_CALL: Final = "passthrough-generic-post-call"
BLOCKING_POST_CALL: Final = "passthrough-post-call-blocks-every-reply"
MODEL_POST_CALL: Final = "passthrough-model-post-call"
GUARDED_MODEL: Final = "passthrough-guarded-model"
POST_CALL_GUARDED_MODEL: Final = "passthrough-post-call-guarded-model"
@ -77,6 +78,10 @@ def _verdict(request: Request) -> Reply:
return Reply(body=json.dumps({"action": action, "blocked_reason": "synthetic block"}).encode())
def _block_every_reply(_request: Request) -> Reply:
return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": "synthetic block"}).encode())
_OPENAI_BODY: Final = json.dumps(
{
**_JSON.validate_json(chat_reply("chatcmpl-passthrough-auth", "gpt-5.6", "Hi", stream=False).body),
@ -148,6 +153,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I
post_call_vendor: Final = stack.enter_context(wire_server(_lakera_clean))
key_post_call_vendor: Final = stack.enter_context(wire_server(_lakera_flagged))
generic_post_call_vendor: Final = stack.enter_context(wire_server(_verdict))
blocking_post_call_vendor: Final = stack.enter_context(wire_server(_block_every_reply))
model_post_call_vendor: Final = stack.enter_context(wire_server(_verdict))
upstreams: Final = MappingProxyType(
{
@ -238,6 +244,16 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I
"api_key": "synthetic-guardrail-key",
},
},
{
"guardrail_name": BLOCKING_POST_CALL,
"litellm_params": {
"guardrail": "generic_guardrail_api",
"mode": "post_call",
"default_on": False,
"api_base": blocking_post_call_vendor.url,
"api_key": "synthetic-guardrail-key",
},
},
{
"guardrail_name": GENERIC_POST_CALL,
"litellm_params": {
@ -592,7 +608,7 @@ def _tagged_body(route: str, tag: str) -> dict[str, JsonValue]:
return {**_route_body(route, "What is the capital of France?"), "metadata": {"tags": [tag]}}
@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE])
@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE])
def test_post_call_policy_pipeline_blocks_a_flagged_pass_through_response(rig: AuthGuardrailRig, route: str) -> None:
gateway: Final = rig.owned.gateway
with gateway.scenario() as scenario:
@ -646,15 +662,41 @@ def test_pre_call_policy_pipeline_relays_the_provider_bytes(rig: AuthGuardrailRi
assert len(rig.vendor.drain()) == 1
@pytest.mark.parametrize("route", [ANTHROPIC_ROUTE, GEMINI_ROUTE])
def test_response_a_post_call_guardrail_blocks_still_charges_the_provider_call(
rig: AuthGuardrailRig, route: str
) -> None:
def test_lakera_post_call_has_nothing_to_read_on_a_gemini_body(rig: AuthGuardrailRig) -> None:
gateway: Final = rig.owned.gateway
with gateway.scenario() as scenario:
key: Final = scenario.key(metadata={"guardrails": [KEY_POST_CALL]})
_drain(rig)
response: Final = _send(gateway, GEMINI_ROUTE, _route_body(GEMINI_ROUTE, "What is the capital of France?"), key)
assert response.status_code == 200, response.text
assert rig.key_post_call_vendor.drain() == ()
assert len(rig.upstreams[GEMINI_ROUTE].drain()) == 1
def test_post_call_policy_pipeline_lakera_step_has_nothing_to_read_on_a_gemini_body(rig: AuthGuardrailRig) -> None:
gateway: Final = rig.owned.gateway
with gateway.scenario() as scenario:
key: Final = scenario.key()
_drain(rig)
response: Final = _send(gateway, GEMINI_ROUTE, _tagged_body(GEMINI_ROUTE, PIPELINE_BLOCKS), key)
assert response.status_code == 200, response.text
assert rig.key_post_call_vendor.drain() == ()
assert len(rig.upstreams[GEMINI_ROUTE].drain()) == 1
@pytest.mark.parametrize(("route", "guardrail"), [(ANTHROPIC_ROUTE, KEY_POST_CALL), (GEMINI_ROUTE, BLOCKING_POST_CALL)])
def test_response_a_post_call_guardrail_blocks_still_charges_the_provider_call(
rig: AuthGuardrailRig, route: str, guardrail: str
) -> None:
gateway: Final = rig.owned.gateway
with gateway.scenario() as scenario:
key: Final = scenario.key(metadata={"guardrails": [guardrail]})
_drain(rig)
response: Final = _send(gateway, route, _route_body(route, "What is the capital of France?"), key)
assert response.status_code == 400, response.text

11
tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py
View file

@ -2692,3 +2692,14 @@ async def test_during_mcp_call_blocked_raises() -> None:
)
assert exc_info.value.message == "credential detected in tool input"
@pytest.mark.parametrize("value", (None, 0, False, "", [], "text", 7))
def test_extract_user_api_key_metadata_reads_past_a_litellm_metadata_that_is_not_an_object(value: object) -> None:
guardrail: Final = GenericGuardrailAPI(api_base="http://vendor.invalid", guardrail_name="generic")
metadata: Final = guardrail._extract_user_api_key_metadata(
{"metadata": {"user_api_key_hash": "hash"}, "litellm_metadata": value}
)
assert metadata["user_api_key_hash"] == "hash"

83
tests/unit/proxy/guardrails/guardrail_hooks/test_guardrails_on_unexpected_request_fields.py
View file

@ -1,83 +0,0 @@
"""Guardrails on built-in pass-through routes read client JSON that request setup hasn't normalized."""
from typing import Final
from unittest.mock import AsyncMock, patch
import httpx
import pytest
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import (
GenericGuardrailAPI,
)
from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import lakeraAI_Moderation
from litellm.proxy.guardrails.guardrail_hooks.noma.noma_v2 import NomaV2Guardrail
NOT_AN_OBJECT: Final = (None, 0, False, "", [], "text", 7)
@pytest.mark.parametrize("value", NOT_AN_OBJECT)
def test_generic_guardrail_reads_metadata_past_a_litellm_metadata_that_is_not_an_object(value: object) -> None:
guardrail: Final = GenericGuardrailAPI(api_base="http://vendor.invalid", guardrail_name="generic")
metadata: Final = guardrail._extract_user_api_key_metadata(
{"metadata": {"user_api_key_hash": "hash"}, "litellm_metadata": value}
)
assert metadata["user_api_key_hash"] == "hash"
@pytest.mark.asyncio
@pytest.mark.parametrize("value", NOT_AN_OBJECT)
async def test_noma_v2_scans_a_request_whose_litellm_metadata_is_not_an_object(value: object) -> None:
guardrail: Final = NomaV2Guardrail(api_key="key", api_base="http://vendor.invalid", guardrail_name="noma")
verdict: Final = httpx.Response(
200, json={"action": "NONE"}, request=httpx.Request("POST", "http://vendor.invalid")
)
with patch.object(guardrail.async_handler, "post", new=AsyncMock(return_value=verdict)):
result: Final = await guardrail.apply_guardrail(
inputs={"texts": ["hello"]},
request_data={"metadata": {}, "litellm_metadata": value},
input_type="request",
)
assert result["texts"] == ["hello"]
@pytest.mark.asyncio
@pytest.mark.parametrize("token_arrays", [[1, 2, 3], [[1, 2], [3, 4]]])
async def test_legacy_lakera_has_no_text_to_scan_in_token_array_embeddings(token_arrays: list[object]) -> None:
guardrail: Final = lakeraAI_Moderation(api_key="key", api_base="http://vendor.invalid")
vendor_error: Final = httpx.HTTPStatusError(
"vendor called",
request=httpx.Request("POST", "http://vendor.invalid"),
response=httpx.Response(500, text="vendor called"),
)
with patch.object(guardrail.async_handler, "post", new=AsyncMock(side_effect=vendor_error)):
await guardrail._check(
data={"input": token_arrays},
user_api_key_dict=UserAPIKeyAuth(api_key="key"),
call_type="pass_through_endpoint",
)
@pytest.mark.asyncio
async def test_legacy_lakera_still_scans_string_inputs() -> None:
guardrail: Final = lakeraAI_Moderation(api_key="key", api_base="http://vendor.invalid")
vendor_error: Final = httpx.HTTPStatusError(
"vendor called",
request=httpx.Request("POST", "http://vendor.invalid"),
response=httpx.Response(500, text="vendor called"),
)
with (
patch.object(guardrail.async_handler, "post", new=AsyncMock(side_effect=vendor_error)),
pytest.raises(Exception, match="vendor called"),
):
await guardrail._check(
data={"input": ["hello"]},
user_api_key_dict=UserAPIKeyAuth(api_key="key"),
call_type="pass_through_endpoint",
)

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

@ -0,0 +1,57 @@
"""The legacy Lakera guardrail reads client JSON that request setup hasn't normalized on pass-through routes."""
import json
from typing import Final
from unittest.mock import AsyncMock, patch
import httpx
import pytest
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import lakeraAI_Moderation
VENDOR: Final = "http://vendor.invalid"
CLEAN: Final = httpx.Response(200, json={"results": [{"flagged": False}]}, request=httpx.Request("POST", VENDOR))
async def _sent_to_lakera(data: dict[str, object], call_type: str = "pass_through_endpoint") -> tuple[str, ...]:
guardrail: Final = lakeraAI_Moderation(api_key="key", api_base=VENDOR)
with patch.object(guardrail.async_handler, "post", new=AsyncMock(return_value=CLEAN)) as post:
await guardrail._check(data=data, user_api_key_dict=UserAPIKeyAuth(api_key="key"), call_type=call_type)
return tuple(call.kwargs["data"] for call in post.await_args_list)
@pytest.mark.asyncio
@pytest.mark.parametrize("token_arrays", [[1, 2, 3], [[1, 2], [3, 4]]])
async def test_token_arrays_have_no_text_to_send(token_arrays: list[object]) -> None:
assert await _sent_to_lakera({"input": token_arrays}) == ()
@pytest.mark.asyncio
async def test_string_entries_of_a_list_input_are_sent() -> None:
sent: Final = await _sent_to_lakera({"input": ["hello", "there"]})
assert [json.loads(body)["input"] for body in sent] == ["hello\nthere"]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"content",
["card 4111", [{"type": "input_text", "text": "card 4111"}, {"type": "input_image", "image_url": "http://x"}]],
)
async def test_a_responses_message_in_a_list_input_is_sent(content: object) -> None:
sent: Final = await _sent_to_lakera({"input": [{"role": "user", "content": content}]})
assert [json.loads(body)["input"] for body in sent] == ["card 4111"]
@pytest.mark.asyncio
async def test_a_pass_through_body_with_no_messages_or_input_sends_nothing() -> None:
contents: Final = {"contents": [{"role": "user", "parts": [{"text": "hello"}]}]}
assert await _sent_to_lakera(contents) == ()
@pytest.mark.asyncio
async def test_a_native_body_with_no_messages_or_input_still_reaches_lakera() -> None:
assert await _sent_to_lakera({"prompt": "hello"}, call_type="text_completion") == ("",)

37
tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py
View file

@ -1940,6 +1940,43 @@ class TestGeminiContents:
call_type="pass_through_endpoint",
)
@pytest.mark.asyncio
@pytest.mark.parametrize("call_type", ["pass_through_endpoint", "agenerate_content_stream"])
async def test_streaming_only_global_guardrail_reads_a_gemini_stream(self, call_type: str):
guardrail: Final = LakeraAIGuardrail(
api_key="test_key", guardrail_name="lakera-global", default_on=True, stream_scope="streaming"
)
with (
patch.object(guardrail, "call_v2_guard", new=AsyncMock(return_value=(GEMINI_FLAGGED, {}))),
pytest.raises(HTTPException),
):
await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test_key"),
cache=DualCache(),
data={**_gemini_request(), "stream": True},
call_type=call_type,
)
@pytest.mark.asyncio
async def test_tag_scoped_global_guardrail_reads_a_gemini_body_with_its_tag(self):
guardrail: Final = LakeraAIGuardrail(
api_key="test_key",
guardrail_name="lakera-global",
default_on=True,
event_hook=Mode(tags={"prod": "pre_call"}, default="during_call"),
)
with (
patch.object(guardrail, "call_v2_guard", new=AsyncMock(return_value=(GEMINI_FLAGGED, {}))),
patch("litellm.proxy.proxy_server.premium_user", True),
pytest.raises(HTTPException, match="Violated"),
):
await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test_key"),
cache=DualCache(),
data={**_gemini_request(), "metadata": {"tags": ["prod"]}},
call_type="pass_through_endpoint",
)
@pytest.mark.parametrize("chat_field", ["messages", "input"])
def test_a_pass_through_chat_body_with_stray_contents_is_read_as_chat(self, chat_field: str):
data: Final = {chat_field: [USER_MSG], "contents": [{"parts": [{"text": "stray"}]}]}

20
tests/unit/proxy/guardrails/guardrail_hooks/test_noma_v2.py
View file

@ -1,7 +1,9 @@
import asyncio
import os
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from litellm.proxy.guardrails.guardrail_hooks.noma import NomaV2Guardrail
@ -747,3 +749,21 @@ class TestNomaV2ApplicationIdResolution:
payload = call_mock.call_args.kwargs["payload"]
assert "application_id" not in payload
@pytest.mark.asyncio
@pytest.mark.parametrize("value", (None, 0, False, "", [], "text", 7))
async def test_noma_v2_scans_a_request_whose_litellm_metadata_is_not_an_object(value: object) -> None:
guardrail: Final = NomaV2Guardrail(api_key="key", api_base="http://vendor.invalid", guardrail_name="noma")
verdict: Final = httpx.Response(
200, json={"action": "NONE"}, request=httpx.Request("POST", "http://vendor.invalid")
)
with patch.object(guardrail.async_handler, "post", new=AsyncMock(return_value=verdict)):
result: Final = await guardrail.apply_guardrail(
inputs={"texts": ["hello"]},
request_data={"metadata": {}, "litellm_metadata": value},
input_type="request",
)
assert result["texts"] == ["hello"]