fix(realtime): keep every per-part mask when a joined text masks differently

litellm/litellm_core_utils/realtime_guardrails.py

@@ -176,17 +176,32 @@ async def _screen(
     return await _screen(guardrails[1:], outcome, user_api_key_dict)
 
 
+def _joins(texts: tuple[str, ...]) -> tuple[str, str]:
+    return "".join(texts), " ".join(texts)
+
+
 def _screening_texts(texts: tuple[str, ...]) -> tuple[str, ...]:
-    return texts if len(texts) < 2 else (*texts, "".join(texts), " ".join(texts))
+    return texts if len(texts) < 2 else (*texts, *_joins(texts))
 
 
-def _rewritten_texts(texts: tuple[str, ...], screened: tuple[str, ...]) -> tuple[str, ...]:
+async def _rewritten_texts(
+    guardrails: Sequence["CustomGuardrail"],
+    texts: tuple[str, ...],
+    screened: tuple[str, ...],
+    user_api_key_dict: object | None,
+) -> RealtimeGuardrailBlock | tuple[str, ...]:
     parts: Final = screened[: len(texts)]
-    joins: Final = screened[len(texts) :]
-    if not joins or joins == ("".join(parts), " ".join(parts)):
+    if len(texts) < 2:
         return parts
-    unspaced: Final = joins[0]
-    return (unspaced,) if unspaced != "".join(texts) else (joins[1],)
+    joins: Final = _joins(parts)
+    screened_joins: Final = (
+        screened[len(texts) :] if parts == texts else await _screen(guardrails, joins, user_api_key_dict)
+    )
+    if isinstance(screened_joins, RealtimeGuardrailBlock):
+        return screened_joins
+    if screened_joins == joins:
+        return parts
+    return (screened_joins[0],) if screened_joins[0] != joins[0] else (screened_joins[1],)
 
 
 def _in_role_scope(guardrail: "CustomGuardrail", role: str) -> bool:
@@ -235,10 +250,13 @@ async def realtime_guardrail_verdict(
         for guardrail in runner_guardrails
         if guardrail is not None and _in_role_scope(guardrail, role)
     }
-    screened: Final = await _screen(tuple(guardrails.values()), _screening_texts(texts), user_api_key_dict)
+    in_scope: Final = tuple(guardrails.values())
+    screened: Final = await _screen(in_scope, _screening_texts(texts), user_api_key_dict)
     if isinstance(screened, RealtimeGuardrailBlock):
         return screened
-    rewritten: Final = _rewritten_texts(texts, screened)
+    rewritten: Final = await _rewritten_texts(in_scope, texts, screened, user_api_key_dict)
+    if isinstance(rewritten, RealtimeGuardrailBlock):
+        return rewritten
     return None if rewritten == texts else RealtimeGuardrailRewrite(texts=rewritten)
 
 

tests/unit/litellm_core_utils/test_realtime_guardrails.py

@@ -15,7 +15,7 @@ from litellm.litellm_core_utils.realtime_guardrails import (
     realtime_guardrail_verdict,
     rewritten_item,
 )
-from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks
+from litellm.types.guardrails import BlockedWord, ContentFilterAction, ContentFilterPattern, GuardrailEventHooks
 from litellm.types.utils import GenericGuardrailAPIInputs
 
 
@@ -104,7 +104,9 @@ def test_items_without_client_text_are_not_screened(client_event: dict) -> None:
 async def test_a_plain_string_block_reason_becomes_the_violation_message() -> None:
     guardrail: Final = RaisingGuardrail(HTTPException(status_code=400, detail="hello is not allowed"))
 
-    block: Final = await realtime_guardrail_verdict(("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,))
+    block: Final = await realtime_guardrail_verdict(
+        ("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,)
+    )
 
     assert isinstance(block, RealtimeGuardrailBlock)
     assert block.message == "hello is not allowed"
@@ -117,7 +119,9 @@ async def test_a_configured_violation_message_replaces_the_block_reason() -> Non
         HTTPException(status_code=400, detail="hello is not allowed"), realtime_violation_message="Please rephrase"
     )
 
-    block: Final = await realtime_guardrail_verdict(("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,))
+    block: Final = await realtime_guardrail_verdict(
+        ("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,)
+    )
     configured_block: Final = await realtime_guardrail_verdict(
         ("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (configured,)
     )
@@ -220,6 +224,38 @@ async def test_a_mask_across_parts_replaces_them_with_the_masked_text_the_model_
     assert verdict == RealtimeGuardrailRewrite(texts=("my [MASKED] is here",))
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "parts",
+    [
+        pytest.param(("abc", "123-45-6789 hunter2"), id="masks inside one part"),
+        pytest.param(("abc", "123-45-6789 hunt", "er2"), id="a mask inside a part and one across parts"),
+    ],
+)
+async def test_no_mask_one_text_part_gets_is_lost_when_its_joined_text_masks_differently(
+    parts: tuple[str, ...],
+) -> None:
+    content_filter: Final = ContentFilterGuardrail(
+        guardrail_name="ssn-and-keyword-mask",
+        patterns=[
+            ContentFilterPattern(pattern_type="prebuilt", pattern_name="us_ssn", action=ContentFilterAction.MASK)
+        ],
+        blocked_words=[BlockedWord(keyword="hunter2", action=ContentFilterAction.MASK)],
+        event_hook=GuardrailEventHooks.pre_call,
+        default_on=True,
+    )
+
+    verdict: Final = await realtime_guardrail_verdict(
+        parts, "user", (GuardrailEventHooks.pre_call,), {}, None, (content_filter,)
+    )
+
+    assert isinstance(verdict, RealtimeGuardrailRewrite)
+    forwarded: Final = "".join(verdict.texts)
+    assert "123-45-6789" not in forwarded
+    assert "hunter2" not in forwarded
+    assert forwarded.startswith("abc")
+
+
 class ScopedGuardrail(RaisingGuardrail):
     def __init__(
         self, *, skip_system_message: bool | None, skip_tool_message: bool | None, scan_only_tool_results: bool | None