From c69bac991be9f8690cd35d11d9bb95c020a37bb3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 24 Sep 2025 10:55:56 -0700 Subject: [PATCH] [Fix] LakeraAI v2 Guardrail - Ensure exception is raised correctly (#14867) * fix exception lakera * test lakera ai v2 * ruff fix --- .../guardrail_hooks/lakera_ai_v2.py | 47 +++-- litellm/proxy/proxy_config.yaml | 13 ++ tests/guardrails_tests/test_lakera_v2.py | 179 +++++++++++++++++- 3 files changed, 226 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index e167e73ac4..b65664e00b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -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, + }, + ) diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index addf2443d3..921b564407 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -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 diff --git a/tests/guardrails_tests/test_lakera_v2.py b/tests/guardrails_tests/test_lakera_v2.py index 19c4424bee..f3b2795a27 100644 --- a/tests/guardrails_tests/test_lakera_v2.py +++ b/tests/guardrails_tests/test_lakera_v2.py @@ -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 +