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
Steve G 9806e21871
Add Lakera v2 post-call hook and tests (fixed PII masking) (#21783) 
* Add post-call hook for Lakera guardrail and mask PII in responses

* Add post-call hook for Lakera and mask PII in responses

* Fix post-call hook: pass event_type to call_v2_guard

* Address Greptile review: return ModelResponse, fix mutation, add header, test location, mask order

- PII masking path: return ModelResponse instead of dict so deployment hook accepts it
- Avoid mutating request data: deep copy original_messages and messages in _mask_pii_in_messages
- Add guardrail header in PII-only return path
- Add test in tests/test_litellm/ (test_lakera_ai_v2.py) per PR checklist
- Sort PII payload spans by (start,end) descending so multiple spans in one message mask correctly

Co-authored-by: Cursor <cursoragent@cursor.com>

* Updated ponteital for index mismatch when choices have null content and inconsistent on_flagged access pattern

* Update litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* Update to explicitly state supported endpoints - chat completions

* Fix minor lint error on masked_entity_count

---------

Co-authored-by: Steve <steve.giguere@lakera.ai>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
2026-02-25 17:20:38 -08:00

510 lines
21 KiB
Python
Raw Permalink Blame History

import sys
import os
import io, asyncio
import pytest
import time
from litellm import mock_completion
from unittest.mock import MagicMock, AsyncMock, patch
sys.path.insert(0, os.path.abspath("../.."))
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-3.5-turbo",
"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-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
@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-3.5-turbo",
"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-3.5-turbo",
"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-3.5-turbo",
"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-3.5-turbo",
"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-3.5-turbo",
"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-3.5-turbo",
"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