mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
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>
This commit is contained in:
Steve G
Sameer Kankute
parent
a03d2a308d
commit
68512bdc05
4 changed files with 295 additions and 2 deletions
|
|
@ -4,6 +4,8 @@ import TabItem from '@theme/TabItem';
|
|||
|
||||
# Lakera AI
|
||||
|
||||
**Supported endpoints:** The Lakera v2 integration only supports the **chat completions** endpoint (`/v1/chat/completions`). It is not supported for the Responses API, `/v1/messages`, MCP, A2A, or other proxy endpoints.
|
||||
|
||||
## Quick Start
|
||||
### 1. Define Guardrails on your LiteLLM config.yaml
|
||||
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import (
|
|||
LakeraAIRequest,
|
||||
LakeraAIResponse,
|
||||
)
|
||||
from litellm.types.utils import CallTypesLiteral, GuardrailStatus
|
||||
from litellm.types.utils import CallTypesLiteral, GuardrailStatus, ModelResponse
|
||||
|
||||
|
||||
class LakeraAIGuardrail(CustomGuardrail):
|
||||
|
|
@ -39,6 +39,9 @@ class LakeraAIGuardrail(CustomGuardrail):
|
|||
"""
|
||||
Initialize the LakeraAIGuardrail class.
|
||||
|
||||
This guardrail only supports the chat completions endpoint (/v1/chat/completions).
|
||||
It is not supported for the Responses API, /v1/messages, MCP, A2A, or other endpoints.
|
||||
|
||||
This calls: https://api.lakera.ai/v2/guard
|
||||
|
||||
Args:
|
||||
|
|
@ -146,6 +149,7 @@ class LakeraAIGuardrail(CustomGuardrail):
|
|||
if not payload:
|
||||
return messages
|
||||
|
||||
messages = copy.deepcopy(messages)
|
||||
# For each message, find its detections on the fly
|
||||
for idx, msg in enumerate(messages):
|
||||
content = msg.get("content", "")
|
||||
|
|
@ -161,6 +165,13 @@ class LakeraAIGuardrail(CustomGuardrail):
|
|||
if not detected_modifications:
|
||||
continue
|
||||
|
||||
# Apply masks from end to start so earlier indices remain valid after each replacement
|
||||
detected_modifications = sorted(
|
||||
detected_modifications,
|
||||
key=lambda d: (d.get("start", 0), d.get("end", 0)),
|
||||
reverse=True,
|
||||
)
|
||||
|
||||
for modification in detected_modifications:
|
||||
start, end = modification.get("start", 0), modification.get("end", 0)
|
||||
|
||||
|
|
@ -321,6 +332,92 @@ class LakeraAIGuardrail(CustomGuardrail):
|
|||
|
||||
return data
|
||||
|
||||
async def async_post_call_success_hook(
|
||||
self,
|
||||
data: dict,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
response,
|
||||
):
|
||||
"""
|
||||
Post-call hook for Lakera guardrail.
|
||||
"""
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
add_guardrail_to_applied_guardrails_header,
|
||||
)
|
||||
|
||||
event_type: GuardrailEventHooks = GuardrailEventHooks.post_call
|
||||
if self.should_run_guardrail(data=data, event_type=event_type) is not True:
|
||||
return response
|
||||
|
||||
original_messages: Optional[List[AllMessageValues]] = data.get("messages", [])
|
||||
if original_messages is None:
|
||||
original_messages = []
|
||||
|
||||
# Extract assistant messages from the response, keeping only role/content.
|
||||
# Track choice indices so we write masked content back to the correct choice
|
||||
# when some choices have null content (e.g. tool-call-only).
|
||||
response_messages: List[AllMessageValues] = []
|
||||
choice_indices: List[int] = []
|
||||
response_dict = (
|
||||
response.model_dump() if hasattr(response, "model_dump") else {}
|
||||
)
|
||||
for i, choice in enumerate(response_dict.get("choices", [])):
|
||||
msg = choice.get("message")
|
||||
if not msg:
|
||||
continue
|
||||
role = msg.get("role")
|
||||
content = msg.get("content")
|
||||
if role and content:
|
||||
response_messages.append({"role": role, "content": content})
|
||||
choice_indices.append(i)
|
||||
|
||||
# Use a copy of original_messages so _mask_pii_in_messages does not mutate data["messages"]
|
||||
post_call_messages = copy.deepcopy(original_messages) + response_messages
|
||||
|
||||
# Call Lakera guardrail
|
||||
lakera_guardrail_response, _ = await self.call_v2_guard(
|
||||
messages=post_call_messages,
|
||||
request_data=data,
|
||||
event_type=GuardrailEventHooks.post_call,
|
||||
)
|
||||
|
||||
# Handle flagged content
|
||||
if lakera_guardrail_response.get("flagged") is True:
|
||||
# If only PII violations exist, mask the PII in the response and allow
|
||||
if self._is_only_pii_violation(lakera_guardrail_response):
|
||||
masked_entity_count: Dict[str, int] = {}
|
||||
masked_messages = self._mask_pii_in_messages(
|
||||
messages=post_call_messages,
|
||||
lakera_response=lakera_guardrail_response,
|
||||
masked_entity_count=masked_entity_count,
|
||||
)
|
||||
assistant_messages = masked_messages[len(original_messages) :]
|
||||
for idx, msg in enumerate(assistant_messages):
|
||||
if idx < len(choice_indices):
|
||||
choice_idx = choice_indices[idx]
|
||||
response_dict["choices"][choice_idx]["message"]["content"] = msg.get("content", "")
|
||||
add_guardrail_to_applied_guardrails_header(
|
||||
request_data=data, guardrail_name=self.guardrail_name
|
||||
)
|
||||
return ModelResponse(**response_dict)
|
||||
|
||||
if self.on_flagged == "monitor":
|
||||
verbose_proxy_logger.warning(
|
||||
"Lakera Guardrail: Post-call violation detected in monitor mode"
|
||||
)
|
||||
# Allow response to proceed
|
||||
elif self.on_flagged == "block":
|
||||
raise self._get_http_exception_for_blocked_guardrail(
|
||||
lakera_guardrail_response
|
||||
)
|
||||
|
||||
# Record applied guardrail
|
||||
add_guardrail_to_applied_guardrails_header(
|
||||
request_data=data, guardrail_name=self.guardrail_name
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
def _is_only_pii_violation(
|
||||
self, lakera_response: Optional[LakeraAIResponse]
|
||||
) -> bool:
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ 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
|
||||
from litellm.types.utils import CallTypes as LitellmCallTypes, ModelResponse
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -380,3 +380,131 @@ async def test_lakera_monitor_mode_during_call():
|
|||
|
||||
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"]
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,66 @@
|
|||
"""
|
||||
Tests for Lakera AI v2 guardrail hook (post-call and shared behavior).
|
||||
|
||||
PR checklist requires at least one test in tests/test_litellm/.
|
||||
Additional tests live in tests/guardrails_tests/test_lakera_v2.py.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import LakeraAIGuardrail
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lakera_post_call_success_hook_returns_model_response_when_pii_masked():
|
||||
"""
|
||||
Post-call hook must return a ModelResponse (not a dict) when PII is masked,
|
||||
so the parent async_post_call_success_deployment_hook accepts it via _is_valid_response_type.
|
||||
"""
|
||||
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
|
||||
), "Must return ModelResponse so deployment hook does not discard masked response"
|
||||
result_dict = result.model_dump()
|
||||
assert "[MASKED" in result_dict["choices"][0]["message"]["content"]
|
||||
assert "test@example.com" not in result_dict["choices"][0]["message"]["content"]
|
||||
Loading…
Add table
Reference in a new issue