From ceb400eee9aa2163ee63ee908f7ec57a600d1830 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 23 Sep 2025 17:45:44 -0700 Subject: [PATCH] Revert "Added model armor testing files basic" This reverts commit 9d1878029dc52f861cf8f592f59c9b50f821d78a. --- 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 deletions(-) create mode 100644 git_model_armor.py create mode 100644 test_model_armor.py delete mode 100644 tests/guardrails_tests/test_model_armor_file_sanitization.py delete mode 100644 tests/guardrails_tests/test_model_armor_guardrail.py diff --git a/git_model_armor.py b/git_model_armor.py new file mode 100644 index 000000000..e69de29bb diff --git a/test_model_armor.py b/test_model_armor.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/guardrails_tests/test_model_armor_file_sanitization.py b/tests/guardrails_tests/test_model_armor_file_sanitization.py deleted file mode 100644 index 4c9826ee1..000000000 --- a/tests/guardrails_tests/test_model_armor_file_sanitization.py +++ /dev/null @@ -1,75 +0,0 @@ -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 deleted file mode 100644 index 51e5859d0..000000000 --- a/tests/guardrails_tests/test_model_armor_guardrail.py +++ /dev/null @@ -1,99 +0,0 @@ -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)