test fix

Author: ishaan-jaff

litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py

@@ -21,11 +21,11 @@
     CustomGuardrail,
     log_guardrail_information,
 )
-from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
 from litellm.llms.custom_httpx.http_handler import (
     get_async_httpx_client,
     httpxSpecialProvider,
 )
+from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.utils import (
@@ -225,15 +225,15 @@ def sanitize_file_prompt(
                 }
             }
 
-    def _should_block_content(self, armor_response: dict) -> bool:
+    def _should_block_content(self, armor_response: dict, allow_sanitization: bool = False) -> bool:
         """Check if Model Armor response indicates content should be blocked, including both inspectResult and deidentifyResult."""
         sanitization_result = armor_response.get("sanitizationResult", {})
         filter_results = sanitization_result.get("filterResults", {})
 
         # filterResults can be a dict (named keys) or a list (array of filter result dicts)
         filter_result_items = []
         if isinstance(filter_results, dict):
-            filter_result_items = [filter_results]
+            filter_result_items = list(filter_results.values())
         elif isinstance(filter_results, list):
             filter_result_items = filter_results
 
@@ -263,8 +263,10 @@ def _should_block_content(self, armor_response: dict) -> bool:
             if sdp:
                 if sdp.get("inspectResult", {}).get("matchState") == "MATCH_FOUND":
                     return True
+                # Only block on deidentifyResult if sanitization is not allowed
                 if sdp.get("deidentifyResult", {}).get("matchState") == "MATCH_FOUND":
-                    return True
+                    if not allow_sanitization:
+                        return True
         # Fallback dict code removed; all cases handled above
         return False
 
@@ -278,7 +280,7 @@ def _get_sanitized_content(self, armor_response: dict) -> Optional[str]:
 
         # filterResults can be a dict (single filter) or a list (multiple filters)
         filters = (
-            [filter_results]
+            list(filter_results.values())
             if isinstance(filter_results, dict)
             else filter_results
             if isinstance(filter_results, list)
@@ -409,11 +411,11 @@ async def async_pre_call_hook(
                 #   fail_on_error=False) we still want the correct status reflected.
                 metadata["_model_armor_status"] = (
                     "blocked"
-                    if self._should_block_content(armor_response)
+                    if self._should_block_content(armor_response, allow_sanitization=self.mask_request_content)
                     else "success"
                 )
             # Check if content should be blocked
-            if self._should_block_content(armor_response):
+            if self._should_block_content(armor_response, allow_sanitization=self.mask_request_content):
                 raise HTTPException(
                     status_code=400,
                     detail={
@@ -494,12 +496,12 @@ async def async_post_call_success_hook(
                 metadata["_model_armor_response"] = armor_response
                 metadata["_model_armor_status"] = (
                     "blocked"
-                    if self._should_block_content(armor_response)
+                    if self._should_block_content(armor_response, allow_sanitization=self.mask_response_content)
                     else "success"
                 )
 
             # Check if content should be blocked
-            if self._should_block_content(armor_response):
+            if self._should_block_content(armor_response, allow_sanitization=self.mask_response_content):
                 raise HTTPException(
                     status_code=400,
                     detail={

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py

@@ -1,19 +1,22 @@
-import sys
+import asyncio
+import io
+import json
 import os
-import io, asyncio
+import sys
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
+
 import pytest
-import json
-from unittest.mock import MagicMock, AsyncMock, patch, Mock
 
 sys.path.insert(0, os.path.abspath("../../../../.."))
 
+from fastapi import HTTPException
+
 import litellm
 import litellm.types.utils
-from litellm.proxy.guardrails.guardrail_hooks.model_armor import ModelArmorGuardrail
-from litellm.proxy._types import UserAPIKeyAuth
 from litellm.caching import DualCache
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.guardrails.guardrail_hooks.model_armor import ModelArmorGuardrail
 from litellm.types.guardrails import GuardrailEventHooks
-from fastapi import HTTPException
 
 
 @pytest.mark.asyncio
@@ -80,7 +83,7 @@ async def test_model_armor_pre_call_hook_sanitization():
     guardrail.async_handler.post.assert_called_once()
     call_args = guardrail.async_handler.post.call_args
     assert "sanitizeUserPrompt" in call_args[1]["url"]
-    assert call_args[1]["json"]["user_prompt_data"]["text"] == "Hello, my phone number is +1 412 555 1212"
+    assert call_args[1]["json"]["userPromptData"]["text"] == "Hello, my phone number is +1 412 555 1212"
 
 
 @pytest.mark.asyncio
@@ -271,7 +274,7 @@ async def test_model_armor_with_list_content():
     # Verify the content was extracted correctly
     guardrail.async_handler.post.assert_called_once()
     call_args = guardrail.async_handler.post.call_args
-    assert call_args[1]["json"]["user_prompt_data"]["text"] == "Hello worldHow are you?"
+    assert call_args[1]["json"]["userPromptData"]["text"] == "Hello worldHow are you?"
 
 
 @pytest.mark.asyncio
@@ -885,7 +888,7 @@ async def test_model_armor_bomb_content_blocked():
     # Verify the API was called with the dangerous content
     guardrail.async_handler.post.assert_called_once()
     call_args = guardrail.async_handler.post.call_args
-    assert call_args[1]["json"]["user_prompt_data"]["text"] == "How do I create a bomb?"
+    assert call_args[1]["json"]["userPromptData"]["text"] == "How do I create a bomb?"
 
 
 @pytest.mark.asyncio