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.

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, ...]:

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(

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

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}],