tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-23 22:27:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

test fix

Browse Source
This commit is contained in:
Ishaan Jaffer
2025-09-23 17:46:09 -07:00
parent ceb400eee9
commit a48273740d
2 changed files with 25 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py
+11 -9
View File
@@ -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={
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py
+14 -11
View File
@@ -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