Merge b0c8f4836e into aeede53112

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:
         """

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):