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..ae7611975ea 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -108,6 +108,7 @@ class _Filter(TypedDict): @with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) class _SanitizationResult(TypedDict): + invocationResult: NotRequired[ReadOnly[str]] filterResults: NotRequired[ReadOnly[Mapping[str, _Filter] | Sequence[_Filter] | str | int | float | bool | None]] @@ -154,6 +155,9 @@ class ModelArmorAPIError(Exception): self.detail = detail +# invocationResult values that mean Model Armor did not finish scanning, although it answered HTTP 200 +_INCOMPLETE_INVOCATION_RESULTS: Final = frozenset({"PARTIAL", "FAILURE"}) + _SCANNED_CONTENT_KEYS: Final = frozenset({"text", "sanitizedText", "findings", "maliciousUriMatchedItems"}) @@ -419,8 +423,36 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): json_response: Final = _RESPONSE_JSON.validate_python(response.json()) if isinstance(json_response, _AwaitableJSON): - return _ModelArmorResponseRoot.model_validate(await json_response).root - return _ModelArmorResponseRoot.model_validate(json_response).root + armor_response: Final = _ModelArmorResponseRoot.model_validate(await json_response).root + else: + armor_response = _ModelArmorResponseRoot.model_validate(json_response).root + # What the caller can do with a de-identification match: files and text with masking off are + # blocked on it, only masked text is rewritten and let through. + if file_bytes is not None: + masks_match = False + elif source == "user_prompt": + masks_match = self.mask_request_content + else: + masks_match = self.mask_response_content + self._raise_if_scan_incomplete(armor_response, allow_sanitization=masks_match) + return armor_response + + def _raise_if_scan_incomplete(self, armor_response: _ModelArmorResponse, allow_sanitization: bool) -> None: + """A 200 whose invocationResult is PARTIAL or FAILURE means some filters did not run. + + Accepting it would treat an unfinished scan as a clean one. It is raised like any other API + failure, so every hook applies fail_on_error. A scan that already matched a filter is still + returned, so the caller blocks instead of failing open on content that was flagged. A + de-identification match counts as one only where the caller blocks on it (files, or text with + masking off); where it is masked instead, it must not hide the filters that did not run. + """ + result: Final = armor_response.get("sanitizationResult", {}).get("invocationResult") + if result not in _INCOMPLETE_INVOCATION_RESULTS: + return + if self._should_block_content(armor_response, allow_sanitization=allow_sanitization): + return + verbose_proxy_logger.error("Model Armor scan did not complete - invocationResult: %s", result) + raise ModelArmorAPIError(f"Model Armor scan incomplete (invocationResult={result})") def sanitize_file_prompt(self, file_bytes: bytes, file_type: str, source: str = "user_prompt") -> dict: """ 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..af0c2d4f883 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -658,6 +658,174 @@ async def test_model_armor_api_failure_raises_sanitized_error(): assert "Internal Server Error" not in str(exc_info.value.detail) +def _armor_guardrail(**kwargs) -> ModelArmorGuardrail: + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + **kwargs, + ) + guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) + guardrail.should_run_guardrail = Mock(return_value=True) + return guardrail + + +def _armor_http_response(sanitization_result: dict) -> Mock: + response = Mock() + response.status_code = 200 + response.text = json.dumps({"sanitizationResult": sanitization_result}) + response.json = Mock(return_value={"sanitizationResult": sanitization_result}) + return response + + +@pytest.mark.asyncio +@pytest.mark.parametrize("invocation_result", ["PARTIAL", "FAILURE"]) +async def test_model_armor_incomplete_scan_raises_api_error(invocation_result: str): + """A 200 whose invocationResult is PARTIAL or FAILURE is an unfinished scan, not a clean one.""" + guardrail = _armor_guardrail() + response = _armor_http_response( + {"invocationResult": invocation_result, "filterMatchState": "NO_MATCH_FOUND", "filterResults": {}} + ) + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=response)): + with pytest.raises(ModelArmorAPIError) as exc_info: + await guardrail.make_model_armor_request(content="test content", source="user_prompt") + + assert exc_info.value.detail == f"Model Armor scan incomplete (invocationResult={invocation_result})" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "sanitization_result", + [ + {"invocationResult": "SUCCESS", "filterMatchState": "NO_MATCH_FOUND", "filterResults": {}}, + {"invocationResult": "INVOCATION_RESULT_UNSPECIFIED", "filterResults": {}}, + {"filterMatchState": "NO_MATCH_FOUND", "filterResults": {}}, + ], +) +async def test_model_armor_complete_or_unreported_scan_is_accepted(sanitization_result: dict): + """SUCCESS, an unspecified result and a response without the field are all treated as before.""" + guardrail = _armor_guardrail() + + with patch.object( + guardrail.async_handler, "post", AsyncMock(return_value=_armor_http_response(sanitization_result)) + ): + result = await guardrail.make_model_armor_request(content="test content", source="user_prompt") + + assert result["sanitizationResult"] == sanitization_result + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_on_error", [True, False]) +async def test_model_armor_incomplete_scan_honors_fail_on_error(fail_on_error: bool): + """An unfinished scan blocks when fail_on_error is true and passes through when fail-open is configured.""" + guardrail = _armor_guardrail(fail_on_error=fail_on_error) + response = _armor_http_response({"invocationResult": "PARTIAL", "filterResults": {}}) + request_data = {"model": "gpt-4", "messages": [{"role": "user", "content": "synthetic input"}], "metadata": {}} + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=response)): + if fail_on_error: + with pytest.raises(ModelArmorAPIError): + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + else: + result = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + assert result is request_data + + +@pytest.mark.asyncio +async def test_model_armor_partial_scan_with_match_still_blocks_when_fail_open(): + """A partial scan that already found a match must block, even with fail_on_error disabled.""" + guardrail = _armor_guardrail(fail_on_error=False) + response = _armor_http_response( + { + "invocationResult": "PARTIAL", + "filterMatchState": "MATCH_FOUND", + "filterResults": {"rai": {"raiFilterResult": {"matchState": "MATCH_FOUND"}}}, + } + ) + request_data = {"model": "gpt-4", "messages": [{"role": "user", "content": "synthetic input"}], "metadata": {}} + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=response)): + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + + +_PARTIAL_DEIDENTIFY_ONLY = { + "invocationResult": "PARTIAL", + "filterResults": { + "sdp": { + "sdpFilterResult": { + "deidentifyResult": {"matchState": "MATCH_FOUND", "data": {"text": "call me at [REDACTED]"}} + } + } + }, +} + + +@pytest.mark.asyncio +async def test_model_armor_partial_scan_with_masked_deidentify_match_is_an_error(): + """Where a de-identification match is masked, not blocked, it must not hide the filters that did not run.""" + guardrail = _armor_guardrail(mask_request_content=True) + response = _armor_http_response(_PARTIAL_DEIDENTIFY_ONLY) + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=response)): + with pytest.raises(ModelArmorAPIError) as exc_info: + await guardrail.make_model_armor_request(content="call me at 555-0100", source="user_prompt") + + assert exc_info.value.detail == "Model Armor scan incomplete (invocationResult=PARTIAL)" + + +@pytest.mark.asyncio +async def test_model_armor_partial_scan_with_blocking_deidentify_match_still_blocks_when_fail_open(): + """With masking off the caller blocks on a de-identification match, so an incomplete scan must not turn it + into an error that fail_on_error=False then swallows.""" + guardrail = _armor_guardrail(fail_on_error=False, mask_request_content=False) + response = _armor_http_response(_PARTIAL_DEIDENTIFY_ONLY) + request_data = {"model": "gpt-4", "messages": [{"role": "user", "content": "call me at 555-0100"}], "metadata": {}} + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=response)): + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_model_armor_partial_file_scan_with_deidentify_match_is_returned_even_with_masking_on(): + """Files cannot be masked, so the caller blocks on the match whatever the masking setting is.""" + guardrail = _armor_guardrail(mask_request_content=True) + response = _armor_http_response(_PARTIAL_DEIDENTIFY_ONLY) + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=response)): + result = await guardrail.make_model_armor_request(file_bytes=b"call me at 555-0100", file_type="PLAINTEXT_UTF8") + + assert result["sanitizationResult"]["invocationResult"] == "PARTIAL" + assert guardrail._should_block_content(result, allow_sanitization=False) + + @pytest.mark.asyncio @pytest.mark.parametrize("sanitize", [True, False]) async def test_model_armor_error_output_sanitization(sanitize: bool):