From 9df7aa5cdb60fe435d8b17ae3f843813eb9475ae Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 01:40:19 -0700 Subject: [PATCH 1/2] fix(proxy): keep inherited post-call guardrails flat so model guardrails merge with them --- .../pass_through_endpoints.py | 2 +- .../test_passthrough_auth_guardrails.py | 67 +++++++++++++++++-- .../test_pass_through_endpoints.py | 16 ++++- 3 files changed, 78 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 07353793b86..cb37c6f30c3 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -888,7 +888,7 @@ async def _inherit_auth_guardrails( merged: Final = (as_str_mapping(parsed_body["metadata"]) or {}).get("guardrails") if not isinstance(merged, list) or not merged or not _named_guardrail_runs_post_call(parsed_body): return None - return tuple(_GUARDRAIL_ENTRIES.validate_python(merged)) + return _GUARDRAIL_ENTRIES.validate_python(merged) def _named_guardrail_runs_post_call(data: Mapping[str, object]) -> bool: diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index 09b20bf6653..221b9b4c885 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -19,12 +19,15 @@ from pydantic import JsonValue, TypeAdapter GUARDRAIL: Final = "passthrough-auth-guardrail" GLOBAL_POST_CALL: Final = "passthrough-global-post-call" +KEY_POST_CALL: Final = "passthrough-key-post-call" +GUARDED_MODEL: Final = "passthrough-guarded-model" BLOCK_MARKER: Final = "block-this-prompt" OPENAI_ROUTE: Final = "/openai/v1/chat/completions" ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" GEMINI_ROUTE: Final = "/gemini/v1beta/models/gemini-2.5-flash:generateContent" _CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" _JSON: Final = TypeAdapter(dict[str, JsonValue]) +_JSON_LIST: Final = TypeAdapter(list[JsonValue]) Level = Literal["key", "team", "project"] @@ -53,8 +56,20 @@ def _openai_reply(_request: Request) -> Reply: return Reply(body=_OPENAI_BODY) -def _vendor_down(_request: Request) -> Reply: - return Reply(status=500, body=b'{"error": "guardrail vendor down"}') +def _lakera_clean(_request: Request) -> Reply: + return Reply(body=json.dumps({"flagged": False, "breakdown": [], "payload": []}).encode()) + + +def _lakera_flagged(_request: Request) -> Reply: + return Reply( + body=json.dumps( + { + "flagged": True, + "breakdown": [{"detector_type": "prompt_attack", "detected": True, "message_id": 0}], + "payload": [], + } + ).encode() + ) def _anthropic_reply(_request: Request) -> Reply: @@ -77,6 +92,7 @@ class AuthGuardrailRig: owned: OwnedProxy vendor: Wire post_call_vendor: Wire + key_post_call_vendor: Wire upstreams: Mapping[str, Wire] @@ -85,7 +101,8 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I directory: Final = tmp_path_factory.mktemp("passthrough-auth-guardrails") with ExitStack() as stack: vendor: Final = stack.enter_context(wire_server(_verdict)) - post_call_vendor: Final = stack.enter_context(wire_server(_vendor_down)) + post_call_vendor: Final = stack.enter_context(wire_server(_lakera_clean)) + key_post_call_vendor: Final = stack.enter_context(wire_server(_lakera_flagged)) upstreams: Final = MappingProxyType( { OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)), @@ -99,6 +116,17 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I yaml.safe_dump( { **config, + "model_list": [ + *_JSON_LIST.validate_python(config.get("model_list") or []), + { + "model_name": GUARDED_MODEL, + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "synthetic-openai-key", + "guardrails": [GUARDRAIL], + }, + }, + ], "guardrails": [ { "guardrail_name": GUARDRAIL, @@ -110,10 +138,20 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "api_key": "synthetic-guardrail-key", }, }, + { + "guardrail_name": KEY_POST_CALL, + "litellm_params": { + "guardrail": "lakera_v2", + "mode": "post_call", + "default_on": False, + "api_base": key_post_call_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, { "guardrail_name": GLOBAL_POST_CALL, "litellm_params": { - "guardrail": "generic_guardrail_api", + "guardrail": "lakera_v2", "mode": "post_call", "default_on": True, "api_base": post_call_vendor.url, @@ -133,7 +171,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "GEMINI_API_KEY": "synthetic-gemini-key", } owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path)) - yield AuthGuardrailRig(owned, vendor, post_call_vendor, upstreams) + yield AuthGuardrailRig(owned, vendor, post_call_vendor, key_post_call_vendor, upstreams) def _key_with_guardrail(scenario: Scenario, level: Level) -> str: @@ -155,6 +193,7 @@ def _send(gateway: Gateway, route: str, body: Mapping[str, JsonValue], key: str) def _drain(rig: AuthGuardrailRig) -> None: rig.vendor.drain() rig.post_call_vendor.drain() + rig.key_post_call_vendor.drain() for upstream in rig.upstreams.values(): upstream.drain() @@ -228,3 +267,21 @@ def test_pre_call_guardrail_relays_the_provider_reply_untouched(rig: AuthGuardra assert response.content == _OPENAI_BODY assert len(rig.vendor.drain()) == 1 assert rig.post_call_vendor.drain() == () + + +def test_key_post_call_guardrail_runs_on_a_model_with_its_own_guardrails(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, + OPENAI_ROUTE, + {**_route_body(OPENAI_ROUTE, "What is the capital of France?"), "model": GUARDED_MODEL}, + key, + ) + + assert response.status_code == 400, response.text + assert len(rig.key_post_call_vendor.drain()) == 1 + assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1 diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 3ebd2f32953..2f2955d063f 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -68,7 +68,7 @@ from litellm.proxy.pass_through_endpoints.success_handler import ( ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import ProxyLogging, _merge_guardrails_with_existing from litellm.types import utils as types_utils from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypes @@ -11829,6 +11829,20 @@ class TestInheritAuthGuardrails: "team-guard", ] + @pytest.mark.asyncio + async def test_post_call_guardrails_stay_named_after_a_model_guardrail_merge(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + _, post_call = await self._run( + self._chat_body(), + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("key-guard", GuardrailEventHooks.post_call),), + ) + + merged: Final = _merge_guardrails_with_existing(post_call.call_args.kwargs["data"], ["model-guard"]) + assert merged["metadata"]["guardrails"] == ["key-guard", "model-guard"] + @pytest.mark.asyncio async def test_pre_call_only_guardrails_relay_the_response_untouched(self): auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) From f697e2ec959097e3eaf98df458cc7414f77551ce Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 01:49:08 -0700 Subject: [PATCH 2/2] fix(guardrails): read a single Gemini contents or parts object in the Lakera v2 guardrail --- .../guardrail_hooks/lakera_ai_v2.py | 5 +- .../test_lakera_v2_gemini_contents.py | 10 +++ .../guardrail_hooks/test_lakera_ai_v2.py | 66 ++++++++++++------- 3 files changed, 56 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 2cc3e5383ca..456a3697219 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -221,6 +221,9 @@ def _unmaskable_reason( def _gemini_items(value: object) -> Sequence[object]: + """A repeated Gemini field as a sequence: a JSON array, or the single object Google also accepts in its place.""" + if isinstance(value, dict): + return (value,) return _OBJECT_LIST.validate_python(value) if isinstance(value, list) else () @@ -235,7 +238,7 @@ def _has_gemini_contents(data: object) -> bool: """True if ``data`` is a Gemini ``generateContent`` body, which carries its prompt in ``contents`` and ``systemInstruction`` instead of ``messages``.""" fields: Final = as_str_mapping(data) or {} - return isinstance(fields.get("contents"), list) or isinstance( + return isinstance(fields.get("contents"), list | dict) or isinstance( fields.get("systemInstruction") or fields.get("system_instruction"), dict ) diff --git a/tests/integration/authorization/test_lakera_v2_gemini_contents.py b/tests/integration/authorization/test_lakera_v2_gemini_contents.py index 63eb05fea64..265e8a47376 100644 --- a/tests/integration/authorization/test_lakera_v2_gemini_contents.py +++ b/tests/integration/authorization/test_lakera_v2_gemini_contents.py @@ -184,3 +184,13 @@ def test_monitor_mode_forwards_the_body_unchanged(rig: GeminiLakeraRig) -> None: assert response.status_code == 200, response.text assert len(_lakera_messages(rig)) == 1 assert [_JSON.validate_json(request.body) for request in rig.upstream.drain()] == [body] + + +def test_card_in_a_single_contents_object_is_blocked(rig: GeminiLakeraRig) -> None: + body: Final[dict[str, JsonValue]] = {"contents": {"role": "user", "parts": {"text": PII_PROMPT}}} + + response: Final = _send(rig, "block", body) + + assert response.status_code == 400, response.text + assert _lakera_messages(rig) == ((("user", PII_PROMPT),),) + assert rig.upstream.drain() == () 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 c0f805e6431..350ebd80733 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 @@ -5,16 +5,18 @@ PR checklist requires at least one test in tests/test_litellm/. Additional tests live in tests/guardrails_tests/test_lakera_v2.py. """ -import importlib, logging, os -from collections.abc import Mapping +import importlib, json, logging, os +from collections.abc import Mapping, Sequence from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import HTTPException import litellm from litellm.caching.caching import DualCache +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.base_llm.guardrail_translation.utils import ( filter_messages_by_skip_flags, ) @@ -1742,23 +1744,39 @@ PROMPT_ATTACK_ONLY_LAKERA_RESPONSE: Final = { async def _run_gemini_pre_call( guardrail: LakeraAIGuardrail, lakera_response: Mapping[str, object], data: dict[str, object] -) -> AsyncMock: - with patch.object(guardrail, "call_v2_guard", new_callable=AsyncMock, return_value=(lakera_response, {})) as call: - await guardrail.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), - cache=MagicMock(), - data=data, - call_type="pass_through_endpoint", - ) - return call +) -> Sequence[object]: + """Run the pre-call hook against a Lakera double and return the messages each guard request carried.""" + sent: Final = list[object]() + + def respond(request: httpx.Request) -> httpx.Response: + sent.append(json.loads(request.content)["messages"]) + return httpx.Response(200, json={"payload": [], **lakera_response}) + + guardrail.async_handler = AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + cache=MagicMock(), + data=data, + call_type="pass_through_endpoint", + ) + return sent + + +def _without_metadata(data: Mapping[str, object]) -> Mapping[str, object]: + """The request body as forwarded, without the guardrail logging the hook records under ``metadata``.""" + return {key: value for key, value in data.items() if key != "metadata"} class TestGeminiContents: def test_reads_system_instruction_and_text_parts(self): assert _build_lakera_inspection_messages(_gemini_request()) == GEMINI_INSPECTION_MESSAGES + def test_reads_a_single_contents_object_with_a_single_parts_object(self): + data: Final = {"contents": {"role": "user", "parts": {"text": "hello world"}}} + assert _build_lakera_inspection_messages(data) == [{"role": "user", "content": "hello world"}] + def test_reads_snake_case_system_instruction_and_joins_text_parts(self): - data = { + data: Final = { "system_instruction": {"parts": [{"text": "a"}, {"text": "b"}]}, "contents": [{"parts": [{"text": "c"}]}, {"role": "function", "parts": [{"text": "d"}]}], } @@ -1770,41 +1788,41 @@ class TestGeminiContents: @pytest.mark.asyncio async def test_clean_body_reaches_lakera_and_is_forwarded_unchanged(self): - data = _gemini_request() - call = await _run_gemini_pre_call(LakeraAIGuardrail(api_key="test_key"), {"flagged": False}, data) - assert list(call.call_args.kwargs["messages"]) == GEMINI_INSPECTION_MESSAGES - assert data == _gemini_request() + data: Final = _gemini_request() + sent: Final = await _run_gemini_pre_call(LakeraAIGuardrail(api_key="test_key"), {"flagged": False}, data) + assert sent == [GEMINI_INSPECTION_MESSAGES] + assert _without_metadata(data) == _gemini_request() @pytest.mark.asyncio @pytest.mark.parametrize("on_flagged", ["block", "mask"]) - async def test_pii_is_blocked_since_contents_cannot_be_masked(self, on_flagged): - data = _gemini_request() + async def test_pii_is_blocked_since_contents_cannot_be_masked(self, on_flagged: str): + data: Final = _gemini_request() with pytest.raises(HTTPException) as blocked: await _run_gemini_pre_call( LakeraAIGuardrail(api_key="test_key", on_flagged=on_flagged), PII_ONLY_LAKERA_RESPONSE, data ) assert blocked.value.status_code == 400 - assert data == _gemini_request() + assert _without_metadata(data) == _gemini_request() @pytest.mark.asyncio @pytest.mark.parametrize("lakera_response", [PROMPT_ATTACK_ONLY_LAKERA_RESPONSE, PII_ONLY_LAKERA_RESPONSE]) - async def test_advisory_blocks_instead_of_adding_a_messages_field(self, lakera_response): - data = _gemini_request() + async def test_advisory_blocks_instead_of_adding_a_messages_field(self, lakera_response: Mapping[str, object]): + data: Final = _gemini_request() with pytest.raises(HTTPException) as blocked: await _run_gemini_pre_call( LakeraAIGuardrail(api_key="test_key", on_flagged="inject_system_message"), lakera_response, data ) assert blocked.value.status_code == 400 - assert data == _gemini_request() + assert _without_metadata(data) == _gemini_request() @pytest.mark.asyncio async def test_monitor_forwards_the_body_and_logs_why_it_was_not_masked(self): - data = _gemini_request() + data: Final = _gemini_request() with patch("litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2.verbose_proxy_logger") as logger: await _run_gemini_pre_call( LakeraAIGuardrail(api_key="test_key", on_flagged="monitor"), PII_ONLY_LAKERA_RESPONSE, data ) - assert data == _gemini_request() + assert _without_metadata(data) == _gemini_request() assert logger.error.call_args.args[1] == ( "Gemini contents, which have no write-back path for the redacted content" )