diff --git a/litellm/litellm_core_utils/realtime_guardrails.py b/litellm/litellm_core_utils/realtime_guardrails.py index 161b266cd3a..0547e5b72b3 100644 --- a/litellm/litellm_core_utils/realtime_guardrails.py +++ b/litellm/litellm_core_utils/realtime_guardrails.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Final, TypeAlias from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_logger +from litellm.exceptions import ModifyResponseException if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail @@ -132,7 +133,7 @@ def _violation_text(error: Exception) -> str: def _is_guardrail_block(error: Exception) -> bool: - return hasattr(error, "status_code") or isinstance(error, ValueError) + return hasattr(error, "status_code") or isinstance(error, (ValueError, ModifyResponseException)) def _screened_texts(screened: object, texts: tuple[str, ...]) -> tuple[str, ...]: @@ -145,12 +146,12 @@ def _screened_texts(screened: object, texts: tuple[str, ...]) -> tuple[str, ...] async def _guardrail_outcome( - guardrail: "CustomGuardrail", texts: tuple[str, ...], user_api_key_dict: object | None + guardrail: "CustomGuardrail", texts: tuple[str, ...], request_data: Mapping[str, object] ) -> RealtimeGuardrailBlock | tuple[str, ...]: try: screened: Final = await guardrail.apply_guardrail( inputs={"texts": list(texts), "images": []}, - request_data={"user_api_key_dict": user_api_key_dict}, + request_data=dict(request_data), input_type="request", ) except Exception as error: @@ -166,42 +167,27 @@ async def _guardrail_outcome( async def _screen( - guardrails: Sequence["CustomGuardrail"], texts: tuple[str, ...], user_api_key_dict: object | None + guardrails: Sequence["CustomGuardrail"], texts: tuple[str, ...], request_data: Mapping[str, object] ) -> RealtimeGuardrailBlock | tuple[str, ...]: if not guardrails: return texts - outcome: Final = await _guardrail_outcome(guardrails[0], texts, user_api_key_dict) + outcome: Final = await _guardrail_outcome(guardrails[0], texts, request_data) if isinstance(outcome, RealtimeGuardrailBlock): return outcome - return await _screen(guardrails[1:], outcome, user_api_key_dict) + return await _screen(guardrails[1:], outcome, request_data) -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, *_joins(texts)) - - -async def _rewritten_texts( - guardrails: Sequence["CustomGuardrail"], - texts: tuple[str, ...], - screened: tuple[str, ...], - user_api_key_dict: object | None, +async def _rewritten_parts( + guardrails: Sequence["CustomGuardrail"], texts: tuple[str, ...], request_data: Mapping[str, object] ) -> RealtimeGuardrailBlock | tuple[str, ...]: - parts: Final = screened[: len(texts)] - if len(texts) < 2: + parts: Final = await _screen(guardrails, texts, request_data) + if isinstance(parts, RealtimeGuardrailBlock): return parts - 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],) + rejoined: Final = (" ".join(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 def _in_role_scope(guardrail: "CustomGuardrail", role: str) -> bool: @@ -251,10 +237,14 @@ async def realtime_guardrail_verdict( if guardrail is not None and _in_role_scope(guardrail, role) } in_scope: Final = tuple(guardrails.values()) - screened: Final = await _screen(in_scope, _screening_texts(texts), user_api_key_dict) + guardrail_request_data: Final = {**request_data, "user_api_key_dict": user_api_key_dict} + joined: Final = (" ".join(texts),) + screened: Final = await _screen(in_scope, joined, guardrail_request_data) if isinstance(screened, RealtimeGuardrailBlock): return screened - rewritten: Final = await _rewritten_texts(in_scope, texts, screened, user_api_key_dict) + if screened == joined: + return None + rewritten: Final = screened if len(texts) < 2 else await _rewritten_parts(in_scope, texts, guardrail_request_data) if isinstance(rewritten, RealtimeGuardrailBlock): return rewritten return None if rewritten == texts else RealtimeGuardrailRewrite(texts=rewritten) diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 8aae82abc27..7ee3a7cb248 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -452,7 +452,7 @@ class BedrockRealtimeGuardrails: guarded_item: Final = realtime_guarded_item(client_event) if guarded_item is None: return _ClientEventRoute.FORWARD - role: Final = "tool" if isinstance(guarded_item, RealtimeToolOutput) else guarded_item.role + role: Final = "tool" if isinstance(guarded_item, RealtimeToolOutput) else "user" match await self._verdict(guarded_item.texts, role, GuardrailEventHooks.pre_call): case None: return _ClientEventRoute.FORWARD diff --git a/tests/unit/litellm_core_utils/test_realtime_guardrails.py b/tests/unit/litellm_core_utils/test_realtime_guardrails.py index 72c43c8e56c..b80abc49ff4 100644 --- a/tests/unit/litellm_core_utils/test_realtime_guardrails.py +++ b/tests/unit/litellm_core_utils/test_realtime_guardrails.py @@ -7,6 +7,7 @@ from fastapi import HTTPException from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail +from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import CustomCodeGuardrail from litellm.litellm_core_utils.realtime_guardrails import ( RealtimeGuardrailBlock, RealtimeGuardrailRewrite, @@ -142,7 +143,19 @@ 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][:3] == ("my [MASKED]", "and swordfish", "hi") + assert second.screened[0] == ("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: + guardrail: Final = MaskingGuardrail("hunter2") + + verdict: Final = await realtime_guardrail_verdict( + ("Compare this photo", "with that one"), "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,) + ) + + assert verdict is None + assert guardrail.screened == [("Compare this photo with that one",)] @pytest.mark.asyncio @@ -155,23 +168,20 @@ async def test_texts_a_guardrail_leaves_unchanged_need_no_rewrite() -> None: @pytest.mark.asyncio -@pytest.mark.parametrize( - ("keyword", "parts"), - [ - pytest.param("secret plan", ("tell me the secret", "plan"), id="split at a space"), - pytest.param("secret-project", ("Tell me about secret-", "project"), id="split mid-word"), - ], -) -async def test_a_phrase_split_across_typed_text_parts_is_still_blocked(keyword: str, parts: tuple[str, ...]) -> None: +async def test_a_phrase_split_across_typed_text_parts_is_still_blocked() -> None: client_event: Final = _item_create( - {"type": "message", "role": "user", "content": [{"type": "input_text", "text": part} for part in parts]} + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "tell me the secret"}, {"type": "input_text", "text": "plan"}], + } ) 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=keyword, action=ContentFilterAction.BLOCK)], + blocked_words=[BlockedWord(keyword="secret plan", action=ContentFilterAction.BLOCK)], event_hook=GuardrailEventHooks.pre_call, default_on=True, ) @@ -218,7 +228,12 @@ 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: verdict: Final = await realtime_guardrail_verdict( - ("my hunt", "er2 is here"), "user", (GuardrailEventHooks.pre_call,), {}, None, (MaskingGuardrail("hunter2"),) + ("my secret", "plan is here"), + "user", + (GuardrailEventHooks.pre_call,), + {}, + None, + (MaskingGuardrail("secret plan"),), ) assert verdict == RealtimeGuardrailRewrite(texts=("my [MASKED] is here",)) @@ -228,8 +243,8 @@ async def test_a_mask_across_parts_replaces_them_with_the_masked_text_the_model_ @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"), + pytest.param(("abc", "123-45-6789 hunter2 rocks"), id="masks inside one part"), + pytest.param(("abc", "123-45-6789 hunter2", "rocks"), 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( @@ -240,7 +255,7 @@ async def test_no_mask_one_text_part_gets_is_lost_when_its_joined_text_masks_dif patterns=[ ContentFilterPattern(pattern_type="prebuilt", pattern_name="us_ssn", action=ContentFilterAction.MASK) ], - blocked_words=[BlockedWord(keyword="hunter2", action=ContentFilterAction.MASK)], + blocked_words=[BlockedWord(keyword="hunter2 rocks", action=ContentFilterAction.MASK)], event_hook=GuardrailEventHooks.pre_call, default_on=True, ) @@ -250,9 +265,9 @@ async def test_no_mask_one_text_part_gets_is_lost_when_its_joined_text_masks_dif ) assert isinstance(verdict, RealtimeGuardrailRewrite) - forwarded: Final = "".join(verdict.texts) + forwarded: Final = " ".join(verdict.texts) assert "123-45-6789" not in forwarded - assert "hunter2" not in forwarded + assert "hunter2 rocks" not in forwarded assert forwarded.startswith("abc") @@ -328,3 +343,33 @@ async def test_a_guardrail_that_crashes_is_not_treated_as_a_block() -> None: with pytest.raises(RuntimeError, match="guardrail backend down"): await realtime_guardrail_verdict(("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,)) + + +_TEAM_POLICY_CODE: Final = ( + "def apply_guardrail(inputs, request_data, input_type):\n" + ' if request_data["team_id"] == "team-strict":\n' + ' return block("blocked for this team")\n' + " return allow()\n" +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_id", "blocked"), + [pytest.param("team-strict", True, id="policy team"), pytest.param("team-other", False, id="other team")], +) +async def test_a_team_scoped_guardrail_sees_the_session_team(team_id: str, blocked: bool) -> None: + team_policy = CustomCodeGuardrail( + custom_code=_TEAM_POLICY_CODE, guardrail_name="team-policy", event_hook="pre_call", default_on=True + ) + + verdict = await realtime_guardrail_verdict( + ("hello",), + "user", + (GuardrailEventHooks.pre_call,), + {"litellm_metadata": {"user_api_key_team_id": team_id}}, + None, + (team_policy,), + ) + + assert isinstance(verdict, RealtimeGuardrailBlock) is blocked diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 81c7d80fa54..97d18a97182 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -1937,7 +1937,9 @@ class TestBedrockRealtimeGuardrails: assert [error["message"] for error in self._guardrail_errors(client_ws)] == [self.VIOLATION] @pytest.mark.asyncio - async def test_a_system_message_the_guardrail_skips_reaches_bedrock(self, stub_aws_models: None) -> None: + async def test_a_system_message_nova_receives_as_a_user_turn_is_screened_by_a_system_skipping_guardrail( + self, stub_aws_models: None + ) -> None: guardrail: Final = SystemSkippingGuardrail( "hunter2", guardrail_name="no_secrets", @@ -1946,13 +1948,10 @@ class TestBedrockRealtimeGuardrails: realtime_violation_message=self.VIOLATION, ) - stream, client_ws = await self._forward_client( - [self._user_text(self.BLOCKED_TEXT, role="system"), self._user_text(self.BLOCKED_TEXT)], guardrail - ) + stream, client_ws = await self._forward_client([self._user_text(self.BLOCKED_TEXT, role="system")], guardrail) - assert self.BLOCKED_TEXT in json.dumps(self._bedrock_events(stream)) + assert "hunter2" not in json.dumps(self._bedrock_events(stream)) assert [error["message"] for error in self._guardrail_errors(client_ws)] == [self.VIOLATION] - assert guardrail.checked == [((self.BLOCKED_TEXT,), None)] @pytest.mark.asyncio async def test_the_next_clean_turn_is_answered_after_a_blocked_item_with_no_response_create(