From 5070f5dd4b575cf111179c663d2741b6ee86722e Mon Sep 17 00:00:00 2001 From: berri-teddy Date: Sat, 20 Sep 2025 22:15:35 -0700 Subject: [PATCH 1/2] Ensure basic detection format proper --- git_model_armor.py | 0 .../model_armor/model_armor.py | 159 ++++++++++++------ test_model_armor.py | 0 3 files changed, 112 insertions(+), 47 deletions(-) create mode 100644 git_model_armor.py create mode 100644 test_model_armor.py diff --git a/git_model_armor.py b/git_model_armor.py new file mode 100644 index 00000000000..e69de29bb2d 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 480be4a651b..31d7d70d5f3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -88,11 +88,11 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): def _create_sanitize_request( self, content: str, source: Literal["user_prompt", "model_response"] ) -> dict: - """Create request body for Model Armor API.""" + """Create request body for Model Armor API with correct camelCase field names.""" if source == "user_prompt": - return {"user_prompt_data": {"text": content}} + return {"userPromptData": {"text": content}} else: - return {"model_response_data": {"text": content}} + return {"modelResponseData": {"text": content}} def _extract_content_from_response( self, response: Union[Any, ModelResponse] @@ -119,11 +119,16 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): async def make_model_armor_request( self, - content: str, - source: Literal["user_prompt", "model_response"], + content: Optional[str] = None, + source: Literal["user_prompt", "model_response"] = "user_prompt", request_data: Optional[dict] = None, + file_bytes: Optional[bytes] = None, + file_type: Optional[str] = None, ) -> dict: - """Make request to Model Armor API.""" + """ + Make request to Model Armor API. Supports both text and file prompt sanitization. + If file_bytes and file_type are provided, file prompt sanitization is performed. + """ # Get access token using VertexBase auth access_token, resolved_project_id = await self._ensure_access_token_async( credentials=self.credentials, @@ -143,7 +148,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): url = f"{endpoint}/v1/projects/{self.project_id}/locations/{self.location}/templates/{self.template_id}:sanitizeModelResponse" # Create request body - body = self._create_sanitize_request(content, source) + if file_bytes is not None and file_type is not None: + body = self.sanitize_file_prompt(file_bytes, file_type, source) + elif content is not None: + body = self._create_sanitize_request(content, source) + else: + raise ValueError( + "Either content or file_bytes and file_type must be provided." + ) # Set headers headers = { @@ -189,57 +201,110 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): return await json_response return json_response + def sanitize_file_prompt( + self, file_bytes: bytes, file_type: str, source: str = "user_prompt" + ) -> dict: + """ + Helper to build the request body for file prompt sanitization for Model Armor. + file_type should be one of: PLAINTEXT_UTF8, PDF, WORD_DOCUMENT, EXCEL_DOCUMENT, POWERPOINT_DOCUMENT, TXT, CSV + Returns the request body dict. + """ + import base64 + + base64_data = base64.b64encode(file_bytes).decode("utf-8") + if source == "user_prompt": + return { + "userPromptData": { + "byteItem": {"byteDataType": file_type, "byteData": base64_data} + } + } + else: + return { + "modelResponseData": { + "byteItem": {"byteDataType": file_type, "byteData": base64_data} + } + } + def _should_block_content(self, armor_response: dict) -> bool: - """Check if Model Armor response indicates content should be blocked.""" - # Check the sanitizationResult from Model Armor API + """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", {}) - # 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 - - # 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 + # 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] + elif isinstance(filter_results, list): + filter_result_items = filter_results + for filt in filter_result_items: + # Check RAI, PI/Jailbreak, Malicious URI, CSAM, Virus scan as before + if filt.get("raiFilterResult", {}).get("matchState") == "MATCH_FOUND": + return True + if ( + filt.get("piAndJailbreakFilterResult", {}).get("matchState") + == "MATCH_FOUND" + ): + return True + if ( + filt.get("maliciousUriFilterResult", {}).get("matchState") + == "MATCH_FOUND" + ): + return True + if ( + filt.get("csamFilterFilterResult", {}).get("matchState") + == "MATCH_FOUND" + ): + return True + if filt.get("virusScanFilterResult", {}).get("matchState") == "MATCH_FOUND": + return True + # Check sdpFilterResult for both inspectResult and deidentifyResult + sdp = filt.get("sdpFilterResult") + if sdp: + if sdp.get("inspectResult", {}).get("matchState") == "MATCH_FOUND": + return True + if sdp.get("deidentifyResult", {}).get("matchState") == "MATCH_FOUND": + return True + # Fallback dict code removed; all cases handled above return False def _get_sanitized_content(self, armor_response: dict) -> Optional[str]: - """Extract sanitized content from Model Armor response.""" - # Model Armor returns sanitized content in the sanitizationResult - sanitization_result = armor_response.get("sanitizationResult", {}) + """ + Get the sanitized content from a Model Armor response, if available. + Looks for sanitized text in deidentifyResult, and falls back to root-level fields if not found. + """ + result = armor_response.get("sanitizationResult", {}) + filter_results = result.get("filterResults", {}) - # Check for sdp structure (for deidentification) - filter_results = sanitization_result.get("filterResults", {}) - sdp = filter_results.get("sdp", {}).get("sdpFilterResult") + # filterResults can be a dict (single filter) or a list (multiple filters) + filters = ( + [filter_results] + if isinstance(filter_results, dict) + else filter_results + if isinstance(filter_results, list) + else [] + ) - 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 + # Prefer sanitized text from deidentifyResult if present + for filter_entry in filters: + sdp = filter_entry.get("sdpFilterResult") + if sdp: + deid = sdp.get("deidentifyResult", {}) + sanitized = deid.get("data", {}).get("text", "") + # If Model Armor found something and returned a sanitized version, use it + if deid.get("matchState") == "MATCH_FOUND" and sanitized: + return sanitized - # Fallback to checking root level + # If no deidentifyResult, optionally check for inspectResult (rare, but could have findings) + for filter_entry in filters: + sdp = filter_entry.get("sdpFilterResult") + if sdp: + inspect = sdp.get("inspectResult", {}) + # If Model Armor flagged something but didn't sanitize, return None + if inspect.get("matchState") == "MATCH_FOUND": + return None + + # Fallback: if Model Armor put sanitized text at the root, use it return armor_response.get("sanitizedText") or armor_response.get("text") def _process_response( diff --git a/test_model_armor.py b/test_model_armor.py new file mode 100644 index 00000000000..e69de29bb2d From 9d1878029dc52f861cf8f592f59c9b50f821d78a Mon Sep 17 00:00:00 2001 From: Teddy Amkie Date: Sat, 20 Sep 2025 22:40:46 -0700 Subject: [PATCH 2/2] Added model armor testing files basic --- git_model_armor.py | 0 test_model_armor.py | 0 .../test_model_armor_file_sanitization.py | 75 ++++++++++++++ .../test_model_armor_guardrail.py | 99 +++++++++++++++++++ 4 files changed, 174 insertions(+) delete mode 100644 git_model_armor.py delete mode 100644 test_model_armor.py create mode 100644 tests/guardrails_tests/test_model_armor_file_sanitization.py create mode 100644 tests/guardrails_tests/test_model_armor_guardrail.py diff --git a/git_model_armor.py b/git_model_armor.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/test_model_armor.py b/test_model_armor.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/guardrails_tests/test_model_armor_file_sanitization.py b/tests/guardrails_tests/test_model_armor_file_sanitization.py new file mode 100644 index 00000000000..4c9826ee14d --- /dev/null +++ b/tests/guardrails_tests/test_model_armor_file_sanitization.py @@ -0,0 +1,75 @@ +import sys +import os +import pytest +from unittest.mock import AsyncMock +from fastapi import HTTPException + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import ModelArmorGuardrail + +def test_sanitize_file_prompt_builds_pdf_body(): + guardrail = ModelArmorGuardrail( + template_id="dummy-template", + project_id="dummy-project", + location="us-central1", + credentials=None, + ) + file_bytes = b"%PDF-1.4 some pdf content" + file_type = "PDF" + body = guardrail.sanitize_file_prompt(file_bytes, file_type, source="user_prompt") + assert "userPromptData" in body + assert body["userPromptData"]["byteItem"]["byteDataType"] == "PDF" + import base64 + assert body["userPromptData"]["byteItem"]["byteData"] == base64.b64encode(file_bytes).decode("utf-8") + +@pytest.mark.asyncio +async def test_make_model_armor_request_file_prompt(): + guardrail = ModelArmorGuardrail( + template_id="dummy-template", + project_id="dummy-project", + location="us-central1", + credentials=None, + ) + file_bytes = b"My SSN is 123-45-6789." + file_type = "PLAINTEXT_UTF8" + armor_response = { + "sanitizationResult": { + "filterResults": [ + { + "sdpFilterResult": { + "inspectResult": { + "executionState": "EXECUTION_SUCCESS", + "matchState": "MATCH_FOUND", + "findings": [ + {"infoType": "US_SOCIAL_SECURITY_NUMBER", "likelihood": "LIKELY"} + ] + }, + "deidentifyResult": { + "executionState": "EXECUTION_SUCCESS", + "matchState": "MATCH_FOUND", + "data": {"text": "My SSN is [REDACTED]."} + } + } + } + ] + } + } + class MockResponse: + def __init__(self, status_code, text, json_data): + self.status_code = status_code + self.text = text + self._json = json_data + def json(self): + return self._json + class MockHandler: + async def post(self, url, json, headers): + return MockResponse(200, str(armor_response), armor_response) + guardrail.async_handler = MockHandler() + guardrail._ensure_access_token_async = AsyncMock(return_value=("dummy-token", "dummy-project")) + result = await guardrail.make_model_armor_request( + file_bytes=file_bytes, + file_type=file_type, + source="user_prompt" + ) + assert result["sanitizationResult"]["filterResults"][0]["sdpFilterResult"]["deidentifyResult"]["data"]["text"] == "My SSN is [REDACTED]." diff --git a/tests/guardrails_tests/test_model_armor_guardrail.py b/tests/guardrails_tests/test_model_armor_guardrail.py new file mode 100644 index 00000000000..51e5859d081 --- /dev/null +++ b/tests/guardrails_tests/test_model_armor_guardrail.py @@ -0,0 +1,99 @@ +import sys +import os +import pytest +from unittest.mock import AsyncMock, patch +from fastapi import HTTPException + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import ModelArmorGuardrail +from litellm.proxy._types import UserAPIKeyAuth +from litellm.caching.caching import DualCache + +@pytest.mark.asyncio +async def test_model_armor_pre_call_hook_inspect_and_deidentify(): + """ + Test Model Armor guardrail pre-call hook for both inspectResult and deidentifyResult handling. + """ + guardrail = ModelArmorGuardrail( + template_id="dummy-template", + project_id="dummy-project", + location="us-central1", + credentials=None, + ) + armor_response = { + "sanitizationResult": { + "filterResults": [ + { + "sdpFilterResult": { + "inspectResult": { + "executionState": "EXECUTION_SUCCESS", + "matchState": "NO_MATCH_FOUND", + "findings": [] + }, + "deidentifyResult": { + "executionState": "EXECUTION_SUCCESS", + "matchState": "MATCH_FOUND", + "data": {"text": "sanitized text here"} + } + } + } + ] + } + } + with patch.object(guardrail, "make_model_armor_request", AsyncMock(return_value=armor_response)): + user_api_key_dict = UserAPIKeyAuth(api_key="test_key") + cache = DualCache() + data = { + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "My SSN is 123-45-6789."} + ], + "model": "gpt-3.5-turbo", + "metadata": {} + } + guardrail.mask_request_content = True + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="completion" + ) + assert exc_info.value.status_code == 400 + assert "Content blocked by Model Armor" in str(exc_info.value.detail) + +def test_model_armor_should_block_content(): + guardrail = ModelArmorGuardrail( + template_id="dummy-template", + project_id="dummy-project", + location="us-central1", + credentials=None, + ) + # Block on inspectResult + armor_response_inspect = { + "sanitizationResult": { + "filterResults": [ + {"sdpFilterResult": {"inspectResult": {"matchState": "MATCH_FOUND"}}} + ] + } + } + assert guardrail._should_block_content(armor_response_inspect) + # Block on deidentifyResult + armor_response_deidentify = { + "sanitizationResult": { + "filterResults": [ + {"sdpFilterResult": {"deidentifyResult": {"matchState": "MATCH_FOUND"}}} + ] + } + } + assert guardrail._should_block_content(armor_response_deidentify) + # No block if neither + armor_response_none = { + "sanitizationResult": { + "filterResults": [ + {"sdpFilterResult": {"inspectResult": {"matchState": "NO_MATCH_FOUND"}, "deidentifyResult": {"matchState": "NO_MATCH_FOUND"}}} + ] + } + } + assert not guardrail._should_block_content(armor_response_none)