From 4031160928b9e043b88c232fed7b0bee3f2422fe Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 08:24:51 -0800 Subject: [PATCH] fix(realtime): keep every per-part mask when a joined text masks differently --- .../litellm_core_utils/realtime_guardrails.py | 34 +++++++++++---- .../test_realtime_guardrails.py | 42 +++++++++++++++++-- 2 files changed, 65 insertions(+), 11 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_guardrails.py b/litellm/litellm_core_utils/realtime_guardrails.py index d4c7c87f307..161b266cd3a 100644 --- a/litellm/litellm_core_utils/realtime_guardrails.py +++ b/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) diff --git a/tests/unit/litellm_core_utils/test_realtime_guardrails.py b/tests/unit/litellm_core_utils/test_realtime_guardrails.py index ad3e9fabb30..72c43c8e56c 100644 --- a/tests/unit/litellm_core_utils/test_realtime_guardrails.py +++ b/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