diff --git a/litellm/proxy/guardrails/guardrail_hooks/realmlabs/realmlabs.py b/litellm/proxy/guardrails/guardrail_hooks/realmlabs/realmlabs.py index 32aa9cedde2..98d1282b235 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/realmlabs/realmlabs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/realmlabs/realmlabs.py @@ -1,4 +1,12 @@ -"""RealmLabs MLS request guardrail with hazard screening and PII protection.""" +"""RealmLabs MLS guardrail. + +On ``pre_call`` MLS checks the request. A matching ``hazard_prompt`` verdict above ``hazard_threshold`` blocks +the request. On ``post_call`` MLS checks the assistant reply with the request's conversation as context. +Both hooks mask detected PII as ``[type]``, or block when ``pii_mask`` is False. + +Conversation context comes from Chat Completions-style ``request_data.messages``; nothing is kept between +hooks. Streaming uses LiteLLM's existing delivery settings, which do not enable text rewrites by default. +""" from __future__ import annotations @@ -20,6 +28,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # helper is untyped in http_handler httpxSpecialProvider, ) +from litellm.proxy.guardrails.guardrail_hooks.content_text import content_to_text from litellm.proxy.guardrails.guardrail_hooks.realmlabs.pii_masking import mask_pii_in_text from litellm.secret_managers.main import get_secret_str from litellm.types.guardrails import GuardrailEventHooks @@ -46,6 +55,8 @@ _DEFAULT_HAZARD_THRESHOLD: Final = 0.703 _DEFAULT_TIMEOUT: Final = 15.0 _RESPONSE_ADAPTER: Final = TypeAdapter(RealmLabsGuardrailResponse) +_MESSAGE_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_MESSAGES_ADAPTER: Final = TypeAdapter(tuple[object, ...]) class RealmLabsMissingCredentials(Exception): @@ -110,8 +121,8 @@ class RealmLabsGuardrail(CustomGuardrail): @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: # mutable-ok: base class returns a list - """Inspect requests on ``pre_call``.""" - return [GuardrailEventHooks.pre_call] + """Inspect requests on ``pre_call`` and replies on ``post_call``.""" + return [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call] @staticmethod def _hazard_score(response: RealmLabsGuardrailResponse) -> float | None: @@ -188,13 +199,13 @@ class RealmLabsGuardrail(CustomGuardrail): input_type: Literal["request", "response"], logging_obj: LiteLLMLoggingObj | None = None, ) -> GenericGuardrailAPIInputs: - """Screen requests for hazardous content and PII through LiteLLM's unified guardrail layer.""" - if input_type != "request": - return inputs + """Inspect request or response text through LiteLLM's unified guardrail layer. + + Enforce hazard scores only on requests with a matching role, then mask or block PII on either side. + Only the current hook's texts are rewritten. MLS failures pass through unless ``block_on_error`` is set. + """ texts: Final = tuple(inputs.get("texts") or ()) - messages: Final = tuple(inputs.get("structured_messages") or ()) or tuple( - RealmLabsChatMessage(role="user", content=text) for text in texts - ) + messages: Final = _messages_for_scan(inputs, request_data, input_type) if not messages: return inputs @@ -207,7 +218,7 @@ class RealmLabsGuardrail(CustomGuardrail): return self._handle_mls_error(inputs, f"Invalid RealmLabs guardrail response: {result.reason}") # Hazard is checked before PII, so a hazardous prompt is rejected rather than masked and forwarded. - hazard_score: Final = self._hazard_score(result) + hazard_score: Final = self._hazard_score(result) if input_type == "request" else None if hazard_score is not None and hazard_score > self.hazard_threshold: verbose_proxy_logger.warning( "RealmLabs MLS blocked request: %s=%s > %s", @@ -240,3 +251,44 @@ class RealmLabsGuardrail(CustomGuardrail): return inputs verbose_proxy_logger.debug("RealmLabs MLS masked PII types in the %s: %s", input_type, self._span_types(spans)) return {**inputs, "texts": list(masked_texts)} + + +def _conversation_message(message: object) -> RealmLabsChatMessage | None: + """Extract a plain-text turn without carrying images or unrelated message fields to MLS.""" + + if not isinstance(message, dict): + return None + row: Final = _MESSAGE_ADAPTER.validate_python(message) + role: Final = row.get("role") + text: Final = content_to_text(row.get("content")) + if isinstance(role, str) and role and text: + return RealmLabsChatMessage(role=role, content=text) + return None + + +def _conversation_messages(request_data: Mapping[str, object]) -> tuple[RealmLabsChatMessage, ...]: + """Preserve roles and text from request history, skipping empty or non-message entries.""" + + raw_messages: Final = request_data.get("messages") + if not isinstance(raw_messages, list): + return () + + messages: Final = _MESSAGES_ADAPTER.validate_python(raw_messages) + return tuple(message for item in messages if (message := _conversation_message(item)) is not None) + + +def _messages_for_scan( + inputs: GenericGuardrailAPIInputs, + request_data: Mapping[str, object], + input_type: Literal["request", "response"], +) -> Sequence[Mapping[str, object]]: + """Prefer scoped request messages; append response texts as assistant turns to the request history.""" + + if input_type == "request" and (structured_messages := inputs.get("structured_messages")): + return tuple(structured_messages) + + conversation: Final = _conversation_messages(request_data) + texts: Final = inputs.get("texts") or () + if input_type == "request": + return conversation or tuple(RealmLabsChatMessage(role="user", content=text) for text in texts) + return (*conversation, *(RealmLabsChatMessage(role="assistant", content=text) for text in texts)) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/realmlabs/test_realmlabs.py b/tests/unit/proxy/guardrails/guardrail_hooks/realmlabs/test_realmlabs.py index 8448bf16c53..a87a0db3a46 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/realmlabs/test_realmlabs.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/realmlabs/test_realmlabs.py @@ -10,6 +10,7 @@ from pydantic import TypeAdapter import litellm from litellm.exceptions import GuardrailRaisedException +from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler from litellm.proxy.guardrails.guardrail_hooks.realmlabs import guardrail_initializer_registry from litellm.proxy.guardrails.guardrail_hooks.realmlabs.realmlabs import ( RealmLabsGuardrail, @@ -18,7 +19,7 @@ from litellm.proxy.guardrails.guardrail_hooks.realmlabs.realmlabs import ( from litellm.types.guardrails import GuardrailEventHooks, LitellmParams from litellm.types.llms.openai import AllMessageValues from litellm.types.proxy.guardrails.guardrail_hooks.realmlabs import RealmLabsGuardrailOptionalParams -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.types.utils import Choices, GenericGuardrailAPIInputs, Message, ModelResponse _API_KEY = "mls_gr_test" _API_BASE = "https://mls.example.test" @@ -232,6 +233,7 @@ async def test_span_from_another_turn_leaves_the_prompt_as_is(respx_mock: respx. pytest.param( "request", ["Alex alex@example.com"], [_NAME_SPAN, _EMAIL_SPAN], "name, email", id="request-types" ), + pytest.param("response", ["alex@example.com"], [_EMAIL_SPAN], "email", id="response-type"), pytest.param("request", ["Hello Alex."], [{"type": "name"}], "name", id="missing-pii-text"), pytest.param("request", ["Hello Alex."], [{"type": "name", "text": None}], "name", id="null-pii-text"), pytest.param("request", ["Hello Alex."], [{"type": "name", "text": ""}], "name", id="empty-pii-text"), @@ -287,7 +289,7 @@ async def test_plain_texts_are_sent_as_user_turns(respx_mock: respx.MockRouter) @pytest.mark.asyncio -@pytest.mark.parametrize("input_type", ["request"]) +@pytest.mark.parametrize("input_type", ["request", "response"]) async def test_empty_inputs_are_returned_without_calling_mls( input_type: Literal["request", "response"], respx_mock: respx.MockRouter ) -> None: @@ -299,7 +301,7 @@ async def test_empty_inputs_are_returned_without_calling_mls( @pytest.mark.asyncio -@pytest.mark.parametrize("input_type", ["request"]) +@pytest.mark.parametrize("input_type", ["request", "response"]) @pytest.mark.parametrize("block_on_error", [None, True], ids=["default-fail-open", "fail-closed"]) @pytest.mark.parametrize( ("mls_reply", "reason"), @@ -344,7 +346,7 @@ async def test_mls_failures_follow_the_error_policy( @pytest.mark.asyncio -@pytest.mark.parametrize("input_type", ["request"]) +@pytest.mark.parametrize("input_type", ["request", "response"]) @pytest.mark.parametrize("block_on_error", [False, True], ids=["fail-open", "fail-closed"]) @pytest.mark.parametrize( "body", @@ -413,7 +415,7 @@ async def test_invalid_probabilities_cannot_bypass_fail_closed(prob: object, res @pytest.mark.asyncio -@pytest.mark.parametrize("input_type", ["request"]) +@pytest.mark.parametrize("input_type", ["request", "response"]) async def test_empty_verdict_arrays_are_valid_in_fail_closed_mode( input_type: Literal["request", "response"], respx_mock: respx.MockRouter ) -> None: @@ -687,6 +689,38 @@ async def test_configured_timeout_reaches_mls( assert timeouts["read"] == expected +@pytest.mark.asyncio +@pytest.mark.parametrize( + "use_structured_messages", [False, True], ids=["request-fallback", "structured-takes-priority"] +) +async def test_request_uses_structured_messages_before_conversation_fallback( + use_structured_messages: bool, respx_mock: respx.MockRouter +) -> None: + conversation: Final[list[AllMessageValues]] = [ + {"role": "system", "content": "Be brief."}, + {"role": "user", "content": "Hello."}, + ] + inputs: Final[GenericGuardrailAPIInputs] = ( + {"texts": ["Hello."], "structured_messages": conversation} if use_structured_messages else {"texts": ["Hello."]} + ) + request_data: Final = { + "messages": [{"role": "user", "content": "Outside the selected scope."}] + if use_structured_messages + else conversation + } + route: Final = _serve(respx_mock, _mls_body()) + + result: Final = await _apply(_guardrail(), inputs, request_data=request_data) + + assert result is inputs + assert _sent_body(route) == { + "messages": conversation, + "probes": ["hazard_prompt"], + "pii": True, + "enable_thinking": False, + } + + @pytest.mark.asyncio async def test_role_mismatch_skips_request_hazard_but_still_masks_pii(respx_mock: respx.MockRouter) -> None: _serve( @@ -717,9 +751,134 @@ async def test_missing_role_mismatch_still_enforces_request_hazard( @pytest.mark.asyncio -async def test_response_text_is_returned_without_calling_mls(respx_mock: respx.MockRouter) -> None: - route: Final = _serve(respx_mock, _mls_body()) - inputs: Final[GenericGuardrailAPIInputs] = {"texts": ["hello"]} +@pytest.mark.parametrize("name", ["Alex", "[name]"], ids=["unmasked-name", "existing-placeholder"]) +async def test_post_call_masks_the_reply_in_conversation_context(name: str, respx_mock: respx.MockRouter) -> None: + guardrail: Final = _guardrail(event_hook=GuardrailEventHooks.post_call) + conversation: Final[list[AllMessageValues]] = [ + {"role": "system", "content": "Be brief."}, + {"role": "user", "content": "My name is Alex."}, + {"role": "assistant", "content": "Hello!"}, + {"role": "user", "content": "What is my name?"}, + ] + route: Final = _serve(respx_mock, _mls_body(pii_spans=[{"type": "name", "text": name.removesuffix("]")}])) + response: Final = ModelResponse( + choices=[ + Choices(index=0, message=Message(content=f"Your name is {name}.", role="assistant"), finish_reason="stop") + ] + ) - assert await _apply(_guardrail(), inputs, input_type="response") is inputs - assert route.call_count == 0 + result: Final = await OpenAIChatCompletionsHandler().process_output_response( # pyright: ignore[reportUnknownMemberType] # upstream request_data parameter uses an unparameterized dict + response=response, guardrail_to_apply=guardrail, request_data={"messages": conversation} + ) + + assert result.choices == [ + Choices(index=0, message=Message(content="Your name is [name].", role="assistant"), finish_reason="stop") + ], result.choices + assert conversation == [ + {"role": "system", "content": "Be brief."}, + {"role": "user", "content": "My name is Alex."}, + {"role": "assistant", "content": "Hello!"}, + {"role": "user", "content": "What is my name?"}, + ], conversation + assert _sent_body(route) == { + "messages": [*conversation, {"role": "assistant", "content": f"Your name is {name}."}], + "probes": ["hazard_prompt"], + "pii": True, + "enable_thinking": False, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("input_type", ["request", "response"]) +async def test_conversation_text_parts_use_blank_lines_without_changing_internal_paragraphs( + input_type: Literal["request", "response"], respx_mock: respx.MockRouter +) -> None: + conversation: Final = [ + {"role": "system", "content": "Be brief."}, + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this.\n\nKeep this paragraph."}, + {"type": "image_url", "image_url": {"url": "https://example.test/image.png"}}, + {"type": "text", "text": "In one sentence."}, + ], + }, + {"role": "assistant", "content": None}, + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "https://example.test/other.png"}}]}, + {"content": "No role."}, + None, + ] + route: Final = _serve(respx_mock, _mls_body()) + inputs: Final[GenericGuardrailAPIInputs] = {"texts": ["A landscape."]} + + result: Final = await _apply(_guardrail(), inputs, input_type=input_type, request_data={"messages": conversation}) + + expected_history: Final = [ + {"role": "system", "content": "Be brief."}, + {"role": "user", "content": "Describe this.\n\nKeep this paragraph.\n\nIn one sentence."}, + ] + expected_messages: Final = ( + [*expected_history, {"role": "assistant", "content": "A landscape."}] + if input_type == "response" + else expected_history + ) + + assert result is inputs + assert _sent_body(route) == { + "messages": expected_messages, + "probes": ["hazard_prompt"], + "pii": True, + "enable_thinking": False, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_data", [{}, {"messages": []}, {"messages": None}], ids=["missing", "empty", "null"]) +async def test_response_without_history_is_still_scanned_as_assistant_text( + request_data: dict[str, object], respx_mock: respx.MockRouter +) -> None: + route: Final = _serve(respx_mock, _mls_body(pii_spans=[_NAME_SPAN])) + + result: Final = await _apply( + _guardrail(), {"texts": ["Alex was here."]}, input_type="response", request_data=request_data + ) + + assert result == {"texts": ["[name] was here."]}, result + assert _sent_body(route) == { + "messages": [{"role": "assistant", "content": "Alex was here."}], + "probes": ["hazard_prompt"], + "pii": True, + "enable_thinking": False, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("role_mismatch", [False, True, None], ids=["matching-role", "mismatched-role", "null-role"]) +async def test_response_ignores_hazard_scores_and_merges_overlapping_pii( + role_mismatch: bool | None, respx_mock: respx.MockRouter +) -> None: + _serve( + respx_mock, + { + "results": [{"probe": "hazard_prompt", "prob": 0.99, "role_mismatch": role_mismatch}], + "pii_spans": [{"type": "name", "text": "Ann"}, {"type": "email", "text": "Ann.Smith@example.com"}], + }, + ) + inputs: Final[GenericGuardrailAPIInputs] = {"texts": ["Contact Ann at Ann.Smith@example.com.", "No PII here."]} + + result: Final = await _apply(_guardrail(), inputs, input_type="response") + + assert result == {"texts": ["Contact [name] at [email].", "No PII here."]}, result + assert inputs == {"texts": ["Contact Ann at Ann.Smith@example.com.", "No PII here."]}, inputs + + +@pytest.mark.asyncio +async def test_pii_only_in_history_leaves_the_reply_unchanged(respx_mock: respx.MockRouter) -> None: + _serve(respx_mock, _mls_body(pii_spans=[_NAME_SPAN])) + inputs: Final[GenericGuardrailAPIInputs] = {"texts": ["Hello there."]} + conversation: Final = [{"role": "user", "content": "My name is Alex."}] + + result: Final = await _apply(_guardrail(), inputs, input_type="response", request_data={"messages": conversation}) + + assert result is inputs + assert conversation == [{"role": "user", "content": "My name is Alex."}], conversation