diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py index 366ea7326ae..a714ab618ba 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py @@ -86,6 +86,17 @@ class _FileBlock(LiteLLMBaseModel): file: _FileData +class _ResponsesFileBlock(LiteLLMBaseModel): + """A Responses API `input_file` item: inline `file_data`, or a `file_id` / `file_url` reference.""" + + model_config = ConfigDict(extra="ignore") + type: Literal["input_file"] + file_data: str | None = None + file_url: str | None = None + file_id: str | None = None + filename: str | None = None + + class _DocumentSource(LiteLLMBaseModel): model_config = ConfigDict(extra="ignore") data: str | None = None @@ -98,8 +109,8 @@ class _DocumentBlock(LiteLLMBaseModel): source: _DocumentSource -_AttachmentBlock = Annotated[_FileBlock | _DocumentBlock, Field(discriminator="type")] -_BLOCK_ADAPTER: Final[TypeAdapter[_FileBlock | _DocumentBlock]] = TypeAdapter(_AttachmentBlock) +_AttachmentBlock = Annotated[_FileBlock | _DocumentBlock | _ResponsesFileBlock, Field(discriminator="type")] +_BLOCK_ADAPTER: Final[TypeAdapter[_FileBlock | _DocumentBlock | _ResponsesFileBlock]] = TypeAdapter(_AttachmentBlock) def plan_file_scans(messages: Sequence[AllMessageValues]) -> FileScanPlan: @@ -152,9 +163,17 @@ def _is_reference(block: _FileBlock | _DocumentBlock) -> bool: def _parse_block(block: object) -> _FileBlock | _DocumentBlock | None: try: - return _BLOCK_ADAPTER.validate_python(block) + parsed: Final = _BLOCK_ADAPTER.validate_python(block) except ValidationError: return None + if isinstance(parsed, _ResponsesFileBlock): + # Same shape as a chat `file` block; a file_url or file_id alone has no inline bytes, so it is + # classified as a reference and fails closed like its chat equivalent. + return _FileBlock( + type="file", + file=_FileData(file_data=parsed.file_data or parsed.file_url, filename=parsed.filename), + ) + return parsed def _block_byte_data_type_and_data( 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..476f0dd2c4f 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, @@ -154,6 +155,8 @@ class ModelArmorAPIError(Exception): self.detail = detail +_RESPONSES_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value}) + _SCANNED_CONTENT_KEYS: Final = frozenset({"text", "sanitizedText", "findings", "maliciousUriMatchedItems"}) @@ -701,6 +704,47 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): detail=self._build_block_error_detail("Content blocked by Model Armor", armor_response), ) + @staticmethod + def _scans_responses_input(data: dict, call_type: CallTypesLiteral) -> bool: + """Whether the scan reads the Responses API `input` rather than `messages`.""" + if call_type in _RESPONSES_CALL_TYPES: + return True + return not data.get("messages") and data.get("input") is not None + + @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 +766,13 @@ 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. For those + # call types `input` is what reaches the provider, so a `messages` field sent alongside it is ignored. + scans_responses_input: Final = self._scans_responses_input(data, call_type) + responses_input: Final = data.get("input") if scans_responses_input else None + messages: Final = ( + self._responses_input_messages(responses_input) if scans_responses_input else data.get("messages") + ) if not messages: verbose_proxy_logger.warning("Model Armor: not running guardrail. No messages in data") return data @@ -787,7 +837,10 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): set_last_user_message, ) - data["messages"] = set_last_user_message(messages, sanitized_content) + if scans_responses_input: + 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 +874,13 @@ 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. For those + # call types `input` is what reaches the provider, so a `messages` field sent alongside it is ignored. + scans_responses_input: Final = self._scans_responses_input(data, call_type) + responses_input: Final = data.get("input") if scans_responses_input else None + messages: Final = ( + self._responses_input_messages(responses_input) if scans_responses_input else data.get("messages") + ) if not messages: verbose_proxy_logger.warning("Model Armor: not running guardrail. No messages in data") return data @@ -878,7 +937,10 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): set_last_user_message, ) - data["messages"] = set_last_user_message(messages, sanitized_content) + if scans_responses_input: + 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..eb4bfdab714 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,242 @@ 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, call_type: str = "responses") -> dict: + return await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=data, + call_type=call_type, + ) + + +@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_chat_call_scans_messages_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, call_type="completion") + + 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 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["responses", "aresponses"]) +async def test_model_armor_responses_call_scans_input_not_a_decoy_messages_field(call_type: str): + """`input` is what reaches the provider on a Responses call, so a benign `messages` sent next to it is ignored.""" + guardrail = _responses_guardrail(_RAI_MATCH) + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello"}], + "input": "synthetic harmful prompt", + "metadata": {}, + } + + with pytest.raises(HTTPException) as exc_info: + await _run_pre_call(guardrail, data, call_type=call_type) + + 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_masking_replaces_input_even_with_a_decoy_messages_field(): + guardrail = _responses_guardrail(_MASKED, mask_request_content=True) + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello"}], + "input": "call me at 555-0100", + "metadata": {}, + } + + result = await _run_pre_call(guardrail, data) + + assert result["input"] == "call me at [REDACTED]" + assert result["messages"] == [{"role": "user", "content": "hello"}] + + +@pytest.mark.asyncio +async def test_model_armor_responses_inline_input_file_is_scanned(): + import base64 + + guardrail = _responses_guardrail(_RAI_MATCH) + document = b"synthetic document text" + data = { + "model": "gpt-4", + "input": [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "summarize this"}, + { + "type": "input_file", + "filename": "notes.txt", + "file_data": "data:text/plain;base64," + base64.b64encode(document).decode(), + }, + ], + } + ], + "metadata": {}, + } + + with pytest.raises(HTTPException) as exc_info: + await _run_pre_call(guardrail, data) + + assert exc_info.value.status_code == 400 + first_call = guardrail.make_model_armor_request.await_args_list[0].kwargs + assert first_call["file_bytes"] == document + assert first_call["file_type"] == "TXT" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "file_item", + [ + {"type": "input_file", "file_id": "file-abc123"}, + {"type": "input_file", "file_url": "https://example.test/report.pdf"}, + ], +) +async def test_model_armor_responses_input_file_reference_fails_closed(file_item: dict): + guardrail = _responses_guardrail(_NO_MATCH) + data = { + "model": "gpt-4", + "input": [{"role": "user", "content": [{"type": "input_text", "text": "summarize this"}, file_item]}], + "metadata": {}, + } + + with pytest.raises(HTTPException) as exc_info: + await _run_pre_call(guardrail, data) + + assert exc_info.value.status_code == 400 + guardrail.make_model_armor_request.assert_not_awaited() +