scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-17 23:51:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 67

fix(guardrails): block during_call when PII can't be safely masked

Browse source 
Greptile finding (P1, security): async_moderation_hook's inject_system_message
branch had no equivalent to async_pre_call_hook's degrade-to-blocking case for
a PII violation on input that can't be safely masked (e.g. combined
messages+input). It fell through to the advisory no-op branch and let raw,
unredacted PII reach the model with no protection at all. Raising still
blocks the response from reaching the caller even though during_call races
with the LLM dispatch, the same mechanism on_flagged="block" already relies
on for this hook, so add the same block-instead-of-advisory branch pre_call
already has.
This commit is contained in:
Deepanshu 2026-08-27 18:20:10 -04:00
parent 41f03b9c19
commit 50c7f107e1
2 changed files with 49 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -668,6 +668,15 @@ 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 _breakdown_has_pii_violation(lakera_guardrail_response) and is_multimodal_input:
# Same as async_pre_call_hook: there's PII in the mix and
# nothing here can be safely masked, so allowing it through
# unprotected would be worse than blocking. Unlike mutating
# data["messages"] below, raising still blocks the response
# from reaching the caller even though during_call races with
# the LLM dispatch -- same mechanism on_flagged="block" already
# relies on for this hook.
raise self._get_http_exception_for_blocked_guardrail(lakera_guardrail_response)
if _breakdown_has_pii_violation(lakera_guardrail_response) and not is_multimodal_input:
# A mixed violation (PII plus something else): mask whatever's
# maskable even though the advisory note below has no effect

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

@ -1370,6 +1370,46 @@ class TestAdvisoryModeWiring:
assert result["messages"][0]["content"] != original_content
assert len(result["messages"]) == 1, "no advisory note is appended during during_call"
@pytest.mark.asyncio
async def test_moderation_hook_blocks_instead_of_advisory_when_pii_is_not_maskable(self):
"""
Greptile finding (P1, security) on BerriAI/litellm#34940: a PII violation
on input that can't be safely masked (combined messages+input) fell
through to the during_call no-op branch and let raw, unredacted PII reach
the model with no protection at all. async_pre_call_hook already degrades
to blocking for this exact case (see
test_pre_call_blocks_instead_of_advisory_when_pii_is_not_maskable) --
async_moderation_hook must too, since raising here still blocks the
response from reaching the caller (same mechanism on_flagged="block"
already relies on), unlike mutating data["messages"] which races with
the concurrent LLM dispatch.
"""
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_moderation_hook(
data=data,
user_api_key_dict=UserAPIKeyAuth(api_key="test_key"),
call_type="completion",
)
assert data["messages"][0]["content"] == "My email is test@example.com", (
"the raw content must be untouched, not partially rewritten before the block"
)
class TestAdvisoryModePostCall:
"""