From 285b14bd69e0b7907a4e31c8a7c2d1a9a785f490 Mon Sep 17 00:00:00 2001 From: pinkgu <39329656+bjornjee@users.noreply.github.com> Date: Wed, 10 Sep 2025 00:18:37 +0800 Subject: [PATCH] Fix/gcp model armor (#14381) * chore: fix model armor always success Signed-off-by: bjorn * chore: fix sanitisation field extraction Signed-off-by: bjorn --------- Signed-off-by: bjorn --- litellm/integrations/custom_guardrail.py | 2 +- .../guardrail_hooks/model_armor/__init__.py | 2 +- .../model_armor/model_armor.py | 128 +++++++- litellm/types/utils.py | 2 +- .../guardrail_hooks/test_model_armor.py | 300 +++++++++++++++++- 5 files changed, 404 insertions(+), 30 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 1ca45f907e..40d2137a7f 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/__init__.py index 4be98a19db..7398e8defe 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/__init__.py @@ -39,4 +39,4 @@ guardrail_initializer_registry = { guardrail_class_registry = { SupportedGuardrailIntegrations.MODEL_ARMOR.value: ModelArmorGuardrail, -} \ No newline at end of file +} 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 ee04e899c6..480be4a651 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/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( diff --git a/litellm/types/utils.py b/litellm/types/utils.py index b1a32469b9..c6f7098a1a 100644 --- a/litellm/types/utils.py +++ b/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] 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 fc408efdf3..1865f58290 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 @@ -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():