fix(guardrails): keep global opt-outs and Responses instructions on Lakera Gemini pass-through bodies

litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py

@@ -279,11 +279,14 @@ def _gemini_tool_text(content: object) -> str:
 def _has_gemini_contents(data: object, call_type: str | None) -> bool:
     """True if ``data`` is a Gemini ``generateContent`` body on a pass-through or Google ``generateContent`` call,
     which carries its prompt in ``contents`` and ``systemInstruction``. A pass-through body with chat ``messages`` or
-    Responses ``input`` is read as chat, as before."""
+    Responses ``input`` or ``instructions`` is read as chat, as before."""
     fields: Final = as_str_mapping(data) or {}
     return (
         call_type in _GEMINI_BODY_CALL_TYPES
-        and not (call_type == CallTypes.pass_through.value and (fields.get("messages") or fields.get("input")))
+        and not (
+            call_type == CallTypes.pass_through.value
+            and (fields.get("messages") or fields.get("input") or fields.get("instructions"))
+        )
         and (
             isinstance(fields.get("contents"), list | dict)
             or isinstance(fields.get("systemInstruction") or fields.get("system_instruction"), dict)
@@ -291,6 +294,25 @@ def _has_gemini_contents(data: object, call_type: str | None) -> bool:
     )
 
 
+def _opted_out_of_global_guardrail(guardrail: "LakeraAIGuardrail", user_api_key_dict: UserAPIKeyAuth) -> bool:
+    """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
+    )
+
+
 def _gemini_inspection_messages(data: Mapping[str, object], call_type: str | None) -> Sequence[Mapping[str, str]]:
     """Chat-style messages for a Gemini ``generateContent`` body on a pass-through or Google ``generateContent``
     call: ``systemInstruction`` first as a system message, then each ``contents`` turn with its text parts joined and
@@ -647,7 +669,10 @@ class LakeraAIGuardrail(CustomGuardrail):
         verbose_proxy_logger.debug("Lakera AI: pre_call_hook")
 
         event_type: Final[GuardrailEventHooks] = GuardrailEventHooks.pre_call
-        if self.should_run_guardrail(data=data, event_type=event_type) is not True:
+        is_gemini_body: Final = _has_gemini_contents(data, call_type)  # pyright: ignore[reportUnknownArgumentType]  # hook data is an untyped dict
+        if self.should_run_guardrail(data=data, event_type=event_type) is not True or (
+            is_gemini_body and _opted_out_of_global_guardrail(self, user_api_key_dict)
+        ):
             verbose_proxy_logger.debug("Lakera AI: not running guardrail. Guardrail is disabled.")
             return data
 
@@ -679,7 +704,6 @@ class LakeraAIGuardrail(CustomGuardrail):
         # handled safely by _apply_redacted_messages_back_preserving_fields's
         # scope-index merge, which never touches a message outside the scope
         # it actually redacted instead of reconstructing the list from scratch.
-        is_gemini_body: Final = _has_gemini_contents(data, call_type)  # pyright: ignore[reportUnknownArgumentType]  # hook data is an untyped dict
         is_multimodal_input: Final = (
             has_non_string_content(data)
             or _has_combined_messages_and_input(data)

tests/integration/authorization/test_lakera_v2_gemini_global_opt_out.py

@@ -0,0 +1,121 @@
+"""A key or team opted out of a global Lakera guardrail keeps skipping it on a Gemini pass-through body."""
+
+import json
+from collections.abc import Callable, Iterator
+from contextlib import ExitStack
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Final
+
+import pytest
+import yaml
+from integration._support.client import Gateway, Scenario
+from integration._support.process import OwnedProxy, owned_proxy_process
+from integration._support.wire import Reply, Request, Wire, wire_server
+from pydantic import JsonValue, TypeAdapter
+
+GEMINI_ROUTE: Final = "/gemini/v1beta/models/gemini-2.5-flash:generateContent"
+CARD: Final = "4111 1111 1111 1111"
+GLOBAL_LAKERA: Final = "lakera-global"
+_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
+_JSON: Final = TypeAdapter(dict[str, JsonValue])
+_BODY: Final[dict[str, JsonValue]] = {"contents": [{"role": "user", "parts": [{"text": f"My card is {CARD}"}]}]}
+
+
+def _lakera_verdict(request: Request) -> Reply:
+    flagged: Final = CARD.encode() in request.body
+    breakdown: Final = [{"detector_type": "pii/credit_card", "detected": flagged, "message_id": 0}]
+    return Reply(body=json.dumps({"flagged": flagged, "breakdown": breakdown, "payload": []}).encode())
+
+
+def _gemini_reply(_request: Request) -> Reply:
+    return Reply(body=json.dumps({"candidates": [{"content": {"role": "model", "parts": [{"text": "Hi"}]}}]}).encode())
+
+
+@dataclass(frozen=True, slots=True)
+class GlobalLakeraRig:
+    owned: OwnedProxy
+    lakera: Wire
+    upstream: Wire
+
+
+@pytest.fixture(scope="module")
+def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[GlobalLakeraRig]:
+    directory: Final = tmp_path_factory.mktemp("lakera-gemini-global-opt-out")
+    with ExitStack() as stack:
+        lakera: Final = stack.enter_context(wire_server(_lakera_verdict))
+        upstream: Final = stack.enter_context(wire_server(_gemini_reply))
+        config: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text()))
+        path: Final = directory / "lakera-gemini-global-opt-out.yaml"
+        path.write_text(
+            yaml.safe_dump(
+                {
+                    **config,
+                    "guardrails": [
+                        {
+                            "guardrail_name": GLOBAL_LAKERA,
+                            "litellm_params": {
+                                "guardrail": "lakera_v2",
+                                "mode": "pre_call",
+                                "default_on": True,
+                                "on_flagged": "block",
+                                "api_base": lakera.url,
+                                "api_key": "synthetic-lakera-key",
+                            },
+                        }
+                    ],
+                }
+            )
+        )
+        environment: Final = {"GEMINI_API_BASE": upstream.url, "GEMINI_API_KEY": "synthetic-gemini-key"}
+        owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path))
+        yield GlobalLakeraRig(owned, lakera, upstream)
+
+
+def _key_disabling_global(scenario: Scenario) -> str:
+    return scenario.key(disable_global_guardrails=True)
+
+
+def _key_naming_it(scenario: Scenario) -> str:
+    return scenario.key(metadata={"opted_out_global_guardrails": [GLOBAL_LAKERA]})
+
+
+def _team_disabling_global(scenario: Scenario) -> str:
+    return scenario.key(team_id=scenario.team(disable_global_guardrails=True))
+
+
+def _team_naming_it(scenario: Scenario) -> str:
+    return scenario.key(team_id=scenario.team(metadata={"opted_out_global_guardrails": [GLOBAL_LAKERA]}))
+
+
+@pytest.mark.parametrize(
+    "opted_out_key", [_key_disabling_global, _key_naming_it, _team_disabling_global, _team_naming_it]
+)
+def test_opted_out_key_skips_the_global_guardrail_on_a_gemini_body(
+    rig: GlobalLakeraRig, opted_out_key: Callable[[Scenario], str]
+) -> None:
+    gateway: Final = rig.owned.gateway
+    with gateway.scenario() as scenario:
+        key: Final = opted_out_key(scenario)
+        rig.lakera.drain()
+        rig.upstream.drain()
+
+        response: Final = gateway.request("POST", GEMINI_ROUTE, _BODY, key=key, headers={"x-goog-api-key": key})
+
+        assert response.status_code == 200, response.text
+        assert rig.lakera.drain() == ()
+        assert len(rig.upstream.drain()) == 1
+
+
+def test_key_that_did_not_opt_out_is_inspected_by_the_global_guardrail(rig: GlobalLakeraRig) -> None:
+    gateway: Final = rig.owned.gateway
+    with gateway.scenario() as scenario:
+        key: Final = scenario.key()
+        rig.lakera.drain()
+        rig.upstream.drain()
+
+        response: Final = gateway.request("POST", GEMINI_ROUTE, _BODY, key=key, headers={"x-goog-api-key": key})
+
+        assert response.status_code == 400, response.text
+        assert len(rig.lakera.drain()) == 1
+        assert rig.upstream.drain() == ()

tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py

@@ -1775,6 +1775,12 @@ def _without_metadata(data: Mapping[str, object]) -> Mapping[str, object]:
     return {key: value for key, value in data.items() if key != "metadata"}
 
 
+GEMINI_FLAGGED: Final = {
+    "flagged": True,
+    "breakdown": [{"detector_type": "prompt_attack", "detected": True, "message_id": 0}],
+}
+
+
 class TestGeminiContents:
     def test_reads_system_instruction_and_text_parts(self):
         assert (
@@ -1855,6 +1861,85 @@ class TestGeminiContents:
             {"role": "user", "content": "card"}
         ]
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "gemini_fields",
+        [
+            pytest.param({"contents": []}, id="contents-empty"),
+            pytest.param({"contents": [{"parts": [{"text": "hello"}]}]}, id="contents-benign"),
+            pytest.param({"systemInstruction": {"parts": [{"text": "hello"}]}}, id="system-instruction"),
+        ],
+    )
+    async def test_monitor_masks_instructions_on_a_pass_through_body_with_gemini_fields(
+        self, gemini_fields: dict[str, object]
+    ):
+        guardrail: Final = LakeraAIGuardrail(api_key="test_key", on_flagged="monitor")
+        data: Final[dict[str, object]] = {  # mutable-ok: the hook masks the body in place
+            "instructions": "a@b.com is the contact",
+            "input": "",
+            **gemini_fields,
+        }
+        flagged: Final = {
+            "flagged": True,
+            "breakdown": [{"detector_type": "pii/email", "detected": True, "message_id": 0}],
+            "payload": [{"detector_type": "pii/email", "start": 0, "end": 7, "message_id": 0}],
+        }
+        with patch.object(guardrail, "call_v2_guard", new=AsyncMock(return_value=(flagged, {}))):
+            await guardrail.async_pre_call_hook(
+                user_api_key_dict=UserAPIKeyAuth(api_key="test_key"),
+                cache=DualCache(),
+                data=data,
+                call_type="pass_through_endpoint",
+            )
+        assert data["instructions"] == "[MASKED EMAIL] is the contact"
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "key_fields",
+        [
+            pytest.param({"metadata": {"disable_global_guardrails": True}}, id="key-disables-global"),
+            pytest.param({"metadata": {"opted_out_global_guardrails": ["lakera-global"]}}, id="key-names-it"),
+            pytest.param({"team_metadata": {"disable_global_guardrails": True}}, id="team-disables-global"),
+            pytest.param({"team_metadata": {"opted_out_global_guardrails": ["lakera-global"]}}, id="team-names-it"),
+        ],
+    )
+    async def test_key_opted_out_of_a_global_guardrail_skips_it_on_a_gemini_pass_through_body(
+        self, key_fields: dict[str, dict[str, object]]
+    ):
+        guardrail: Final = LakeraAIGuardrail(api_key="test_key", guardrail_name="lakera-global", default_on=True)
+        data: Final = _gemini_request()
+        with patch.object(guardrail, "call_v2_guard", new=AsyncMock(return_value=(GEMINI_FLAGGED, {}))):
+            result: Final = await guardrail.async_pre_call_hook(
+                user_api_key_dict=UserAPIKeyAuth(api_key="test_key", **key_fields),
+                cache=DualCache(),
+                data=data,
+                call_type="pass_through_endpoint",
+            )
+        assert result == data
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "key_fields",
+        [
+            pytest.param({}, id="no-opt-out"),
+            pytest.param({"metadata": {"opted_out_global_guardrails": ["another-guardrail"]}}, id="another-named"),
+        ],
+    )
+    async def test_key_not_opted_out_of_a_global_guardrail_is_blocked_on_a_gemini_pass_through_body(
+        self, key_fields: dict[str, dict[str, object]]
+    ):
+        guardrail: Final = LakeraAIGuardrail(api_key="test_key", guardrail_name="lakera-global", default_on=True)
+        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", **key_fields),
+                cache=DualCache(),
+                data=_gemini_request(),
+                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"}]}]}