diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index faecb8e31e6..f9942004aff 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/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) diff --git a/tests/integration/authorization/test_lakera_v2_gemini_global_opt_out.py b/tests/integration/authorization/test_lakera_v2_gemini_global_opt_out.py new file mode 100644 index 00000000000..adc3a0edd82 --- /dev/null +++ b/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() == () 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 e8785551e4f..c83c976f9db 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 @@ -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"}]}]}