mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
feat: support during_call for model armor guardrails (#15970)
Signed-off-by: bjornjee <bjornjee95@gmail.com>
This commit is contained in:
pinkgu
GitHub
parent
49b9bd3cad
commit
4758e2998b
2 changed files with 444 additions and 6 deletions
|
|
@ -226,7 +226,9 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
}
|
||||
}
|
||||
|
||||
def _should_block_content(self, armor_response: dict, allow_sanitization: bool = False) -> 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", {})
|
||||
|
|
@ -413,11 +415,15 @@ 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, allow_sanitization=self.mask_request_content)
|
||||
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, allow_sanitization=self.mask_request_content):
|
||||
if self._should_block_content(
|
||||
armor_response, allow_sanitization=self.mask_request_content
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
|
|
@ -456,6 +462,109 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
|
||||
return data
|
||||
|
||||
@log_guardrail_information
|
||||
async def async_moderation_hook(
|
||||
self,
|
||||
data: dict,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
call_type: Literal[
|
||||
"completion",
|
||||
"embeddings",
|
||||
"image_generation",
|
||||
"moderation",
|
||||
"audio_transcription",
|
||||
"responses",
|
||||
"mcp_call",
|
||||
"anthropic_messages",
|
||||
],
|
||||
) -> Union[Exception, str, dict, None]:
|
||||
"""During-call hook to sanitize user prompts in parallel with LLM call."""
|
||||
verbose_proxy_logger.debug("Inside Model Armor Moderation Hook")
|
||||
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
add_guardrail_to_applied_guardrails_header,
|
||||
)
|
||||
|
||||
event_type = GuardrailEventHooks.during_call
|
||||
if self.should_run_guardrail(data=data, event_type=event_type) is not True:
|
||||
return data
|
||||
|
||||
messages = data.get("messages")
|
||||
if not messages:
|
||||
verbose_proxy_logger.warning(
|
||||
"Model Armor: not running guardrail. No messages in data"
|
||||
)
|
||||
return data
|
||||
|
||||
# Extract content from messages
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
get_last_user_message,
|
||||
)
|
||||
|
||||
content = get_last_user_message(messages)
|
||||
if not content:
|
||||
return data
|
||||
|
||||
# Make Model Armor request
|
||||
try:
|
||||
armor_response = await self.make_model_armor_request(
|
||||
content=content,
|
||||
source="user_prompt",
|
||||
request_data=data,
|
||||
)
|
||||
|
||||
# Store the armor response for logging
|
||||
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, allow_sanitization=self.mask_request_content
|
||||
)
|
||||
else "success"
|
||||
)
|
||||
|
||||
# Check if content should be blocked
|
||||
if self._should_block_content(
|
||||
armor_response, allow_sanitization=self.mask_request_content
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": "Content blocked by Model Armor",
|
||||
"model_armor_response": armor_response,
|
||||
},
|
||||
)
|
||||
|
||||
# If mask_request_content is enabled, update messages with sanitized content
|
||||
if self.mask_request_content:
|
||||
sanitized_content = self._get_sanitized_content(armor_response)
|
||||
if sanitized_content and sanitized_content != content:
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
set_last_user_message,
|
||||
)
|
||||
|
||||
data["messages"] = set_last_user_message(
|
||||
messages, sanitized_content
|
||||
)
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
"Model Armor moderation error: %s", str(e), exc_info=True
|
||||
)
|
||||
if self.optional_params.get("fail_on_error", True):
|
||||
raise
|
||||
|
||||
# Add guardrail to headers
|
||||
add_guardrail_to_applied_guardrails_header(
|
||||
request_data=data, guardrail_name=self.guardrail_name
|
||||
)
|
||||
|
||||
return data
|
||||
|
||||
@log_guardrail_information
|
||||
async def async_post_call_success_hook(
|
||||
self,
|
||||
|
|
@ -498,12 +607,16 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
metadata["_model_armor_response"] = armor_response
|
||||
metadata["_model_armor_status"] = (
|
||||
"blocked"
|
||||
if self._should_block_content(armor_response, allow_sanitization=self.mask_response_content)
|
||||
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, allow_sanitization=self.mask_response_content):
|
||||
if self._should_block_content(
|
||||
armor_response, allow_sanitization=self.mask_response_content
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
|
|
|
|||
|
|
@ -1166,4 +1166,329 @@ async def test_model_armor_with_default_credentials():
|
|||
# Verify the project_id was used correctly in the API call
|
||||
guardrail.async_handler.post.assert_called_once()
|
||||
call_args = guardrail.async_handler.post.call_args
|
||||
assert "cloud-test-project" in call_args[1]["url"]
|
||||
assert "cloud-test-project" in call_args[1]["url"]
|
||||
|
||||
|
||||
# ===== ASYNC MODERATION HOOK TESTS =====
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_moderation_hook_success_no_blocking():
|
||||
"""Test async_moderation_hook with successful response (no blocking)"""
|
||||
mock_user_api_key_dict = UserAPIKeyAuth()
|
||||
|
||||
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",
|
||||
"filterResults": {
|
||||
"rai": {
|
||||
"raiFilterResult": {
|
||||
"matchState": "NO_MATCH_FOUND"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
# Mock the access token method and async handler
|
||||
guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project"))
|
||||
guardrail.async_handler = AsyncMock()
|
||||
guardrail.async_handler.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
request_data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello, how are you?"}
|
||||
],
|
||||
"metadata": {"guardrails": ["model-armor-test"]}
|
||||
}
|
||||
|
||||
result = await guardrail.async_moderation_hook(
|
||||
data=request_data,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
call_type="completion"
|
||||
)
|
||||
|
||||
# Should return the original data unchanged
|
||||
assert result == request_data
|
||||
# Should have metadata added
|
||||
assert "_model_armor_response" in request_data["metadata"]
|
||||
assert request_data["metadata"]["_model_armor_status"] == "success"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_moderation_hook_content_blocked():
|
||||
"""Test async_moderation_hook when content should be blocked"""
|
||||
mock_user_api_key_dict = UserAPIKeyAuth()
|
||||
|
||||
guardrail = ModelArmorGuardrail(
|
||||
template_id="test-template",
|
||||
project_id="test-project",
|
||||
location="us-central1",
|
||||
guardrail_name="model-armor-test",
|
||||
)
|
||||
|
||||
# Mock response that indicates content should be blocked
|
||||
mock_response = AsyncMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json = AsyncMock(return_value={
|
||||
"sanitizationResult": {
|
||||
"filterMatchState": "MATCH_FOUND",
|
||||
"filterResults": {
|
||||
"rai": {
|
||||
"raiFilterResult": {
|
||||
"matchState": "MATCH_FOUND"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
# Mock the access token method and async handler
|
||||
guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project"))
|
||||
guardrail.async_handler = AsyncMock()
|
||||
guardrail.async_handler.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
request_data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Some harmful content"}
|
||||
],
|
||||
"metadata": {"guardrails": ["model-armor-test"]}
|
||||
}
|
||||
|
||||
# Should raise HTTPException for blocked content
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await guardrail.async_moderation_hook(
|
||||
data=request_data,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
call_type="completion"
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "Content blocked by Model Armor" in str(exc_info.value.detail)
|
||||
# Should have metadata added even when blocked
|
||||
assert "_model_armor_response" in request_data["metadata"]
|
||||
assert request_data["metadata"]["_model_armor_status"] == "blocked"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_moderation_hook_with_sanitization():
|
||||
"""Test async_moderation_hook with content sanitization enabled"""
|
||||
mock_user_api_key_dict = UserAPIKeyAuth()
|
||||
|
||||
guardrail = ModelArmorGuardrail(
|
||||
template_id="test-template",
|
||||
project_id="test-project",
|
||||
location="us-central1",
|
||||
guardrail_name="model-armor-test",
|
||||
mask_request_content=True, # Enable sanitization
|
||||
)
|
||||
|
||||
# Mock response with sanitized content
|
||||
mock_response = AsyncMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json = AsyncMock(return_value={
|
||||
"sanitizationResult": {
|
||||
"filterMatchState": "MATCH_FOUND",
|
||||
"filterResults": {
|
||||
"sdp": {
|
||||
"sdpFilterResult": {
|
||||
"deidentifyResult": {
|
||||
"matchState": "MATCH_FOUND",
|
||||
"data": {
|
||||
"text": "Hello, my phone number is [REDACTED]"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
# Mock the access token method and async handler
|
||||
guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project"))
|
||||
guardrail.async_handler = AsyncMock()
|
||||
guardrail.async_handler.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
original_content = "Hello, my phone number is 555-123-4567"
|
||||
request_data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [
|
||||
{"role": "user", "content": original_content}
|
||||
],
|
||||
"metadata": {"guardrails": ["model-armor-test"]}
|
||||
}
|
||||
|
||||
result = await guardrail.async_moderation_hook(
|
||||
data=request_data,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
call_type="completion"
|
||||
)
|
||||
|
||||
# Should return data with sanitized content
|
||||
assert result == request_data
|
||||
# Content should be sanitized
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import get_last_user_message
|
||||
sanitized_content = get_last_user_message(request_data["messages"])
|
||||
assert sanitized_content == "Hello, my phone number is [REDACTED]"
|
||||
assert sanitized_content != original_content
|
||||
# Should have metadata added
|
||||
assert "_model_armor_response" in request_data["metadata"]
|
||||
assert request_data["metadata"]["_model_armor_status"] == "success"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_moderation_hook_no_user_messages():
|
||||
"""Test async_moderation_hook when there are no user messages to check"""
|
||||
mock_user_api_key_dict = UserAPIKeyAuth()
|
||||
|
||||
guardrail = ModelArmorGuardrail(
|
||||
template_id="test-template",
|
||||
project_id="test-project",
|
||||
location="us-central1",
|
||||
guardrail_name="model-armor-test",
|
||||
)
|
||||
|
||||
request_data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant"},
|
||||
{"role": "assistant", "content": "How can I help you?"}
|
||||
],
|
||||
"metadata": {"guardrails": ["model-armor-test"]}
|
||||
}
|
||||
|
||||
result = await guardrail.async_moderation_hook(
|
||||
data=request_data,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
call_type="completion"
|
||||
)
|
||||
|
||||
# Should return the original data unchanged since no user messages to check
|
||||
assert result == request_data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_moderation_hook_should_not_run():
|
||||
"""Test async_moderation_hook when guardrail should not run due to missing guardrail name"""
|
||||
try:
|
||||
import google.auth
|
||||
except ImportError:
|
||||
pytest.skip("google.auth not installed")
|
||||
return
|
||||
|
||||
mock_user_api_key_dict = UserAPIKeyAuth()
|
||||
|
||||
guardrail = ModelArmorGuardrail(
|
||||
template_id="test-template",
|
||||
project_id="test-project",
|
||||
location="us-central1",
|
||||
guardrail_name="different-guardrail-name", # Different name than what's in metadata
|
||||
)
|
||||
|
||||
# Request data with a different guardrail name
|
||||
request_data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello, how are you?"}
|
||||
],
|
||||
"metadata": {"guardrails": ["some-other-guardrail"]} # Different guardrail name
|
||||
}
|
||||
|
||||
result = await guardrail.async_moderation_hook(
|
||||
data=request_data,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
call_type="completion"
|
||||
)
|
||||
|
||||
# Should return the original data unchanged since guardrail name doesn't match
|
||||
assert result == request_data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_moderation_hook_api_error_fail_on_error_true():
|
||||
"""Test async_moderation_hook when API call fails and fail_on_error is True"""
|
||||
mock_user_api_key_dict = UserAPIKeyAuth()
|
||||
|
||||
guardrail = ModelArmorGuardrail(
|
||||
template_id="test-template",
|
||||
project_id="test-project",
|
||||
location="us-central1",
|
||||
guardrail_name="model-armor-test",
|
||||
optional_params={"fail_on_error": True}
|
||||
)
|
||||
|
||||
# Mock the access token method
|
||||
guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project"))
|
||||
|
||||
# Mock the async handler to raise an exception
|
||||
guardrail.async_handler = AsyncMock()
|
||||
guardrail.async_handler.post = AsyncMock(side_effect=Exception("API Error"))
|
||||
|
||||
request_data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello, how are you?"}
|
||||
],
|
||||
"metadata": {"guardrails": ["model-armor-test"]}
|
||||
}
|
||||
|
||||
# Should raise the exception since fail_on_error is True
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
await guardrail.async_moderation_hook(
|
||||
data=request_data,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
call_type="completion"
|
||||
)
|
||||
|
||||
assert "API Error" in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_moderation_hook_api_error_fail_on_error_false():
|
||||
"""Test async_moderation_hook when API call fails and fail_on_error is False"""
|
||||
mock_user_api_key_dict = UserAPIKeyAuth()
|
||||
|
||||
guardrail = ModelArmorGuardrail(
|
||||
template_id="test-template",
|
||||
project_id="test-project",
|
||||
location="us-central1",
|
||||
guardrail_name="model-armor-test",
|
||||
optional_params={"fail_on_error": False}
|
||||
)
|
||||
|
||||
# Mock the access token method
|
||||
guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project"))
|
||||
|
||||
# Mock the async handler to raise an exception
|
||||
guardrail.async_handler = AsyncMock()
|
||||
guardrail.async_handler.post = AsyncMock(side_effect=Exception("API Error"))
|
||||
|
||||
request_data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello, how are you?"}
|
||||
],
|
||||
"metadata": {"guardrails": ["model-armor-test"]}
|
||||
}
|
||||
|
||||
# Even with fail_on_error=False, the decorator may still raise the exception
|
||||
# This test verifies that the exception is properly logged and handled
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
await guardrail.async_moderation_hook(
|
||||
data=request_data,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
call_type="completion"
|
||||
)
|
||||
|
||||
assert "API Error" in str(exc_info.value)
|
||||
Loading…
Add table
Reference in a new issue