fix(guardrails): read Gemini contents only on pass-through and Google generateContent calls

litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py

@@ -40,13 +40,22 @@ from litellm.types.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import (
     LakeraAIRequest,
     LakeraAIResponse,
 )
-from litellm.types.utils import CallTypesLiteral, GuardrailStatus, ModelResponse
+from litellm.types.utils import CallTypes, CallTypesLiteral, GuardrailStatus, ModelResponse
 
 _DECODED_JSON: Final = TypeAdapter(object)
 _RESPONSE_FIELDS: Final = TypeAdapter(Mapping[str, object])
 _OBJECT_LIST: Final = TypeAdapter(list[object])
 
 _GEMINI_ROLES: Final[Mapping[str, str]] = MappingProxyType({"user": "user", "model": "assistant", "function": "tool"})
+_GEMINI_BODY_CALL_TYPES: Final = frozenset(
+    {
+        CallTypes.pass_through.value,
+        CallTypes.generate_content.value,
+        CallTypes.agenerate_content.value,
+        CallTypes.generate_content_stream.value,
+        CallTypes.agenerate_content_stream.value,
+    }
+)
 
 _DETECTOR_CATEGORY_PHRASES: Final[Mapping[str, str]] = MappingProxyType(
     {
@@ -203,6 +212,7 @@ def _unmaskable_reason(
     guardrail: "LakeraAIGuardrail",
     data: dict[str, object],
     lakera_response: LakeraAIResponse | None,
+    call_type: str,
 ) -> str | None:
     """Why a PII-only violation on ``data`` can't be masked in place, or None when it can."""
     if has_non_string_content(data):
@@ -211,7 +221,7 @@ def _unmaskable_reason(
         return "messages and input are both present, so the write-back is positionally ambiguous"
     if "messages" in data and not isinstance(data.get("messages"), list):
         return "a messages key that isn't a list, so there's nothing to merge the redacted content into"
-    if _has_gemini_contents(data):
+    if _has_gemini_contents(data, call_type):
         return "Gemini contents, which have no write-back path for the redacted content"
     if not _has_responses_instructions(guardrail, data):
         return "no write-back path for the redacted content"
@@ -235,19 +245,22 @@ def _gemini_text(content: object) -> str:
     return "\n".join(text for text in texts if isinstance(text, str) and text)
 
 
-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``."""
+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`` instead of ``messages``."""
     fields: Final = as_str_mapping(data) or {}
-    return isinstance(fields.get("contents"), list | dict) or isinstance(
-        fields.get("systemInstruction") or fields.get("system_instruction"), dict
+    return call_type in _GEMINI_BODY_CALL_TYPES and (
+        isinstance(fields.get("contents"), list | dict)
+        or isinstance(fields.get("systemInstruction") or fields.get("system_instruction"), dict)
     )
 
 
-def _gemini_inspection_messages(data: Mapping[str, object]) -> Sequence[Mapping[str, str]]:
-    """Chat-style messages for a Gemini ``generateContent`` body: ``systemInstruction`` first as a
-    system message, then each ``contents`` turn with its text parts joined and ``model`` read as
-    ``assistant``. Turns with no text are dropped."""
+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
+    ``model`` read as ``assistant``. Turns with no text are dropped."""
+    if call_type not in _GEMINI_BODY_CALL_TYPES:
+        return []
     system_text: Final = _gemini_text(data.get("systemInstruction") or data.get("system_instruction"))
     turns: Final = (
         {"role": _GEMINI_ROLES.get(str(content.get("role") or "user"), "user"), "content": _gemini_text(content)}
@@ -260,7 +273,9 @@ def _gemini_inspection_messages(data: Mapping[str, object]) -> Sequence[Mapping[
     ]
 
 
-def _build_lakera_inspection_messages(data: Mapping[str, object]) -> Sequence[Mapping[str, str]]:
+def _build_lakera_inspection_messages(
+    data: Mapping[str, object], *, call_type: str | None = None
+) -> Sequence[Mapping[str, str]]:
     """Like build_inspection_messages, but also covers the Responses-API
     ``instructions`` field, placed first since litellm later converts it
     into the model's leading system message and a prompt-injection detector
@@ -281,7 +296,7 @@ def _build_lakera_inspection_messages(data: Mapping[str, object]) -> Sequence[Ma
     )
     return [
         *leading,
-        *_gemini_inspection_messages(data),
+        *_gemini_inspection_messages(data, call_type),
         *build_inspection_messages(dict(data)),
     ]
 
@@ -557,6 +572,7 @@ class LakeraAIGuardrail(CustomGuardrail):
         inspected_messages: Sequence[AllMessageValues],
         lakera_response: LakeraAIResponse | None,
         masked_entity_count: dict[str, int],
+        call_type: str,
     ) -> bool:
         """Mask a body whose only obstacle to mask-in-place is the Responses-API
         ``instructions`` field, writing the redacted instructions straight into
@@ -564,7 +580,7 @@ class LakeraAIGuardrail(CustomGuardrail):
         that field and would fold the instructions text into ``data["input"]``.
         Returns False without masking anything when _unmaskable_reason names an
         obstacle this can't get around."""
-        if _unmaskable_reason(self, data, lakera_response) is not None:
+        if _unmaskable_reason(self, data, lakera_response, call_type) is not None:
             return False
         redacted: Final = self._mask_pii_in_messages(
             messages=inspected_messages,
@@ -596,7 +612,7 @@ class LakeraAIGuardrail(CustomGuardrail):
             return data
 
         # Covers multimodal list content + Responses-API input/instructions.
-        inspection_messages: Final = _build_lakera_inspection_messages(data)
+        inspection_messages: Final = _build_lakera_inspection_messages(data, call_type=call_type)
         if not inspection_messages:
             verbose_proxy_logger.warning("Lakera AI: not running guardrail. No inspectable text in data")
             return data
@@ -623,7 +639,7 @@ 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)  # pyright: ignore[reportUnknownArgumentType]  # hook data is an untyped dict
+        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)
@@ -697,6 +713,7 @@ class LakeraAIGuardrail(CustomGuardrail):
                     inspected_messages=new_messages,
                     lakera_response=lakera_guardrail_response,
                     masked_entity_count=masked_entity_count,
+                    call_type=call_type,
                 )
                 if masked_in_place:
                     verbose_proxy_logger.warning(
@@ -706,7 +723,7 @@ class LakeraAIGuardrail(CustomGuardrail):
                     verbose_proxy_logger.error(
                         "Lakera Guardrail: Monitoring mode - PII detected but NOT masked, forwarding unredacted "
                         "content to the model (reason: %s)",
-                        _unmaskable_reason(self, data, lakera_guardrail_response),
+                        _unmaskable_reason(self, data, lakera_guardrail_response, call_type),
                     )
                 else:
                     verbose_proxy_logger.warning(
@@ -737,7 +754,7 @@ class LakeraAIGuardrail(CustomGuardrail):
             return
 
         # Covers multimodal list content + Responses-API input/instructions.
-        inspection_messages: Final = _build_lakera_inspection_messages(data)
+        inspection_messages: Final = _build_lakera_inspection_messages(data, call_type=call_type)
         if not inspection_messages:
             verbose_proxy_logger.warning("Lakera AI: not running guardrail. No inspectable text in data")
             return

tests/integration/authorization/test_lakera_v2_gemini_native_route.py

@@ -200,3 +200,21 @@ def test_monitor_on_native_route_forwards_the_card_unchanged(rig: NativeGeminiRi
     assert response.status_code == 200, response.text
     assert len(_lakera_messages(rig)) == 1
     assert [_JSON.validate_json(request.body)["contents"] for request in rig.upstream.drain()] == [body["contents"]]
+
+
+def test_a_chat_call_carrying_contents_is_masked_not_blocked(rig: NativeGeminiRig) -> None:
+    gateway: Final = rig.owned.gateway
+    with gateway.scenario() as scenario:
+        key: Final = scenario.key(metadata={"guardrails": ["lakera-native-mask"]})
+        rig.lakera.drain()
+        rig.upstream.drain()
+        response: Final = gateway.request(
+            "POST",
+            "/v1/chat/completions",
+            {"model": MODEL, "messages": [{"role": "user", "content": PII_PROMPT}], "contents": []},
+            key=key,
+        )
+
+    assert response.status_code == 200, response.text
+    assert _lakera_messages(rig) == ((("user", PII_PROMPT),),)
+    assert all(CARD.encode() not in request.body for request in rig.upstream.drain())

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

@@ -27,7 +27,7 @@ from litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import (
     humanize_lakera_block_reasons,
 )
 from litellm.types.guardrails import LitellmParams, Mode
-from litellm.types.utils import ModelResponse
+from litellm.types.utils import CallTypesLiteral, ModelResponse
 from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
 
 
@@ -1736,6 +1736,11 @@ GEMINI_INSPECTION_MESSAGES: Final = [
     {"role": "user", "content": "hello world"},
     {"role": "assistant", "content": "hi"},
 ]
+EMAIL_AT_START_LAKERA_RESPONSE: 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}],
+}
 PROMPT_ATTACK_ONLY_LAKERA_RESPONSE: Final = {
     "flagged": True,
     "breakdown": [{"detector_type": "prompt_attack", "detected": True, "message_id": 1}],
@@ -1743,7 +1748,10 @@ PROMPT_ATTACK_ONLY_LAKERA_RESPONSE: Final = {
 
 
 async def _run_gemini_pre_call(
-    guardrail: LakeraAIGuardrail, lakera_response: Mapping[str, object], data: dict[str, object]
+    guardrail: LakeraAIGuardrail,
+    lakera_response: Mapping[str, object],
+    data: dict[str, object],
+    call_type: CallTypesLiteral = "pass_through_endpoint",
 ) -> Sequence[object]:
     """Run the pre-call hook against a Lakera double and return the messages each guard request carried."""
     sent: Final = list[object]()  # mutable-ok: the Lakera double records each guard request it receives
@@ -1757,7 +1765,7 @@ async def _run_gemini_pre_call(
         user_api_key_dict=UserAPIKeyAuth(api_key="test_key"),
         cache=MagicMock(),
         data=data,
-        call_type="pass_through_endpoint",
+        call_type=call_type,
     )
     return sent
 
@@ -1769,18 +1777,23 @@ def _without_metadata(data: Mapping[str, object]) -> Mapping[str, object]:
 
 class TestGeminiContents:
     def test_reads_system_instruction_and_text_parts(self):
-        assert _build_lakera_inspection_messages(_gemini_request()) == GEMINI_INSPECTION_MESSAGES
+        assert (
+            _build_lakera_inspection_messages(_gemini_request(), call_type="pass_through_endpoint")
+            == 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"}]
+        assert _build_lakera_inspection_messages(data, call_type="pass_through_endpoint") == [
+            {"role": "user", "content": "hello world"}
+        ]
 
     def test_reads_snake_case_system_instruction_and_joins_text_parts(self):
         data: Final = {
             "system_instruction": {"parts": [{"text": "a"}, {"text": "b"}]},
             "contents": [{"parts": [{"text": "c"}]}, {"role": "function", "parts": [{"text": "d"}]}],
         }
-        assert _build_lakera_inspection_messages(data) == [
+        assert _build_lakera_inspection_messages(data, call_type="pass_through_endpoint") == [
             {"role": "system", "content": "a\nb"},
             {"role": "user", "content": "c"},
             {"role": "tool", "content": "d"},
@@ -1789,7 +1802,41 @@ class TestGeminiContents:
     @pytest.mark.parametrize("chat_field", ["messages", "input"])
     def test_a_chat_field_next_to_contents_does_not_hide_the_contents(self, chat_field: str):
         data: Final = {chat_field: [], "contents": [{"parts": [{"text": "card"}]}]}
-        assert _build_lakera_inspection_messages(data) == [{"role": "user", "content": "card"}]
+        assert _build_lakera_inspection_messages(data, call_type="pass_through_endpoint") == [
+            {"role": "user", "content": "card"}
+        ]
+
+    @pytest.mark.parametrize("call_type", ["agenerate_content", "agenerate_content_stream"])
+    def test_reads_contents_on_the_google_generate_content_routes(self, call_type: str):
+        assert _build_lakera_inspection_messages(_gemini_request(), call_type=call_type) == GEMINI_INSPECTION_MESSAGES
+
+    @pytest.mark.parametrize("call_type", ["acompletion", "aresponses", "anthropic_messages", None])
+    def test_contents_on_another_call_type_are_not_read(self, call_type: str | None):
+        data: Final = {**_gemini_request(), "messages": [USER_MSG]}
+        assert _build_lakera_inspection_messages(data, call_type=call_type) == [USER_MSG]
+
+    @pytest.mark.asyncio
+    async def test_a_chat_call_carrying_contents_is_still_masked(self):
+        data: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "a@b.com hi"}], "contents": []}
+        sent: Final = await _run_gemini_pre_call(
+            LakeraAIGuardrail(api_key="test_key", on_flagged="mask"),
+            EMAIL_AT_START_LAKERA_RESPONSE,
+            data,
+            "acompletion",
+        )
+        assert sent == [[{"role": "user", "content": "a@b.com hi"}]]
+        assert data["messages"] == [{"role": "user", "content": "[MASKED EMAIL] hi"}]
+
+    @pytest.mark.asyncio
+    async def test_monitor_still_masks_responses_instructions_next_to_contents(self):
+        data: Final[dict[str, object]] = {"instructions": "a@b.com is the contact", "input": "hi", "contents": []}
+        await _run_gemini_pre_call(
+            LakeraAIGuardrail(api_key="test_key", on_flagged="monitor"),
+            EMAIL_AT_START_LAKERA_RESPONSE,
+            data,
+            "responses",
+        )
+        assert data["instructions"] == "[MASKED EMAIL] is the contact"
 
     @pytest.mark.asyncio
     async def test_a_body_with_messages_and_contents_is_blocked_rather_than_masked(self):