diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 1b3571cf242..0c0637aa8f0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -138,7 +138,7 @@ def _pre_masking_scope_indices( def _apply_redacted_messages_back_preserving_fields( guardrail: "LakeraAIGuardrail", data: dict[str, object], # mutable-ok: writes the redacted result back into the caller's request dict in place - redacted_messages: Sequence[Mapping[str, str]], + redacted_messages: Sequence[AllMessageValues], ) -> None: """Write masked content back to ``data["messages"]`` without losing fields the synthetic role/content-only ``redacted_messages`` never carried (e.g. a @@ -346,7 +346,7 @@ class LakeraAIGuardrail(CustomGuardrail): async def call_v2_guard( self, - messages: list[AllMessageValues], + messages: Sequence[AllMessageValues], request_data: dict, event_type: GuardrailEventHooks, ) -> tuple[LakeraAIResponse, dict]: @@ -408,10 +408,10 @@ class LakeraAIGuardrail(CustomGuardrail): def _mask_pii_in_messages( self, - messages: list[AllMessageValues], + messages: Sequence[AllMessageValues], lakera_response: LakeraAIResponse | None, masked_entity_count: dict, - ) -> list[AllMessageValues]: + ) -> Sequence[AllMessageValues]: """ Return a copy of messages with any detected PII replaced by “[MASKED ]” tokens. @@ -540,12 +540,30 @@ class LakeraAIGuardrail(CustomGuardrail): _apply_redacted_messages_back_preserving_fields(self, data, redacted_messages) verbose_proxy_logger.debug("Lakera AI: Masked PII in messages instead of blocking request") elif self.on_flagged == "inject_system_message": + if is_multimodal_input: + # Nothing here can be safely masked, so an advisory note next + # to this raw, unredacted content would be no safer than a + # note next to nothing. Degrade to blocking instead, same as + # this on_flagged setting already does when the advisory + # itself has no field it can be delivered into. + raise self._get_http_exception_for_blocked_guardrail(lakera_guardrail_response) + # A mixed violation (PII plus something else, e.g. prompt injection): + # mask whatever Lakera returned location data for before advising + # about what remains, so the advisory is never shown next to raw + # PII that could have been redacted. + mixed_redacted_messages: Final = self._mask_pii_in_messages( + messages=new_messages, + lakera_response=lakera_guardrail_response, + masked_entity_count=masked_entity_count, + ) + _apply_redacted_messages_back_preserving_fields(self, data, mixed_redacted_messages) advisory_delivered: Final = self.inject_advisory_message( data, self._build_advisory_message(lakera_guardrail_response) ) if advisory_delivered: verbose_proxy_logger.warning( - "Lakera Guardrail: Advisory mode - violation detected, appended advisory system message" + "Lakera Guardrail: Advisory mode - violation detected, masked PII and appended advisory " + "system message" ) else: # Structured Responses-API input (a list, not a plain string) @@ -636,6 +654,17 @@ class LakeraAIGuardrail(CustomGuardrail): _apply_redacted_messages_back_preserving_fields(self, data, redacted_messages) verbose_proxy_logger.debug("Lakera AI: Masked PII in messages instead of blocking request") elif self.on_flagged == "inject_system_message": + if not is_multimodal_input: + # A mixed violation (PII plus something else): mask whatever's + # maskable even though the advisory note below has no effect + # here, so raw PII doesn't pass through untouched just because + # this violation wasn't PII-only. + mixed_redacted_messages: Final = self._mask_pii_in_messages( + messages=new_messages, + lakera_response=lakera_guardrail_response, + masked_entity_count=masked_entity_count, + ) + _apply_redacted_messages_back_preserving_fields(self, data, mixed_redacted_messages) # during_call runs concurrently with the LLM dispatch (see # ProxyLogging.during_call_hook / common_request_processing.py), # with no pre-call barrier -- mutating data["messages"] here races diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py index 0cc27ca169f..977d26e1a40 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py @@ -930,6 +930,83 @@ class TestAdvisoryModeWiring: assert result["messages"][0]["content"] != original_content assert len(result["messages"]) == 1, "no advisory note should be appended once PII is masked" + @pytest.mark.asyncio + async def test_pre_call_mixed_violation_masks_pii_before_appending_advisory(self): + """ + Bugbot finding on BerriAI/litellm#34940: a mixed violation (PII plus a + non-PII flag like prompt injection) isn't PII-only, so it fell straight + through to the advisory branch with the raw PII still in place. It must + mask the maskable PII first, then still append the advisory note for the + remaining, non-PII concern. + """ + lakera_guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="inject_system_message") + mock_response = { + "flagged": True, + "payload": [{"detector_type": "pii/email", "start": 11, "end": 26, "message_id": 0}], + "breakdown": [ + {"detector_type": "pii/email", "detected": True}, + {"detector_type": "prompt_injection", "detected": True}, + ], + } + original_content = "My email is test@example.com" + + with patch.object(lakera_guardrail, "call_v2_guard", new_callable=AsyncMock) as mock_call: + mock_call.return_value = (mock_response, {}) + data = { + "messages": [{"role": "user", "content": original_content}], + "model": "gpt-5-mini", + "metadata": {}, + } + + result = await lakera_guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + cache=DualCache(), + data=data, + call_type="completion", + ) + + assert "[MASKED" in result["messages"][0]["content"] + assert result["messages"][0]["content"] != original_content + assert len(result["messages"]) == 2, "the remaining, non-PII concern still gets an advisory note" + assert result["messages"][1]["role"] == "system" + + @pytest.mark.asyncio + async def test_pre_call_blocks_instead_of_advisory_when_pii_is_not_maskable(self): + """ + Bugbot finding on BerriAI/litellm#34940: a PII-only or mixed violation on + input that can't be safely masked (combined messages+input, multimodal + content) fell through to the advisory branch with raw, unredacted content. + It must degrade to blocking instead, same as block mode already does for + this exact case, rather than showing an advisory note next to raw content. + """ + lakera_guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="inject_system_message") + mock_response = { + "flagged": True, + "payload": [{"detector_type": "pii/email", "start": 11, "end": 26, "message_id": 0}], + "breakdown": [{"detector_type": "pii/email", "detected": True}], + } + + with patch.object(lakera_guardrail, "call_v2_guard", new_callable=AsyncMock) as mock_call: + mock_call.return_value = (mock_response, {}) + data = { + "messages": [{"role": "user", "content": "My email is test@example.com"}], + "input": "responses-api content", + "model": "gpt-5-mini", + "metadata": {}, + } + with pytest.raises(HTTPException): + await lakera_guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + cache=DualCache(), + data=data, + call_type="completion", + ) + + assert "messages" in data + assert data["messages"][0]["content"] == "My email is test@example.com", ( + "the raw content must be untouched, not partially rewritten before the block" + ) + @pytest.mark.asyncio async def test_moderation_hook_inspects_all_message_roles_not_just_user(self): """See test_pre_call_inspects_all_message_roles_not_just_user.""" @@ -1022,6 +1099,42 @@ class TestAdvisoryModeWiring: assert "[MASKED" in result["messages"][0]["content"] assert result["messages"][0]["content"] != original_content + @pytest.mark.asyncio + async def test_moderation_hook_mixed_violation_masks_pii_even_though_advisory_has_no_effect(self): + """ + Bugbot finding on BerriAI/litellm#34940: a mixed violation isn't PII-only, + so it fell through to the during_call no-op branch with the raw PII still + in place. It must mask the maskable PII even though the advisory note + itself still has no effect during during_call. + """ + lakera_guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="inject_system_message") + mock_response = { + "flagged": True, + "payload": [{"detector_type": "pii/email", "start": 11, "end": 26, "message_id": 0}], + "breakdown": [ + {"detector_type": "pii/email", "detected": True}, + {"detector_type": "prompt_injection", "detected": True}, + ], + } + original_content = "My email is test@example.com" + + with patch.object(lakera_guardrail, "call_v2_guard", new_callable=AsyncMock) as mock_call: + mock_call.return_value = (mock_response, {}) + data = { + "messages": [{"role": "user", "content": original_content}], + "model": "gpt-5-mini", + "metadata": {}, + } + result = await lakera_guardrail.async_moderation_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + call_type="completion", + ) + + assert "[MASKED" in result["messages"][0]["content"] + assert result["messages"][0]["content"] != original_content + assert len(result["messages"]) == 1, "no advisory note is appended during during_call" + class TestAdvisoryModePostCall: """