onyx-dot-app/litellm
1
0
Fork 0
mirror of https://github.com/onyx-dot-app/litellm.git synced 2026-06-30 20:47:56 -04:00
Code Issues Packages Projects Releases Wiki Activity

Revert "Added model armor testing files basic"

Browse Source 
This reverts commit 9d1878029d.
This commit is contained in:
Ishaan Jaffer
2025-09-23 17:45:44 -07:00
parent 79a6feb7ac
commit ceb400eee9
4 changed files with 0 additions and 174 deletions
Download Patch File Download Diff File Expand all files Collapse all files
git_model_armor.py
View File
test_model_armor.py
View File
tests/guardrails_tests/test_model_armor_file_sanitization.py
-75
View File
@@ -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]."
tests/guardrails_tests/test_model_armor_guardrail.py
-99
View File
@@ -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)