scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 8
litellm/tests/guardrails_tests/test_lakera_v2.py
yuneng-jiang 6a0d03914c
test: drop the cwd-relative sys.path.insert calls from the test suite (#37802) 
* test: drop the cwd-relative sys.path.insert calls from the test suite

TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.

Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.

Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.

* test: drop the duplicate imports the sys.path sweep exposed to F811

* test(pre-call-utils): restore the os import the new bedrock tests need
2026-08-22 09:25:58 -07:00

733 lines
24 KiB
Python
Raw Permalink Blame History

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"]
Reference in a new issue View git blame Copy permalink