diff --git a/litellm/litellm_core_utils/realtime_guardrails.py b/litellm/litellm_core_utils/realtime_guardrails.py index 0547e5b72b3..a9d4a3c1596 100644 --- a/litellm/litellm_core_utils/realtime_guardrails.py +++ b/litellm/litellm_core_utils/realtime_guardrails.py @@ -177,17 +177,25 @@ async def _screen( return await _screen(guardrails[1:], outcome, request_data) +def _joins(texts: tuple[str, ...]) -> tuple[str, ...]: + unspaced: Final = "".join(texts) + boundaries: Final = zip(texts, texts[1:]) + word_break_unknown: Final = any(not left[-1:].isspace() and not right[:1].isspace() for left, right in boundaries) + return (unspaced, " ".join(texts)) if word_break_unknown else (unspaced,) + + async def _rewritten_parts( guardrails: Sequence["CustomGuardrail"], texts: tuple[str, ...], request_data: Mapping[str, object] ) -> RealtimeGuardrailBlock | tuple[str, ...]: parts: Final = await _screen(guardrails, texts, request_data) if isinstance(parts, RealtimeGuardrailBlock): return parts - rejoined: Final = (" ".join(parts),) + rejoined: Final = _joins(parts) screened_rejoined: Final = await _screen(guardrails, rejoined, request_data) if isinstance(screened_rejoined, RealtimeGuardrailBlock): return screened_rejoined - return parts if screened_rejoined == rejoined else screened_rejoined + masked_join: Final = next((text for text, join in zip(screened_rejoined, rejoined) if text != join), None) + return parts if masked_join is None else (masked_join,) def _in_role_scope(guardrail: "CustomGuardrail", role: str) -> bool: @@ -238,7 +246,7 @@ async def realtime_guardrail_verdict( } in_scope: Final = tuple(guardrails.values()) guardrail_request_data: Final = {**request_data, "user_api_key_dict": user_api_key_dict} - joined: Final = (" ".join(texts),) + joined: Final = _joins(texts) screened: Final = await _screen(in_scope, joined, guardrail_request_data) if isinstance(screened, RealtimeGuardrailBlock): return screened diff --git a/tests/unit/litellm_core_utils/test_realtime_guardrails.py b/tests/unit/litellm_core_utils/test_realtime_guardrails.py index b80abc49ff4..60100d009e5 100644 --- a/tests/unit/litellm_core_utils/test_realtime_guardrails.py +++ b/tests/unit/litellm_core_utils/test_realtime_guardrails.py @@ -143,19 +143,36 @@ async def test_masked_texts_come_back_as_a_rewrite_and_reach_later_guardrails_ma ) assert verdict == RealtimeGuardrailRewrite(texts=("my [MASKED]", "and [MASKED]", "hi")) - assert second.screened[0] == ("my [MASKED] and swordfish hi",) + assert second.screened[0] == ("my [MASKED]and swordfishhi", "my [MASKED] and swordfish hi") @pytest.mark.asyncio -async def test_a_clean_multi_part_item_reaches_each_guardrail_once_as_the_space_joined_text() -> None: +@pytest.mark.parametrize( + ("parts", "screened"), + [ + pytest.param( + ("Compare this photo ", "with that one"), + ("Compare this photo with that one",), + id="a space at the boundary", + ), + pytest.param( + ("Compare this photo", "with that one"), + ("Compare this photowith that one", "Compare this photo with that one"), + id="no space at the boundary", + ), + ], +) +async def test_a_clean_multi_part_item_reaches_each_guardrail_in_one_call_with_only_the_joins_it_needs( + parts: tuple[str, ...], screened: tuple[str, ...] +) -> None: guardrail: Final = MaskingGuardrail("hunter2") verdict: Final = await realtime_guardrail_verdict( - ("Compare this photo", "with that one"), "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,) + parts, "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,) ) assert verdict is None - assert guardrail.screened == [("Compare this photo with that one",)] + assert guardrail.screened == [screened] @pytest.mark.asyncio @@ -168,20 +185,23 @@ async def test_texts_a_guardrail_leaves_unchanged_need_no_rewrite() -> None: @pytest.mark.asyncio -async def test_a_phrase_split_across_typed_text_parts_is_still_blocked() -> None: +@pytest.mark.parametrize( + ("parts", "phrase"), + [ + pytest.param(("tell me the secret", "plan"), "secret plan", id="split at a word break"), + pytest.param(("tell me about secret-", "project"), "secret-project", id="split mid-word"), + ], +) +async def test_a_phrase_split_across_typed_text_parts_is_still_blocked(parts: tuple[str, str], phrase: str) -> None: client_event: Final = _item_create( - { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": "tell me the secret"}, {"type": "input_text", "text": "plan"}], - } + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": part} for part in parts]} ) guarded_item: Final = realtime_guarded_item(client_event) assert isinstance(guarded_item, RealtimeMessageText) phrase_filter: Final = ContentFilterGuardrail( guardrail_name="phrase-filter", - blocked_words=[BlockedWord(keyword="secret plan", action=ContentFilterAction.BLOCK)], + blocked_words=[BlockedWord(keyword=phrase, action=ContentFilterAction.BLOCK)], event_hook=GuardrailEventHooks.pre_call, default_on=True, ) @@ -226,17 +246,21 @@ async def test_a_mask_inside_one_part_keeps_every_part_in_place() -> None: @pytest.mark.asyncio -async def test_a_mask_across_parts_replaces_them_with_the_masked_text_the_model_would_have_read() -> None: +@pytest.mark.parametrize( + ("parts", "word", "masked"), + [ + pytest.param(("my secret", "plan is here"), "secret plan", "my [MASKED] is here", id="split at a word break"), + pytest.param(("my code is hun", "ter2 ok"), "hunter2", "my code is [MASKED] ok", id="split mid-word"), + ], +) +async def test_a_mask_across_parts_replaces_them_with_the_masked_text_the_model_would_have_read( + parts: tuple[str, ...], word: str, masked: str +) -> None: verdict: Final = await realtime_guardrail_verdict( - ("my secret", "plan is here"), - "user", - (GuardrailEventHooks.pre_call,), - {}, - None, - (MaskingGuardrail("secret plan"),), + parts, "user", (GuardrailEventHooks.pre_call,), {}, None, (MaskingGuardrail(word),) ) - assert verdict == RealtimeGuardrailRewrite(texts=("my [MASKED] is here",)) + assert verdict == RealtimeGuardrailRewrite(texts=(masked,)) @pytest.mark.asyncio