tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-18 04:28:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Fix] LakeraAI v2 Guardrail - Ensure exception is raised correctly (#14867)

Browse Source 
* fix exception lakera

* test lakera ai v2

* ruff fix
This commit is contained in:
Ishaan Jaff
2025-09-24 10:55:56 -07:00
committed by GitHub
parent c6cb36186c
commit c69bac991b
3 changed files with 226 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py
+35 -12
View File
@@ -3,9 +3,10 @@ import os
from datetime import datetime
from typing import Dict, List, Literal, Optional, Tuple, Union
from fastapi import HTTPException
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.exceptions import GuardrailRaisedException
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
@@ -237,9 +238,8 @@ class LakeraAIGuardrail(CustomGuardrail):
)
else:
# If there are other violations or not set to mask PII, raise exception
raise GuardrailRaisedException(
guardrail_name=self.guardrail_name,
message="Lakera AI flagged this request. Please review the request and try again.",
raise self._get_http_exception_for_blocked_guardrail(
lakera_guardrail_response
)
#########################################################
@@ -304,9 +304,8 @@ class LakeraAIGuardrail(CustomGuardrail):
)
else:
# If there are other violations or not set to mask PII, raise exception
raise GuardrailRaisedException(
guardrail_name=self.guardrail_name,
message="Lakera AI flagged this request. Please review the request and try again.",
raise self._get_http_exception_for_blocked_guardrail(
lakera_guardrail_response
)
#########################################################
@@ -327,8 +326,32 @@ class LakeraAIGuardrail(CustomGuardrail):
if not lakera_response:
return False
for item in lakera_response.get("payload", []) or []:
detector_type = item.get("detector_type", "") or ""
if not detector_type.startswith("pii/"):
return False
return True
# Check breakdown field for detected violations
breakdown = lakera_response.get("breakdown", []) or []
if not breakdown:
return False
has_violations = False
for item in breakdown:
if item.get("detected", False):
has_violations = True
detector_type = item.get("detector_type", "") or ""
if not detector_type.startswith("pii/"):
return False
# Return True only if there are violations and they are all PII
return has_violations
def _get_http_exception_for_blocked_guardrail(
self, lakera_response: Optional[LakeraAIResponse]
) -> HTTPException:
"""
Get the HTTP exception for a blocked guardrail, similar to Bedrock's implementation.
"""
return HTTPException(
status_code=400,
detail={
"error": "Violated guardrail policy",
"lakera_guardrail_response": lakera_response,
},
)
litellm/proxy/proxy_config.yaml
+13
View File
@@ -23,3 +23,16 @@ model_list:
litellm_params:
model: gemini/*
api_key: os.environ/GEMINI_API_KEY
guardrails:
- guardrail_name: lakera
litellm_params:
guardrail: lakera_v2
mode: pre_call
api_key: os.environ/LAKERA_API_KEY
default_on: false
project_id: project-9770817088
breakdown: true
payload: true
dev_info: true
tests/guardrails_tests/test_lakera_v2.py
+178 -1
View File
@@ -11,7 +11,8 @@ from litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import LakeraAIGuardr
from litellm.types.guardrails import PiiEntityType, PiiAction
from litellm.proxy._types import UserAPIKeyAuth
from litellm.caching.caching import DualCache
from litellm.exceptions import BlockedPiiEntityError
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
from fastapi import HTTPException
from litellm.types.utils import CallTypes as LitellmCallTypes
@@ -54,3 +55,179 @@ async def test_lakera_pre_call_hook_for_pii_masking():
assert "4111-1111-1111-1111" not in user_message
assert "test@example.com" not in user_message
@pytest.mark.asyncio
async def test_lakera_blocks_non_pii_violations():
"""Test that Lakera guardrail blocks requests with non-PII violations like hate speech, violence, etc."""
lakera_guardrail = LakeraAIGuardrail(
api_key="test_key",
)
# Mock the call_v2_guard method to return a response similar to the user's example
mock_response = {
'payload': [],
'flagged': True,
'dev_info': {'git_revision': 'f0bc093a', 'git_timestamp': '2025-09-23T15:28:06+00:00', 'model_version': 'lakera-guard-1', 'version': '2.0.281'},
'metadata': {'request_uuid': 'b7cd4c8a-28aa-4285-a245-2befee514dbf'},
'breakdown': [
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-moderated-content', 'detector_type': 'moderated_content/crime', 'detected': True, 'message_id': 0},
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-moderated-content', 'detector_type': 'moderated_content/hate', 'detected': True, 'message_id': 0},
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-moderated-content', 'detector_type': 'moderated_content/violence', 'detected': True, 'message_id': 0},
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-prompt-attack', 'detector_type': 'prompt_attack', 'detected': True, 'message_id': 0},
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-pii', 'detector_type': 'pii/email', 'detected': False, 'message_id': 0},
]
}
with patch.object(lakera_guardrail, 'call_v2_guard', new_callable=AsyncMock) as mock_call:
mock_call.return_value = (mock_response, {})
# Create a sample request that would trigger violations
data = {
"messages": [
{"role": "user", "content": "Some harmful content that triggers violations"}
],
"model": "gpt-3.5-turbo",
"metadata": {}
}
# Mock objects needed for the pre-call hook
user_api_key_dict = UserAPIKeyAuth(api_key="test_key")
cache = DualCache()
# The guardrail should raise an HTTPException for non-PII violations
with pytest.raises(HTTPException) as exc_info:
await lakera_guardrail.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=cache,
data=data,
call_type="completion"
)
# Verify the exception details include the Lakera response
assert exc_info.value.status_code == 400
assert "Violated guardrail policy" in str(exc_info.value.detail)
assert "lakera_guardrail_response" in exc_info.value.detail
@pytest.mark.asyncio
async def test_lakera_only_pii_violations_are_masked():
"""Test that Lakera guardrail only masks PII violations and doesn't block the request."""
lakera_guardrail = LakeraAIGuardrail(
api_key="test_key",
)
# Mock response with only PII violations
mock_response = {
'payload': [
{'detector_type': 'pii/email', 'start': 10, 'end': 25, 'message_id': 0}
],
'flagged': True,
'breakdown': [
{'project_id': 'project-9770817088', 'detector_type': 'pii/email', 'detected': True, 'message_id': 0},
{'project_id': 'project-9770817088', 'detector_type': 'moderated_content/hate', 'detected': False, 'message_id': 0},
{'project_id': 'project-9770817088', 'detector_type': 'prompt_attack', 'detected': False, 'message_id': 0},
]
}
with patch.object(lakera_guardrail, 'call_v2_guard', new_callable=AsyncMock) as mock_call:
mock_call.return_value = (mock_response, {})
data = {
"messages": [
{"role": "user", "content": "My email test@example.com here"}
],
"model": "gpt-3.5-turbo",
"metadata": {}
}
user_api_key_dict = UserAPIKeyAuth(api_key="test_key")
cache = DualCache()
# Should not raise an exception, just mask the PII
result = await lakera_guardrail.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=cache,
data=data,
call_type="completion"
)
# Verify the request was not blocked
assert result is not None
assert "messages" in result
@pytest.mark.asyncio
async def test_lakera_blocks_flagged_content_with_user_scenario():
"""
Test the exact user scenario where Lakera flagged content but request went through.
This should now be blocked with the fix to check breakdown field instead of payload.
"""
lakera_guardrail = LakeraAIGuardrail(
api_key="test_key",
)
# Mock response matching the exact user scenario
mock_response = {
'payload': [], # Empty payload like in user's case
'flagged': True,
'dev_info': {'git_revision': 'f0bc093a', 'git_timestamp': '2025-09-23T15:28:06+00:00', 'model_version': 'lakera-guard-1', 'version': '2.0.281'},
'metadata': {'request_uuid': 'b7cd4c8a-28aa-4285-a245-2befee514dbf'},
'breakdown': [
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-moderated-content', 'detector_type': 'moderated_content/crime', 'detected': True, 'message_id': 0},
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-moderated-content', 'detector_type': 'moderated_content/hate', 'detected': True, 'message_id': 0},
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-moderated-content', 'detector_type': 'moderated_content/profanity', 'detected': False, 'message_id': 0},
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-moderated-content', 'detector_type': 'moderated_content/sexual', 'detected': False, 'message_id': 0},
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-moderated-content', 'detector_type': 'moderated_content/violence', 'detected': True, 'message_id': 0},
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-moderated-content', 'detector_type': 'moderated_content/weapons', 'detected': True, 'message_id': 0},
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-pii', 'detector_type': 'pii/address', 'detected': False, 'message_id': 0},
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-pii', 'detector_type': 'pii/credit_card', 'detected': False, 'message_id': 0},
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-pii', 'detector_type': 'pii/email', 'detected': False, 'message_id': 0},
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-pii', 'detector_type': 'pii/iban_code', 'detected': False, 'message_id': 0},
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-pii', 'detector_type': 'pii/ip_address', 'detected': False, 'message_id': 0},
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-pii', 'detector_type': 'pii/name', 'detected': False, 'message_id': 0},
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-pii', 'detector_type': 'pii/phone_number', 'detected': False, 'message_id': 0},
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-pii', 'detector_type': 'pii/us_social_security_number', 'detected': False, 'message_id': 0},
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-prompt-attack', 'detector_type': 'prompt_attack', 'detected': True, 'message_id': 0},
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-unknown-links', 'detector_type': 'unknown_links', 'detected': False, 'message_id': 0}
]
}
with patch.object(lakera_guardrail, 'call_v2_guard', new_callable=AsyncMock) as mock_call:
mock_call.return_value = (mock_response, {})
# Create a sample request that would trigger violations
data = {
"messages": [
{"role": "user", "content": "Some harmful content that should be blocked"}
],
"model": "gpt-3.5-turbo",
"metadata": {}
}
# Mock objects needed for the pre-call hook
user_api_key_dict = UserAPIKeyAuth(api_key="test_key")
cache = DualCache()
# With the fix, this should now raise an HTTPException instead of letting the request through
with pytest.raises(HTTPException) as exc_info:
await lakera_guardrail.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=cache,
data=data,
call_type="completion"
)
# Verify the exception details
assert exc_info.value.status_code == 400
assert "Violated guardrail policy" in str(exc_info.value.detail)
assert "lakera_guardrail_response" in exc_info.value.detail
# Verify the full response is included in the exception
lakera_response = exc_info.value.detail["lakera_guardrail_response"]
assert lakera_response["flagged"] is True
assert lakera_response["metadata"]["request_uuid"] == "b7cd4c8a-28aa-4285-a245-2befee514dbf"
assert len(lakera_response["breakdown"]) == 16 # All the breakdown items from the user's scenario