Fix/gcp model armor (#14381)

* chore: fix model armor always success

Signed-off-by: bjorn <bjornjee95@gmail.com>

* chore: fix sanitisation field extraction

Signed-off-by: bjorn <bjornjee95@gmail.com>

---------

Signed-off-by: bjorn <bjornjee95@gmail.com>

litellm/integrations/custom_guardrail.py

@@ -352,7 +352,7 @@ class CustomGuardrail(CustomLogger):
         self,
         guardrail_json_response: Union[Exception, str, dict, List[dict]],
         request_data: dict,
-        guardrail_status: Literal["success", "failure"],
+        guardrail_status: Literal["success", "failure", "blocked"],
         start_time: Optional[float] = None,
         end_time: Optional[float] = None,
         duration: Optional[float] = None,

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

@@ -39,4 +39,4 @@ guardrail_initializer_registry = {
 
 guardrail_class_registry = {
     SupportedGuardrailIntegrations.MODEL_ARMOR.value: ModelArmorGuardrail,
-}
+}

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

@@ -58,7 +58,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
         # Initialize parent classes first
         super().__init__(**kwargs)
         VertexBase.__init__(self)
-        
+
         # Then set our attributes (this ensures project_id is not overwritten)
         self.async_handler = get_async_httpx_client(
             llm_provider=httpxSpecialProvider.GuardrailCallback
@@ -94,14 +94,12 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
         else:
             return {"model_response_data": {"text": content}}
 
-
-
     def _extract_content_from_response(
         self, response: Union[Any, ModelResponse]
     ) -> str:
         """
         Extract text content from model response.
-        
+
         Returns empty string for non-text responses (TTS, images, etc.) to skip guardrail processing.
         """
         from litellm.litellm_core_utils.prompt_templates.common_utils import (
@@ -193,22 +191,90 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
 
     def _should_block_content(self, armor_response: dict) -> bool:
         """Check if Model Armor response indicates content should be blocked."""
-        # Model Armor may return different response structures
-        # This is a basic implementation - adjust based on actual API response
-        if armor_response.get("blocked", False):
+        # Check the sanitizationResult from Model Armor API
+        sanitization_result = armor_response.get("sanitizationResult", {})
+        filter_results = sanitization_result.get("filterResults", {})
+
+        # Check blocking filters (these should cause the request to be blocked)
+        # RAI (Responsible AI) filters
+        rai_results = filter_results.get("rai", {}).get("raiFilterResult", {})
+        if rai_results.get("matchState") == "MATCH_FOUND":
             return True
 
-        # Check for sanitization actions
-        if armor_response.get("action") == "BLOCK":
+        # Prompt injection and jailbreak filters
+        pi_jailbreak = filter_results.get("piAndJailbreakFilterResult", {})
+        if pi_jailbreak.get("matchState") == "MATCH_FOUND":
+            return True
+
+        # Malicious URI filters
+        malicious_uri = filter_results.get("maliciousUriFilterResult", {})
+        if malicious_uri.get("matchState") == "MATCH_FOUND":
+            return True
+
+        # CSAM filters
+        csam = filter_results.get("csamFilterFilterResult", {})
+        if csam.get("matchState") == "MATCH_FOUND":
+            return True
+
+        # Virus scan filters
+        virus_scan = filter_results.get("virusScanFilterResult", {})
+        if virus_scan.get("matchState") == "MATCH_FOUND":
             return True
 
         return False
 
     def _get_sanitized_content(self, armor_response: dict) -> Optional[str]:
         """Extract sanitized content from Model Armor response."""
-        # This depends on the actual Model Armor API response structure
-        # Adjust based on documentation
-        return armor_response.get("sanitized_text") or armor_response.get("text")
+        # Model Armor returns sanitized content in the sanitizationResult
+        sanitization_result = armor_response.get("sanitizationResult", {})
+
+        # Check for sdp structure (for deidentification)
+        filter_results = sanitization_result.get("filterResults", {})
+        sdp = filter_results.get("sdp", {}).get("sdpFilterResult")
+
+        if sdp is not None:
+            # Model Armor returns sanitized text under deidentifyResult in sdp
+            deidentify_result = sdp.get("deidentifyResult", {})
+            sanitized_text = deidentify_result.get("data", {}).get("text", "")
+            if deidentify_result.get("matchState") == "MATCH_FOUND" and sanitized_text:
+                return sanitized_text
+
+        # Fallback to checking root level
+        return armor_response.get("sanitizedText") or armor_response.get("text")
+
+    def _process_response(
+        self,
+        response: Optional[dict],
+        request_data: dict,
+        start_time: Optional[float] = None,
+        end_time: Optional[float] = None,
+        duration: Optional[float] = None,
+    ):
+        """
+        Override to store only the Model Armor API response, not the entire data dict.
+        This prevents circular references in logging.
+        """
+        # Retrieve the Model Armor response & status stored on the per-request `metadata` object.
+        metadata = (
+            request_data.get("metadata", {}) if isinstance(request_data, dict) else {}
+        )
+
+        guardrail_response = metadata.get("_model_armor_response", {})
+
+        # Determine status – default to "success" but prefer the explicit value if present.
+        guardrail_status: Literal["success", "failure", "blocked"] = metadata.get(
+            "_model_armor_status", "success"
+        )  # type: ignore
+
+        self.add_standard_logging_guardrail_information_to_request_data(
+            guardrail_json_response=guardrail_response,
+            request_data=request_data,
+            guardrail_status=guardrail_status,  # type: ignore
+            duration=duration,
+            start_time=start_time,
+            end_time=end_time,
+        )
+        return response
 
     @log_guardrail_information
     async def async_pre_call_hook(
@@ -263,6 +329,24 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
                 request_data=data,
             )
 
+            # Store the armor response for logging
+            # Attach Model Armor response + evaluation status directly to the per-request metadata to avoid
+            #   race-conditions between concurrent requests which share the same guardrail instance.
+            #   This ensures each request logs its own Model Armor response instead of a potentially stale value
+            #   overwritten by another coroutine.
+            if isinstance(data, dict):
+                metadata = data.setdefault(
+                    "metadata", {}
+                )  # ensures metadata exists and is unique per request
+                metadata["_model_armor_response"] = armor_response
+                # Pre-compute guardrail status for downstream logging. A blocked response will eventually raise
+                #   an HTTPException, however in scenarios where the caller decides to ignore the exception (e.g.
+                #   fail_on_error=False) we still want the correct status reflected.
+                metadata["_model_armor_status"] = (
+                    "blocked"
+                    if self._should_block_content(armor_response)
+                    else "success"
+                )
             # Check if content should be blocked
             if self._should_block_content(armor_response):
                 raise HTTPException(
@@ -339,6 +423,16 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
                 request_data=data,
             )
 
+            # Attach Model Armor response & status to this request's metadata to prevent race conditions
+            if isinstance(data, dict):
+                metadata = data.setdefault("metadata", {})
+                metadata["_model_armor_response"] = armor_response
+                metadata["_model_armor_status"] = (
+                    "blocked"
+                    if self._should_block_content(armor_response)
+                    else "success"
+                )
+
             # Check if content should be blocked
             if self._should_block_content(armor_response):
                 raise HTTPException(
@@ -406,6 +500,16 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
                         request_data=request_data,
                     )
 
+                    # Attach Model Armor response & status to this request's metadata to avoid race conditions
+                    if isinstance(request_data, dict):
+                        metadata = request_data.setdefault("metadata", {})
+                        metadata["_model_armor_response"] = armor_response
+                        metadata["_model_armor_status"] = (
+                            "blocked"
+                            if self._should_block_content(armor_response)
+                            else "success"
+                        )
+
                     # Check if blocked
                     if self._should_block_content(armor_response):
                         raise HTTPException(

litellm/types/utils.py

@@ -1995,7 +1995,7 @@ class StandardLoggingGuardrailInformation(TypedDict, total=False):
     ]
     guardrail_request: Optional[dict]
     guardrail_response: Optional[Union[dict, str, List[dict]]]
-    guardrail_status: Literal["success", "failure"]
+    guardrail_status: Literal["success", "failure","blocked"]
     start_time: Optional[float]
     end_time: Optional[float]
     duration: Optional[float]

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

@@ -34,8 +34,21 @@ async def test_model_armor_pre_call_hook_sanitization():
     mock_response = AsyncMock()
     mock_response.status_code = 200
     mock_response.json = AsyncMock(return_value={
-        "sanitized_text": "Hello, my phone number is [REDACTED]",
-        "action": "SANITIZE"
+        "sanitizationResult": {
+            "filterMatchState": "MATCH_FOUND",
+            "filterResults": {
+                "sdp": {
+                    "sdpFilterResult": {
+                        "deidentifyResult": {
+                            "matchState": "MATCH_FOUND",
+                            "data": {
+                                "text":"Hello, my phone number is [REDACTED]"
+                            },
+                        }
+                    }
+                }
+            }
+        }
     })
     
     # Mock the access token method
@@ -87,9 +100,22 @@ async def test_model_armor_pre_call_hook_blocked():
     mock_response = AsyncMock()
     mock_response.status_code = 200
     mock_response.json = AsyncMock(return_value={
-        "action": "BLOCK",
-        "blocked": True,
-        "reason": "Prohibited content detected"
+        "sanitizationResult": {
+            "filterMatchState": "MATCH_FOUND",
+            "filterResults": {
+                "rai": {
+                    "raiFilterResult": {
+                        "matchState": "MATCH_FOUND",
+                        "raiFilterTypeResults": {
+                            "dangerous": {
+                                "matchState": "MATCH_FOUND",
+                                "reason": "Prohibited content detected"
+                            }
+                        }
+                    }
+                }
+            }
+        }
     })
     
     # Mock the access token method
@@ -137,8 +163,21 @@ async def test_model_armor_post_call_hook_sanitization():
     mock_response = AsyncMock()
     mock_response.status_code = 200
     mock_response.json = AsyncMock(return_value={
-        "sanitized_text": "Here is the information: [REDACTED]",
-        "action": "SANITIZE"
+        "sanitizationResult": {
+            "filterMatchState": "MATCH_FOUND",
+            "filterResults": {
+                "sdp": {
+                    "sdpFilterResult": {
+                        "deidentifyResult": {
+                            "matchState": "MATCH_FOUND",
+                            "data": {
+                                "text":"Here is the information: [REDACTED]"
+                            },
+                        }
+                    }
+                }
+            }
+        }
     })
     
     # Mock the access token method
@@ -196,7 +235,9 @@ async def test_model_armor_with_list_content():
     mock_response = AsyncMock()
     mock_response.status_code = 200
     mock_response.json = AsyncMock(return_value={
-        "action": "NONE"
+        "sanitizationResult": {
+            "filterMatchState": "NO_MATCH_FOUND"
+        }
     })
     
     # Mock the access token method
@@ -328,8 +369,10 @@ async def test_model_armor_streaming_response():
     mock_response = AsyncMock()
     mock_response.status_code = 200
     mock_response.json = AsyncMock(return_value={
-        "sanitized_text": "Sanitized response",
-        "action": "SANITIZE"
+        "sanitizationResult": {
+            "filterMatchState": "NO_MATCH_FOUND",
+            "sanitizedText": "Sanitized response"
+        }
     })
     
     # Mock the access token method
@@ -614,10 +657,14 @@ async def test_model_armor_action_none():
         mask_request_content=True,
     )
     
-    # Mock response with action=NONE
+    # Mock response with action=NO_MATCH_FOUND
     mock_response = AsyncMock()
     mock_response.status_code = 200
-    mock_response.json = AsyncMock(return_value={"action": "NONE"})
+    mock_response.json = AsyncMock(return_value={
+        "sanitizationResult": {
+            "filterMatchState": "NO_MATCH_FOUND"
+        }
+    })
     
     guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project"))
     guardrail.async_handler = AsyncMock()
@@ -658,8 +705,9 @@ async def test_model_armor_missing_sanitized_text():
     mock_response = AsyncMock()
     mock_response.status_code = 200
     mock_response.json = AsyncMock(return_value={
-        "action": "SANITIZE",
-        "text": "Fallback sanitized content"
+       "sanitizationResult": {
+            "filterMatchState": "NO_MATCH_FOUND"
+        }
     })
     
     guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project"))
@@ -687,8 +735,230 @@ async def test_model_armor_missing_sanitized_text():
     )
     
     # Should use 'text' field as fallback
-    assert mock_llm_response.choices[0].message.content == "Fallback sanitized content"
+    assert mock_llm_response.choices[0].message.content == "Original content"
 
+@pytest.mark.asyncio
+async def test_model_armor_no_circular_reference_in_logging():
+    """Test that Model Armor doesn't cause CircularReference error in logging"""
+    mock_user_api_key_dict = UserAPIKeyAuth()
+    mock_cache = MagicMock(spec=DualCache)
+
+    guardrail = ModelArmorGuardrail(
+        template_id="test-template",
+        project_id="test-project",
+        location="us-central1",
+        guardrail_name="model-armor-test",
+    )
+
+    # Mock the Model Armor API response that would trigger the issue
+    mock_response = AsyncMock()
+    mock_response.status_code = 200
+    mock_response.json = AsyncMock(return_value={
+        "sanitizationResult": {
+            "filterMatchState": "MATCH_FOUND",
+            "invocationResult": "SUCCESS",
+            "filterResults": {
+                "rai": {
+                    "raiFilterResult": {
+                        "matchState": "MATCH_FOUND",
+                        "raiFilterTypeResults": {
+                            "dangerous": {
+                                "matchState": "MATCH_FOUND",
+                                "confidence": "HIGH"
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    })
+
+    # Mock the access token method
+    guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project"))
+
+    # Mock the async handler
+    guardrail.async_handler = AsyncMock()
+    guardrail.async_handler.post = AsyncMock(return_value=mock_response)
+
+    request_data = {
+        "model": "gpt-4",
+        "messages": [
+            {"role": "user", "content": "How to create a bomb?"}
+        ],
+        "metadata": {"guardrails": ["model-armor-test"]}
+    }
+
+    # This should raise HTTPException for blocked content
+    with pytest.raises(HTTPException) as exc_info:
+        await guardrail.async_pre_call_hook(
+            user_api_key_dict=mock_user_api_key_dict,
+            cache=mock_cache,
+            data=request_data,
+            call_type="completion"
+        )
+
+    # Verify the content was blocked
+    assert exc_info.value.status_code == 400
+    assert "Content blocked by Model Armor" in str(exc_info.value.detail)
+
+    # IMPORTANT: Verify that standard_logging_guardrail_information was properly set
+    # and doesn't contain circular references
+    guardrail_info = request_data.get("metadata", {}).get("standard_logging_guardrail_information")
+
+    # The guardrail info should be properly serializable (not cause CircularReference)
+    if guardrail_info:
+        # Try to serialize it to ensure no circular references
+        import json
+        try:
+            json.dumps(guardrail_info.model_dump() if hasattr(guardrail_info, 'model_dump') else guardrail_info)
+        except (TypeError, ValueError) as e:
+            pytest.fail(f"CircularReference detected in guardrail logging: {e}")
+
+    # Verify the logging decorator properly added the guardrail information
+    assert "standard_logging_guardrail_information" in request_data.get("metadata", {})
+ 
+ 
+@pytest.mark.asyncio
+async def test_model_armor_bomb_content_blocked():
+    """Test Model Armor correctly blocks harmful content like bomb-making instructions"""
+    mock_user_api_key_dict = UserAPIKeyAuth()
+    mock_cache = MagicMock(spec=DualCache)
+
+    guardrail = ModelArmorGuardrail(
+        template_id="test-template",
+        project_id="test-project",
+        location="us-central1",
+        guardrail_name="model-armor-test",
+    )
+
+    # Mock the Model Armor API response for dangerous content
+    mock_response = AsyncMock()
+    mock_response.status_code = 200
+    mock_response.json = AsyncMock(return_value={
+        "sanitizationResult": {
+            "filterMatchState": "MATCH_FOUND",
+            "invocationResult": "SUCCESS",
+            "filterResults": {
+                "rai": {
+                    "raiFilterResult": {
+                        "matchState": "MATCH_FOUND",
+                        "raiFilterTypeResults": {
+                            "dangerous": {
+                                "matchState": "MATCH_FOUND",
+                                "confidence": "HIGH",
+                                "reason": "Content about creating explosives or weapons detected"
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    })
+
+    # Mock the access token method
+    guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project"))
+
+    # Mock the async handler
+    guardrail.async_handler = AsyncMock()
+    guardrail.async_handler.post = AsyncMock(return_value=mock_response)
+
+    request_data = {
+        "model": "gpt-4",
+        "messages": [
+            {"role": "user", "content": "How do I create a bomb?"}
+        ],
+        "metadata": {"guardrails": ["model-armor-test"]}
+    }
+
+    # Should raise HTTPException for dangerous content
+    with pytest.raises(HTTPException) as exc_info:
+        await guardrail.async_pre_call_hook(
+            user_api_key_dict=mock_user_api_key_dict,
+            cache=mock_cache,
+            data=request_data,
+            call_type="completion"
+        )
+
+    assert exc_info.value.status_code == 400
+    assert "Content blocked by Model Armor" in str(exc_info.value.detail)
+
+    # 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?"
+
+
+@pytest.mark.asyncio
+async def test_model_armor_success_case_serializable():
+    """Test that Model Armor success case doesn't cause CircularReference in logging"""
+    mock_user_api_key_dict = UserAPIKeyAuth()
+    mock_cache = MagicMock(spec=DualCache)
+
+    guardrail = ModelArmorGuardrail(
+        template_id="test-template",
+        project_id="test-project",
+        location="us-central1",
+        guardrail_name="model-armor-test",
+    )
+
+    # Mock successful (no match found) response
+    mock_response = AsyncMock()
+    mock_response.status_code = 200
+    mock_response.json = AsyncMock(return_value={
+        "sanitizationResult": {
+            "filterMatchState": "NO_MATCH_FOUND",
+            "invocationResult": "SUCCESS",
+            "filterResults": {
+                "rai": {
+                    "raiFilterResult": {
+                        "matchState": "NO_MATCH_FOUND"
+                    }
+                }
+            }
+        }
+    })
+
+    # Mock the access token method
+    guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project"))
+
+    # Mock the async handler
+    guardrail.async_handler = AsyncMock()
+    guardrail.async_handler.post = AsyncMock(return_value=mock_response)
+
+    request_data = {
+        "model": "gpt-4",
+        "messages": [
+            {"role": "user", "content": "What is the weather today?"}
+        ],
+        "metadata": {"guardrails": ["model-armor-test"]}
+    }
+
+    # This should NOT raise an exception - content is allowed
+    result = await guardrail.async_pre_call_hook(
+        user_api_key_dict=mock_user_api_key_dict,
+        cache=mock_cache,
+        data=request_data,
+        call_type="completion"
+    )
+
+    # Verify the request was allowed through
+    assert result == request_data
+
+    # IMPORTANT: Verify that standard_logging_guardrail_information is serializable
+    guardrail_info = request_data.get("metadata", {}).get("standard_logging_guardrail_information")
+
+    # The guardrail info should exist and be properly serializable
+    assert guardrail_info is not None
+
+    # Try to serialize it to ensure no circular references
+    import json
+    try:
+        # This should NOT raise any exception
+        serialized = json.dumps(guardrail_info.model_dump() if hasattr(guardrail_info, 'model_dump') else guardrail_info)
+        # Verify it's not the string "CircularReference Detected"
+        assert "CircularReference Detected" not in serialized
+    except (TypeError, ValueError) as e:
+        pytest.fail(f"CircularReference detected in guardrail logging for success case: {e}")
 
 @pytest.mark.asyncio
 async def test_model_armor_non_text_response():