diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index b9871b1e276..cfba6e3c9be 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -45,6 +45,7 @@ from litellm.proxy.guardrails.guardrail_hooks.model_armor.file_scanning import ( MODEL_ARMOR_MAX_FILE_SIZE_BYTES, plan_file_scans, ) +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.guardrails import GuardrailEventHooks, LitellmParams from litellm.types.llms.openai import ( AllMessageValues, @@ -701,6 +702,40 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): detail=self._build_block_error_detail("Content blocked by Model Armor", armor_response), ) + @staticmethod + def _responses_input_messages(responses_input: object) -> list[AllMessageValues]: + """The role-carrying items of a Responses API `input`, in the chat message shape the scan reads.""" + if not isinstance(responses_input, (str, list)): + return [] + return ResponsesAPIRequestUtils.responses_input_to_chat_messages(responses_input) + + @staticmethod + def _mask_responses_input(responses_input: object, sanitized_content: str) -> str | list[object]: + """Replace the last block of user messages in a Responses API `input` with the sanitized text. + + Mirrors set_last_user_message for chat messages: the trailing run of user messages becomes one + user message, and items without a role (tool calls, tool outputs) keep their place. + """ + if isinstance(responses_input, str): + return sanitized_content + items: Final = list(responses_input) if isinstance(responses_input, list) else [] + trailing_user: Final[list[int]] = [] + for index in range(len(items) - 1, -1, -1): + item = items[index] + if not isinstance(item, dict) or "role" not in item: + continue + if item["role"] != "user": + break + trailing_user.append(index) + if not trailing_user: + return items + first: Final = min(trailing_user) + return [ + {"role": "user", "content": sanitized_content} if index == first else item + for index, item in enumerate(items) + if index not in trailing_user or index == first + ] + @log_guardrail_information async def async_pre_call_hook( self, @@ -722,7 +757,9 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): if self.should_run_guardrail(data=data, event_type=event_type) is not True: return data - messages: Final = data.get("messages") + # Responses API requests carry `input` instead of `messages`; scan it as chat messages + responses_input: Final = data.get("input") if not data.get("messages") else None + messages: Final = data.get("messages") or self._responses_input_messages(responses_input) if not messages: verbose_proxy_logger.warning("Model Armor: not running guardrail. No messages in data") return data @@ -787,7 +824,10 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): set_last_user_message, ) - data["messages"] = set_last_user_message(messages, sanitized_content) + if responses_input is not None: + data["input"] = self._mask_responses_input(responses_input, sanitized_content) + else: + data["messages"] = set_last_user_message(messages, sanitized_content) except ModelArmorAPIError as e: self._raise_if_fail_closed(e) @@ -821,7 +861,9 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): if self.should_run_guardrail(data=data, event_type=event_type) is not True: return data - messages: Final = data.get("messages") + # Responses API requests carry `input` instead of `messages`; scan it as chat messages + responses_input: Final = data.get("input") if not data.get("messages") else None + messages: Final = data.get("messages") or self._responses_input_messages(responses_input) if not messages: verbose_proxy_logger.warning("Model Armor: not running guardrail. No messages in data") return data @@ -878,7 +920,10 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): set_last_user_message, ) - data["messages"] = set_last_user_message(messages, sanitized_content) + if responses_input is not None: + data["input"] = self._mask_responses_input(responses_input, sanitized_content) + else: + data["messages"] = set_last_user_message(messages, sanitized_content) except ModelArmorAPIError as e: self._raise_if_fail_closed(e) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py index 5aa3c7dfc17..6d515d1a7c0 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -5318,3 +5318,150 @@ async def test_apply_guardrail_raises_on_flagged_when_not_logging_only(): entries = request_data["metadata"]["standard_logging_guardrail_information"] flagged = [e for e in entries if e["guardrail_status"] == "guardrail_flagged"] assert len(flagged) == 1 + + +_RAI_MATCH = {"sanitizationResult": {"filterResults": {"rai": {"raiFilterResult": {"matchState": "MATCH_FOUND"}}}}} +_NO_MATCH = {"sanitizationResult": {"filterMatchState": "NO_MATCH_FOUND", "filterResults": {}}} +_MASKED = { + "sanitizationResult": { + "filterResults": { + "sdp": { + "sdpFilterResult": { + "deidentifyResult": {"matchState": "MATCH_FOUND", "data": {"text": "call me at [REDACTED]"}} + } + } + } + } +} + + +def _responses_guardrail(armor_response: dict, **kwargs) -> ModelArmorGuardrail: + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + **kwargs, + ) + guardrail.should_run_guardrail = Mock(return_value=True) + guardrail.make_model_armor_request = AsyncMock(return_value=armor_response) + return guardrail + + +async def _run_pre_call(guardrail: ModelArmorGuardrail, data: dict) -> dict: + return await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=data, + call_type="responses", + ) + + +@pytest.mark.asyncio +async def test_model_armor_responses_string_input_is_scanned_and_blocked(): + """A Responses API request has `input`, not `messages`, and must not skip the scan.""" + guardrail = _responses_guardrail(_RAI_MATCH) + data = {"model": "gpt-4", "input": "synthetic harmful prompt", "metadata": {}} + + with pytest.raises(HTTPException) as exc_info: + await _run_pre_call(guardrail, data) + + assert exc_info.value.status_code == 400 + assert guardrail.make_model_armor_request.await_args.kwargs["content"] == "synthetic harmful prompt" + + +@pytest.mark.asyncio +async def test_model_armor_responses_list_input_scans_last_user_text(): + guardrail = _responses_guardrail(_NO_MATCH) + data = { + "model": "gpt-4", + "input": [ + {"role": "user", "content": "first question"}, + {"role": "assistant", "content": "an answer"}, + {"type": "function_call_output", "call_id": "call_1", "output": "tool text"}, + {"role": "user", "content": [{"type": "input_text", "text": "follow-up question"}]}, + ], + "metadata": {}, + } + + result = await _run_pre_call(guardrail, data) + + assert guardrail.make_model_armor_request.await_args.kwargs["content"] == "follow-up question" + assert result["input"] == data["input"] + + +@pytest.mark.asyncio +async def test_model_armor_responses_masks_string_input(): + guardrail = _responses_guardrail(_MASKED, mask_request_content=True) + data = {"model": "gpt-4", "input": "call me at 555-0100", "metadata": {}} + + result = await _run_pre_call(guardrail, data) + + assert result["input"] == "call me at [REDACTED]" + assert "messages" not in result + + +@pytest.mark.asyncio +async def test_model_armor_responses_masks_trailing_user_items_and_keeps_other_items(): + guardrail = _responses_guardrail(_MASKED, mask_request_content=True) + tool_output = {"type": "function_call_output", "call_id": "call_1", "output": "tool text"} + data = { + "model": "gpt-4", + "input": [ + {"role": "user", "content": "earlier question"}, + {"role": "assistant", "content": "an answer"}, + {"role": "user", "content": "call me"}, + tool_output, + {"role": "user", "content": "at 555-0100"}, + ], + "metadata": {}, + } + + result = await _run_pre_call(guardrail, data) + + assert result["input"] == [ + {"role": "user", "content": "earlier question"}, + {"role": "assistant", "content": "an answer"}, + {"role": "user", "content": "call me at [REDACTED]"}, + tool_output, + ] + + +@pytest.mark.asyncio +async def test_model_armor_messages_take_precedence_over_input(): + guardrail = _responses_guardrail(_NO_MATCH) + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "chat message"}], + "input": "responses input", + "metadata": {}, + } + + await _run_pre_call(guardrail, data) + + assert guardrail.make_model_armor_request.await_args.kwargs["content"] == "chat message" + + +@pytest.mark.asyncio +async def test_model_armor_no_messages_and_no_usable_input_skips_scan(): + guardrail = _responses_guardrail(_NO_MATCH) + + for data in ({"model": "gpt-4", "metadata": {}}, {"model": "gpt-4", "input": None, "metadata": {}}): + assert await _run_pre_call(guardrail, data) is data + + guardrail.make_model_armor_request.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_model_armor_moderation_hook_scans_responses_input(): + guardrail = _responses_guardrail(_RAI_MATCH) + data = {"model": "gpt-4", "input": "synthetic harmful prompt", "metadata": {}} + + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_moderation_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(), + call_type="responses", + ) + + assert exc_info.value.status_code == 400