import io, asyncio import pytest import time from litellm import mock_completion from unittest.mock import MagicMock, AsyncMock, patch import litellm from litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import LakeraAIGuardrail from litellm.types.guardrails import PiiEntityType, PiiAction from litellm.proxy._types import UserAPIKeyAuth from litellm.caching.caching import DualCache from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from fastapi import HTTPException from litellm.types.utils import CallTypes as LitellmCallTypes, ModelResponse @pytest.mark.asyncio async def test_lakera_pre_call_hook_for_pii_masking(): """Test for Lakera guardrail pre-call hook for PII masking""" # Setup the guardrail with specific entities config litellm._turn_on_debug() lakera_guardrail = LakeraAIGuardrail( api_key="test_key", ) # Mock response with PII detections in payload (with start/end positions for masking) mock_response = { "payload": [ { "detector_type": "pii/credit_card", "start": 18, "end": 37, "message_id": 1, }, # "4111-1111-1111-1111" { "detector_type": "pii/email", "start": 54, "end": 70, "message_id": 1, }, # "test@example.com" ], "flagged": True, "breakdown": [ {"detector_type": "pii/credit_card", "detected": True, "message_id": 1}, {"detector_type": "pii/email", "detected": True, "message_id": 1}, ], } with patch.object( lakera_guardrail, "call_v2_guard", new_callable=AsyncMock ) as mock_call: mock_call.return_value = (mock_response, {}) # Create a sample request with PII data data = { "messages": [ {"role": "system", "content": "You are a helpful assistant."}, { "role": "user", "content": "My credit card is 4111-1111-1111-1111 and my email is test@example.com. My phone number is 555-123-4567", }, ], "model": "gpt-5-mini", "metadata": {}, } # Mock objects needed for the pre-call hook user_api_key_dict = UserAPIKeyAuth(api_key="test_key") cache = DualCache() # Call the pre-call hook with the specified call type modified_data = await lakera_guardrail.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="completion", ) print(modified_data) # Verify the messages have been modified to mask PII assert ( modified_data["messages"][0]["content"] == "You are a helpful assistant." ) # System prompt should be unchanged user_message = modified_data["messages"][1]["content"] # Verify both credit card and email are masked assert "4111-1111-1111-1111" not in user_message assert "test@example.com" not in user_message # Verify masking placeholders are present assert "[MASKED CREDIT_CARD]" in user_message assert "[MASKED EMAIL]" 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-5-mini", "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-5-mini", "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-5-mini", "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 @pytest.mark.asyncio async def test_lakera_monitor_mode_allows_flagged_content(): """Test that monitor mode logs violations but allows requests to proceed.""" lakera_guardrail = LakeraAIGuardrail( api_key="test_key", on_flagged="monitor", # Monitor mode ) # Mock response with violations mock_response = { "payload": [], "flagged": True, "breakdown": [ { "detector_type": "moderated_content/violence", "detected": True, "message_id": 0, }, {"detector_type": "prompt_attack", "detected": True, "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": "Some harmful content"}], "model": "gpt-5-mini", "metadata": {}, } user_api_key_dict = UserAPIKeyAuth(api_key="test_key") cache = DualCache() # Should NOT raise an exception in monitor mode result = await lakera_guardrail.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="completion", ) # Verify request was allowed through assert result is not None assert "messages" in result @pytest.mark.asyncio async def test_lakera_block_mode_raises_exception(): """Test that block mode (default) raises HTTPException for violations.""" lakera_guardrail = LakeraAIGuardrail( api_key="test_key", on_flagged="block", # Block mode (default) ) mock_response = { "payload": [], "flagged": True, "breakdown": [ { "detector_type": "moderated_content/violence", "detected": True, "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": "Harmful content"}], "model": "gpt-5-mini", "metadata": {}, } user_api_key_dict = UserAPIKeyAuth(api_key="test_key") cache = DualCache() # Should raise HTTPException in block mode 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", ) assert exc_info.value.status_code == 400 @pytest.mark.asyncio async def test_lakera_monitor_mode_during_call(): """Test monitor mode works with during_call (moderation_hook).""" lakera_guardrail = LakeraAIGuardrail( api_key="test_key", on_flagged="monitor", ) mock_response = { "payload": [], "flagged": True, "breakdown": [ {"detector_type": "prompt_attack", "detected": True, "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": "Test content"}], "model": "gpt-5-mini", "metadata": {}, } user_api_key_dict = UserAPIKeyAuth(api_key="test_key") # Should NOT raise exception in monitor mode result = await lakera_guardrail.async_moderation_hook( data=data, user_api_key_dict=user_api_key_dict, call_type="completion" ) assert result is not None @pytest.mark.asyncio async def test_lakera_post_call_blocks_flagged_content(): """Post-call hook should block when violations are flagged.""" lakera_guardrail = LakeraAIGuardrail(api_key="test_key") mock_response = { "payload": [], "flagged": True, "breakdown": [ { "detector_type": "moderated_content/violence", "detected": True, "message_id": 0, }, ], } # Mock LLM response object llm_response = MagicMock() llm_response.model_dump.return_value = { "choices": [{"message": {"role": "assistant", "content": "some response"}}] } 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": "Harmful content"}], "model": "gpt-5-mini", "metadata": {}, } user_api_key_dict = UserAPIKeyAuth(api_key="test_key") with pytest.raises(HTTPException) as exc_info: await lakera_guardrail.async_post_call_success_hook( data=data, user_api_key_dict=user_api_key_dict, response=llm_response, ) assert exc_info.value.status_code == 400 @pytest.mark.asyncio async def test_lakera_post_call_allows_clean_content(): """Post-call hook should allow when not flagged.""" lakera_guardrail = LakeraAIGuardrail(api_key="test_key") mock_response = { "payload": [], "flagged": False, "breakdown": [], } llm_response = MagicMock() llm_response.model_dump.return_value = { "choices": [{"message": {"role": "assistant", "content": "clean response"}}] } 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": "Hello"}], "model": "gpt-5-mini", "metadata": {}, } user_api_key_dict = UserAPIKeyAuth(api_key="test_key") result = await lakera_guardrail.async_post_call_success_hook( data=data, user_api_key_dict=user_api_key_dict, response=llm_response, ) assert result is llm_response @pytest.mark.asyncio async def test_lakera_post_call_masks_pii_and_allows(): """Post-call hook should mask PII-only violations and allow response.""" lakera_guardrail = LakeraAIGuardrail(api_key="test_key") mock_response = { "payload": [ {"detector_type": "pii/email", "start": 11, "end": 26, "message_id": 1} ], "flagged": True, "breakdown": [ {"detector_type": "pii/email", "detected": True, "message_id": 1}, ], } llm_response = MagicMock() llm_response.model_dump.return_value = { "choices": [ { "message": { "role": "assistant", "content": "Your email is test@example.com", } }, ] } 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": "Hello"}], "model": "gpt-5-mini", "metadata": {}, } user_api_key_dict = UserAPIKeyAuth(api_key="test_key") result = await lakera_guardrail.async_post_call_success_hook( data=data, user_api_key_dict=user_api_key_dict, response=llm_response, ) assert isinstance( result, ModelResponse ), "PII masking path must return ModelResponse" result_dict = result.model_dump() assert ( result_dict["choices"][0]["message"]["content"] != "Your email is test@example.com" ) assert "[MASKED" in result_dict["choices"][0]["message"]["content"]