From 6219636835c3891cc325cb1ee37b19bec13c0fdc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 18:29:03 -0700 Subject: [PATCH] fix(guardrails): mask findings on both sides of a JSON escape and skip an empty Lakera post-call A systemone entry whose key or value has a JSON escape is scanned as the entry text and decoded. Each field now takes the decoded rewrite when it changed the field, else the entry's, so a card in the key and a tab-split card in the value both mask, and an email after a line break masks instead of answering 500. Lakera v2 in post_call mode returns before calling its API when there are no messages to scan, which is every extraction and systemone response. --- .../guardrail_translation/handler.py | 21 +++++----- .../guardrail_hooks/lakera_ai_v2.py | 2 + .../test_decisions_guardrail_handler.py | 39 +++++++++++++++---- .../guardrail_hooks/test_lakera_ai_v2.py | 20 ++++++++++ 4 files changed, 64 insertions(+), 18 deletions(-) diff --git a/litellm/llms/openai/decisions/guardrail_translation/handler.py b/litellm/llms/openai/decisions/guardrail_translation/handler.py index 5b04bde9572..ead6f3e176c 100644 --- a/litellm/llms/openai/decisions/guardrail_translation/handler.py +++ b/litellm/llms/openai/decisions/guardrail_translation/handler.py @@ -3,7 +3,7 @@ import re from collections.abc import Iterator, Mapping from itertools import chain, groupby from types import MappingProxyType -from typing import TYPE_CHECKING, ClassVar, Final +from typing import TYPE_CHECKING, ClassVar, Final, TypeVar from pydantic import JsonValue, TypeAdapter, ValidationError @@ -22,6 +22,7 @@ _JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _KEY_ADAPTER: Final[TypeAdapter[str]] = TypeAdapter(str) _JSON_KEY_PREFIX: Final = re.compile(r'("(?:[^"\\\x00-\x1f]|\\(?:["\\/bfnrt]|u[0-9a-fA-F]{4}))*"): ') _DESCRIBED_ITEM_KEYS: Final = ("choices", "levels") +_Field: Final = TypeVar("_Field", bound=JsonValue) _QUESTION_LIST_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, JsonValue], ...]]] = TypeAdapter( tuple[Mapping[str, JsonValue], ...] ) @@ -173,6 +174,10 @@ def _decoded_text(text: str, form: str, guardrailed: Mapping[str, str]) -> str: return form if form in text else guardrailed.get(form, form) +def _merged_field(original: _Field, rewritten: _Field, decoded: _Field) -> _Field: + return rewritten if decoded == original else decoded + + def _guardrailed_entry( key: str, item: JsonValue, guardrailed: Mapping[str, str], guardrail_name: str | None ) -> tuple[str, JsonValue]: @@ -180,16 +185,12 @@ def _guardrailed_entry( return guardrailed.get(key, key), _guardrailed_json(item, guardrailed, guardrail_name) text: Final = _entry_text(key, item) rewritten: Final = guardrailed.get(text, text) - decoded: Final = ( - _decoded_text(text, key, guardrailed), - _decoded_text(text, item, guardrailed) if isinstance(item, str) else item, + entry: Final = (key, item) if rewritten == text else _rewritten_entry(rewritten, guardrail_name) + decoded_item: Final = _decoded_text(text, item, guardrailed) if isinstance(item, str) else item + return ( + _merged_field(key, entry[0], _decoded_text(text, key, guardrailed)), + _merged_field(item, entry[1], decoded_item), ) - if rewritten == text: - return decoded - entry: Final = _rewritten_entry(rewritten, guardrail_name) - if decoded not in ((key, item), entry): - raise unappliable_request_rewrite(guardrail_name) - return entry def _instruction_texts(question: Mapping[str, JsonValue]) -> tuple[str, ...]: diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 7c0cdb0d36f..52ae718572c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -783,6 +783,8 @@ class LakeraAIGuardrail(CustomGuardrail): # Use a copy of original_messages so _mask_pii_in_messages does not mutate data["messages"] post_call_messages: Final = list(copy.deepcopy(original_messages)) + response_messages + if not post_call_messages: + return response # Call Lakera guardrail lakera_guardrail_response, _ = await self.call_v2_guard( diff --git a/tests/unit/llms/openai/decisions/guardrail_translation/test_decisions_guardrail_handler.py b/tests/unit/llms/openai/decisions/guardrail_translation/test_decisions_guardrail_handler.py index 080f231f5e7..a7e5fa3711f 100644 --- a/tests/unit/llms/openai/decisions/guardrail_translation/test_decisions_guardrail_handler.py +++ b/tests/unit/llms/openai/decisions/guardrail_translation/test_decisions_guardrail_handler.py @@ -2,6 +2,7 @@ from functools import reduce from typing import Final, Literal import pytest +from pydantic import JsonValue from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.prompt_templates.common_utils import request_contains_image_content @@ -232,15 +233,37 @@ async def test_systemone_masks_a_card_number_that_json_escaping_splits() -> None @pytest.mark.asyncio -async def test_systemone_rejects_an_escaped_entry_whose_two_scanned_forms_disagree() -> None: - data: Final = { - "model": "judge", - "state": {"note": "first\nsecond"}, - "questions": {"q": {"type": "noul", "criteria": {}}}, - } +@pytest.mark.parametrize( + ("pattern_name", "state", "masked"), + ( + pytest.param( + "visa", + {"4111111111111111": "4111\t1111\t1111\t1111"}, + {"[VISA_REDACTED]": "[VISA_REDACTED]"}, + id="key-and-escaped-value", + ), + pytest.param( + "email", + {"note": "Please call me back.\ndana.quill@example.com"}, + {"note": "Please call me back.\n[EMAIL_REDACTED]"}, + id="match-after-an-escape", + ), + ), +) +async def test_systemone_masks_findings_on_both_sides_of_a_json_escape( + pattern_name: str, state: dict[str, JsonValue], masked: dict[str, JsonValue] +) -> None: + guardrail: Final = ContentFilterGuardrail( + guardrail_name="mask", + patterns=[ + ContentFilterPattern(pattern_type="prebuilt", pattern_name=pattern_name, action=ContentFilterAction.MASK) + ], + ) + data: Final = {"model": "judge", "state": state, "questions": {"q": {"type": "noul", "criteria": {}}}} - with pytest.raises(UnappliableRequestRewrite): - await SystemOneHandler().process_input_messages(data, _TaggingGuardrail()) + result: Final = await SystemOneHandler().process_input_messages(data, guardrail) + + assert result["state"] == masked @pytest.mark.asyncio diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py index 23419ebc839..2dcbf32eed3 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py @@ -6,6 +6,7 @@ Additional tests live in tests/guardrails_tests/test_lakera_v2.py. """ import importlib, logging, os +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -317,6 +318,25 @@ class TestAsyncPostCallSuccessHookSkipFlags: assert "a@b.com" not in result_dict["choices"][0]["message"]["content"] +@pytest.mark.asyncio +async def test_post_call_skips_a_response_with_no_assistant_text() -> None: + guardrail: Final = LakeraAIGuardrail(api_key="test_key") + data: Final = {"model": "clairic", "metadata": {}} + llm_response: Final = MagicMock() + llm_response.model_dump.return_value = {"answer": "Berlin"} + + with patch.object(guardrail, "call_v2_guard", new_callable=AsyncMock) as mock_call: + result: Final = await guardrail.async_post_call_success_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + response=llm_response, + ) + + mock_call.assert_not_called() + assert result is llm_response + assert "applied_guardrails" not in data["metadata"] + + PII_ONLY_LAKERA_RESPONSE = { "flagged": True, "breakdown": [{"detector_type": "pii/email", "detected": True, "message_id": 0}],