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 134

fix(guardrails): require payload and breakdown for Lakera v2 advisory mode

Browse source 
Advisory mode's mixed-violation masking safety net can only redact
detected PII when Lakera's response carries both the breakdown (to
detect a PII hit at all) and payload (the location data to mask by).
payload=False or breakdown=False alongside on_flagged='inject_system_message'
silently forwarded raw PII next to the advisory note. Reject that
combination at construction and hot-reload time instead.
This commit is contained in:
Deepanshu 2026-08-26 10:08:37 -04:00
parent 76e7ff9459
commit 3b67436cf9
2 changed files with 87 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -293,6 +293,8 @@ class LakeraAIGuardrail(CustomGuardrail):
on_flagged=self.on_flagged,
advisory_system_message=self.advisory_system_message,
event_hook=self.event_hook,
payload=self.payload,
breakdown=self.breakdown,
)
def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
@ -315,10 +317,14 @@ class LakeraAIGuardrail(CustomGuardrail):
whose real dispatch hook never changed.
"""
new_event_hook: Final = getattr(litellm_params, "mode", None) or self.event_hook
prospective_payload: Final = getattr(litellm_params, "payload", None)
prospective_breakdown: Final = getattr(litellm_params, "breakdown", None)
self._validate_advisory_config(
on_flagged=getattr(litellm_params, "on_flagged", None) or self.on_flagged,
advisory_system_message=getattr(litellm_params, "advisory_system_message", None),
event_hook=new_event_hook,
payload=self.payload if prospective_payload is None else prospective_payload,
breakdown=self.breakdown if prospective_breakdown is None else prospective_breakdown,
)
super().update_in_memory_litellm_params(litellm_params=litellm_params)
self.event_hook = new_event_hook
@ -328,6 +334,8 @@ class LakeraAIGuardrail(CustomGuardrail):
on_flagged: str,
advisory_system_message: str | None,
event_hook: GuardrailEventHooks | Sequence[GuardrailEventHooks] | Mode | str | Sequence[str] | None,
payload: bool | None,
breakdown: bool | None,
) -> None:
if advisory_system_message is not None:
if not _template_uses_reason_placeholder(advisory_system_message):
@ -348,6 +356,13 @@ class LakeraAIGuardrail(CustomGuardrail):
"runs concurrently with the LLM dispatch with no pre-call barrier, so the advisory message "
"cannot reliably reach the request. Use mode='pre_call' instead."
)
if on_flagged == "inject_system_message" and not (payload and breakdown):
raise ValueError(
"on_flagged='inject_system_message' requires payload=True and breakdown=True: advisory "
"mode masks any detected PII before appending the advisory note, and that masking can "
"only happen when Lakera's response carries both the violation breakdown and the "
"payload location data. Without them, PII would be forwarded to the model unredacted."
)
def _build_advisory_message(self, lakera_response: LakeraAIResponse | None) -> str:
"""Format the advisory message shown to the LLM when on_flagged='inject_system_message'."""

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

@ -735,6 +735,78 @@ class TestAdvisoryModeDuringCallUnsupported:
assert guardrail.event_hook == "pre_call"
class TestAdvisoryModeRequiresPayloadAndBreakdown:
"""Veria-ai finding on BerriAI/litellm#34940: the mixed-violation masking
safety net (mask any detected PII before appending the advisory note) only
works when Lakera's response carries both breakdown (to detect a PII hit
at all) and payload (the location data to mask by). payload=False or
breakdown=False alongside on_flagged='inject_system_message' would forward
raw, unredacted PII next to the advisory note with no error and no signal
to the operator, so that combination must be rejected at construction
time, same as the during_call combination already is."""
def test_payload_false_raises_at_construction(self):
with pytest.raises(ValueError, match="requires payload=True and breakdown=True"):
LakeraAIGuardrail(api_key="test_key", on_flagged="inject_system_message", payload=False)
def test_breakdown_false_raises_at_construction(self):
with pytest.raises(ValueError, match="requires payload=True and breakdown=True"):
LakeraAIGuardrail(api_key="test_key", on_flagged="inject_system_message", breakdown=False)
def test_both_false_raises_at_construction(self):
with pytest.raises(ValueError, match="requires payload=True and breakdown=True"):
LakeraAIGuardrail(
api_key="test_key", on_flagged="inject_system_message", payload=False, breakdown=False
)
def test_defaults_construct_without_error(self):
guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="inject_system_message")
assert guardrail.payload is True
assert guardrail.breakdown is True
def test_payload_false_with_block_mode_constructs_without_error(self):
guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="block", payload=False)
assert guardrail.payload is False
def test_in_memory_update_reintroducing_payload_false_raises(self):
guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="block", payload=False)
updated_params = LitellmParams(
guardrail="lakera_v2", mode="pre_call", on_flagged="inject_system_message", payload=False
)
with pytest.raises(ValueError, match="requires payload=True and breakdown=True"):
guardrail.update_in_memory_litellm_params(litellm_params=updated_params)
assert guardrail.on_flagged == "block", "a rejected update must leave the live instance untouched"
def test_in_memory_update_leaving_payload_unspecified_resets_to_the_model_default(self):
"""LitellmParams.payload defaults to True (not None/unset), so an update
that doesn't mention payload at all still carries payload=True through
the base setattr -- it does not preserve the live instance's prior
False value. That's a valid transition, not a bug: it's the same
pydantic-default behavior every other field on this update already has."""
guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="block", payload=False)
updated_params = LitellmParams(guardrail="lakera_v2", mode="pre_call", on_flagged="inject_system_message")
guardrail.update_in_memory_litellm_params(litellm_params=updated_params)
assert guardrail.on_flagged == "inject_system_message"
assert guardrail.payload is True
def test_in_memory_update_disabling_breakdown_on_an_advisory_instance_raises(self):
guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="inject_system_message")
updated_params = LitellmParams(
guardrail="lakera_v2", mode="pre_call", on_flagged="inject_system_message", breakdown=False
)
with pytest.raises(ValueError, match="requires payload=True and breakdown=True"):
guardrail.update_in_memory_litellm_params(litellm_params=updated_params)
assert guardrail.breakdown is True, "a rejected update must leave the live instance untouched"
def test_in_memory_update_enabling_both_while_flipping_on_flagged_is_allowed(self):
guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="block", payload=False, breakdown=False)
updated_params = LitellmParams(
guardrail="lakera_v2", mode="pre_call", on_flagged="inject_system_message", payload=True, breakdown=True
)
guardrail.update_in_memory_litellm_params(litellm_params=updated_params)
assert guardrail.on_flagged == "inject_system_message"
class TestAdvisoryModeWiring:
"""Tests for on_flagged='inject_system_message' wiring in async_pre_call_hook / async_moderation_hook."""