From 9fcf2347505c10037d8a6ee65c2f951d612a0f77 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Fri, 1 May 2026 04:22:57 +0000 Subject: [PATCH] fix(guardrails): degrade Lakera v2 mask mode to block on multimodal input MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mask-in-place uses the offsets that Lakera returns for the inspection payload. ``build_inspection_messages`` flattens multimodal content into joined text before sending to Lakera, so the offsets refer to the flattened representation. Writing those offsets back via ``_mask_pii_in_messages`` and overwriting ``data["messages"]`` would silently strip image/audio parts from the original request — that is a real functional regression for Lakera + mask mode + multimodal input. Detect multimodal input (any list-format ``content`` or non-string ``data["input"]``) up front and skip the mask-in-place branch in that case. The hook then falls into the standard block-on-detect path so PII is still blocked but the multimodal payload is never silently rewritten. Per-part masking that preserves multimodal structure is the right long-term fix; tracking that as a follow-up. Also: add ``has_non_string_content`` to ``_content_utils`` (with tests) and a regression test that asserts multimodal+PII raises an HTTPException instead of returning a flattened request body. Co-Authored-By: Claude Opus 4.7 (1M context) --- litellm/proxy/guardrails/_content_utils.py | 24 ++++++++- .../guardrail_hooks/lakera_ai_v2.py | 36 +++++++++---- .../proxy/guardrails/test_content_utils.py | 28 ++++++++++ .../guardrails/test_guardrail_coverage.py | 52 +++++++++++++++++++ 4 files changed, 130 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 576fefe85a7..6c681c322f1 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -86,7 +86,7 @@ def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int: return visit(content) return content if isinstance(content, list): - new_parts = [] + new_parts: List[Any] = [] for part in content: if isinstance(part, str) and part: visited += 1 @@ -143,6 +143,28 @@ def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int: return visited +def has_non_string_content(data: Dict[str, Any]) -> bool: + """Return True if any inspected content is not a plain string. + + Used by hooks whose mask/redact path operates on string offsets and + therefore cannot preserve multimodal non-text parts. Such hooks should + degrade to block-on-detect when this returns True so image/audio parts + are not silently stripped during in-place masking. + """ + messages = data.get("messages") + if isinstance(messages, list): + for message in messages: + if isinstance(message, dict) and not isinstance( + message.get("content"), str + ): + if message.get("content") is not None: + return True + input_value = data.get("input") + if input_value is not None and not isinstance(input_value, str): + return True + return False + + def build_inspection_messages(data: Dict[str, Any]) -> List[Dict[str, str]]: """Synthesize a chat-style messages list for posting to a guardrail API. diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index d3f19966be2..4d6e6a302b5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -13,7 +13,10 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.guardrails._content_utils import build_inspection_messages +from litellm.proxy.guardrails._content_utils import ( + build_inspection_messages, + has_non_string_content, +) from litellm.secret_managers.main import get_secret_str from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import AllMessageValues @@ -223,6 +226,13 @@ class LakeraAIGuardrail(CustomGuardrail): ) return data + # Mask-in-place uses offsets returned by Lakera and can only + # preserve non-text parts (images, audio, …) when the original + # content is a plain string. For multimodal/Responses-API input + # we degrade to block-on-detect so we never silently strip image + # parts while attempting to redact text. + is_multimodal_input = has_non_string_content(data) + ######################################################### ########## 1. Make the Lakera AI v2 guard API request ########## ######################################################### @@ -236,8 +246,11 @@ class LakeraAIGuardrail(CustomGuardrail): ########## 2. Handle flagged content ########## ######################################################### if lakera_guardrail_response.get("flagged") is True: - # If only PII violations exist, mask the PII - if self._is_only_pii_violation(lakera_guardrail_response): + # If only PII violations exist, mask the PII (string input only). + if ( + self._is_only_pii_violation(lakera_guardrail_response) + and not is_multimodal_input + ): data["messages"] = self._mask_pii_in_messages( messages=new_messages, # type: ignore[arg-type] lakera_response=lakera_guardrail_response, @@ -254,7 +267,9 @@ class LakeraAIGuardrail(CustomGuardrail): ) # Log violation but continue elif self.on_flagged == "block": - # If there are other violations or not set to mask PII, raise exception + # Either non-PII violations, or PII on multimodal input + # (which cannot be masked in place without dropping + # image/audio parts) — raise the standard block error. raise self._get_http_exception_for_blocked_guardrail( lakera_guardrail_response ) @@ -289,6 +304,10 @@ class LakeraAIGuardrail(CustomGuardrail): ) return + # See ``async_pre_call_hook`` — multimodal input degrades to + # block-on-detect because mask-in-place would drop image parts. + is_multimodal_input = has_non_string_content(data) + ######################################################### ########## 1. Make the Lakera AI v2 guard API request ########## ######################################################### @@ -302,8 +321,10 @@ class LakeraAIGuardrail(CustomGuardrail): ########## 2. Handle flagged content ########## ######################################################### if lakera_guardrail_response.get("flagged") is True: - # If only PII violations exist, mask the PII - if self._is_only_pii_violation(lakera_guardrail_response): + if ( + self._is_only_pii_violation(lakera_guardrail_response) + and not is_multimodal_input + ): data["messages"] = self._mask_pii_in_messages( messages=new_messages, # type: ignore[arg-type] lakera_response=lakera_guardrail_response, @@ -313,14 +334,11 @@ class LakeraAIGuardrail(CustomGuardrail): "Lakera AI: Masked PII in messages instead of blocking request" ) else: - # Check on_flagged setting if self.on_flagged == "monitor": verbose_proxy_logger.warning( "Lakera Guardrail: Monitoring mode - violation detected but allowing request" ) - # Log violation but continue elif self.on_flagged == "block": - # If there are other violations or not set to mask PII, raise exception raise self._get_http_exception_for_blocked_guardrail( lakera_guardrail_response ) diff --git a/tests/test_litellm/proxy/guardrails/test_content_utils.py b/tests/test_litellm/proxy/guardrails/test_content_utils.py index e26dc7aa2d4..23683abd1cc 100644 --- a/tests/test_litellm/proxy/guardrails/test_content_utils.py +++ b/tests/test_litellm/proxy/guardrails/test_content_utils.py @@ -2,6 +2,7 @@ from litellm.proxy.guardrails._content_utils import ( build_inspection_messages, + has_non_string_content, iter_message_text, walk_user_text, ) @@ -235,3 +236,30 @@ def test_build_inspection_messages_empty_data(): assert build_inspection_messages({}) == [] assert build_inspection_messages({"messages": []}) == [] assert build_inspection_messages({"input": ""}) == [] + + +# ── has_non_string_content ──────────────────────────────────────────────────── + + +def test_has_non_string_content_string_messages(): + data = {"messages": [{"role": "user", "content": "hello"}]} + assert has_non_string_content(data) is False + + +def test_has_non_string_content_multimodal_messages(): + data = {"messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]} + assert has_non_string_content(data) is True + + +def test_has_non_string_content_responses_api_string_input(): + assert has_non_string_content({"input": "plain string"}) is False + + +def test_has_non_string_content_responses_api_list_input(): + assert has_non_string_content({"input": ["a", "b"]}) is True + + +def test_has_non_string_content_empty_data(): + assert has_non_string_content({}) is False + assert has_non_string_content({"messages": []}) is False + assert has_non_string_content({"input": ""}) is False diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py index 4cfcbf3c6d9..a4fb995ed54 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py @@ -161,6 +161,58 @@ async def test_lakera_v2_inspects_responses_api_input(user_api_key, monkeypatch) assert seen_messages == [[{"role": "user", "content": "responses-api content"}]] +@pytest.mark.asyncio +async def test_lakera_v2_multimodal_pii_degrades_to_block(user_api_key, monkeypatch): + """Mask-in-place uses Lakera offsets and cannot preserve image/audio + parts of multimodal input. When PII is detected on a multimodal + request, the hook must raise the block exception instead of silently + flattening ``data["messages"]`` to text-only.""" + monkeypatch.setenv("LAKERA_API_KEY", "lk-test") + from fastapi import HTTPException + + from litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import ( + LakeraAIGuardrail, + ) + + guard = LakeraAIGuardrail(api_key="lk-test", on_flagged="block") + + async def fake_call_v2_guard(messages, request_data, event_type): + return ( + { + "flagged": True, + "payload": [{"detector_type": "pii/email", "start": 0, "end": 5}], + }, + {"EMAIL": 1}, + ) + + with ( + patch.object(guard, "call_v2_guard", side_effect=fake_call_v2_guard), + patch.object(guard, "_is_only_pii_violation", return_value=True), + patch.object( + guard, + "_get_http_exception_for_blocked_guardrail", + return_value=HTTPException(status_code=400, detail="blocked"), + ), + ): + with pytest.raises(HTTPException): + await guard.async_pre_call_hook( + user_api_key_dict=user_api_key, + cache=DualCache(), + data={ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "leak"}, + {"type": "image_url", "image_url": {"url": "..."}}, + ], + } + ] + }, + call_type="acompletion", + ) + + @pytest.mark.asyncio async def test_lakera_v2_inspects_multimodal_list_content(user_api_key, monkeypatch): monkeypatch.setenv("LAKERA_API_KEY", "lk-test")