scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-20 00:11:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 76

fix(guardrails): mask maskable PII in mixed violations before advising in Lakera v2

Browse source 
on_flagged="inject_system_message" only masked when a violation was
PII-only; a mixed violation (PII plus a non-PII flag like prompt injection)
fell straight through to the advisory branch with the raw PII still in
place, in both async_pre_call_hook and async_moderation_hook. Mask whatever
Lakera returned location data for before appending or logging the advisory,
so a mixed violation never ships raw PII just because something else was
also flagged.

Also degrade to blocking, same as block mode already does, when nothing
can be safely masked at all (multimodal content, or messages combined with
a Responses API input field) instead of showing an advisory note next to
raw, unredacted content.

Widened call_v2_guard/_mask_pii_in_messages/the write-back helper's message
parameters from list to Sequence to match what's actually passed through
from _filter_skipped_messages, instead of duplicating list(...) casts at
every call site.
This commit is contained in:
Deepanshu 2026-08-25 21:46:38 -04:00
parent 0716c55bfe
commit 4c7aac48a4
2 changed files with 147 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

39
litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py
View file

@ -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 <TYPE>]” 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

113
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py
View file

@ -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:
"""