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 31d7d70d5f..787c46d0dd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -21,11 +21,11 @@ from litellm.integrations.custom_guardrail import ( 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,7 +225,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): } } - 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", {}) @@ -233,7 +233,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # 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 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): 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 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # 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 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # 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 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): 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={ diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index 1865f58290..b596d427c5 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -1,19 +1,22 @@ -import sys -import os -import io, asyncio -import pytest +import asyncio +import io import json -from unittest.mock import MagicMock, AsyncMock, patch, Mock +import os +import sys +from unittest.mock import AsyncMock, MagicMock, Mock, patch + +import pytest 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