From 1a5b56c9a279488e46b0affb8b33de8617c7ca1f Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 17:58:15 -0800 Subject: [PATCH 1/2] fix(guardrails): read Responses messages in legacy Lakera and skip bodies it can't read on pass-through A list input of Responses messages now reaches legacy Lakera instead of being dropped, and a pass-through body with no messages or input skips the empty vendor call. Lakera post_call on /gemini has nothing to read, so the rig's Gemini block cells use a generic post_call guardrail. Moves the regression cells into each guardrail's own test file. --- .../guardrails/guardrail_hooks/lakera_ai.py | 22 ++++- .../test_passthrough_auth_guardrails.py | 52 ++++++++++-- .../test_generic_guardrail_api.py | 11 +++ ...guardrails_on_unexpected_request_fields.py | 83 ------------------- .../guardrail_hooks/test_lakera_ai.py | 57 +++++++++++++ .../guardrail_hooks/test_noma_v2.py | 20 +++++ 6 files changed, 156 insertions(+), 89 deletions(-) delete mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/test_guardrails_on_unexpected_request_fields.py create mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index f7eb05e2ab6..228dc417142 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -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 diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index d292fe41059..c8b91077d91 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -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" @@ -76,6 +77,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), @@ -147,6 +152,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( { @@ -236,6 +242,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": { @@ -590,7 +606,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: @@ -630,15 +646,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 diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index b5f794cb136..daa9945e406 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -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" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_guardrails_on_unexpected_request_fields.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_guardrails_on_unexpected_request_fields.py deleted file mode 100644 index f1e1dca04c3..00000000000 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_guardrails_on_unexpected_request_fields.py +++ /dev/null @@ -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", - ) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py new file mode 100644 index 00000000000..afba8f1b39d --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py @@ -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") == ("",) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_noma_v2.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_noma_v2.py index b145b990539..8d430d83787 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_noma_v2.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_noma_v2.py @@ -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"] From 6ef3c0e05605815e6b9989e9b6e96cd6cd7611c2 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 18:05:07 -0800 Subject: [PATCH 2/2] fix(guardrails): read only the admin opt-out settings for a global Lakera guardrail on a Gemini body The opt-out check reran should_run_guardrail on metadata alone, so a streaming-only or tag-scoped global guardrail looked opted out and skipped the Gemini body. It now reads the key's and team's opt-out fields. --- .../guardrail_hooks/lakera_ai_v2.py | 22 +++++------ .../guardrail_hooks/test_lakera_ai_v2.py | 37 +++++++++++++++++++ 2 files changed, 47 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index d818c67d524..5ca64b3c0bd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -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 ) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py index 7a587545807..e9efda2c537 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py @@ -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"}]}]}