From 68512bdc054483700048129cfb0936d9400dabb4 Mon Sep 17 00:00:00 2001 From: Steve G <9045562+eurogig@users.noreply.github.com> Date: Wed, 25 Feb 2026 20:20:38 -0500 Subject: [PATCH] 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 * 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 Co-authored-by: Cursor Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../docs/proxy/guardrails/lakera_ai.md | 2 + .../guardrail_hooks/lakera_ai_v2.py | 99 ++++++++++++- tests/guardrails_tests/test_lakera_v2.py | 130 +++++++++++++++++- .../guardrail_hooks/test_lakera_ai_v2.py | 66 +++++++++ 4 files changed, 295 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py diff --git a/docs/my-website/docs/proxy/guardrails/lakera_ai.md b/docs/my-website/docs/proxy/guardrails/lakera_ai.md index 7aacc3fa924..cd27dd23618 100644 --- a/docs/my-website/docs/proxy/guardrails/lakera_ai.md +++ b/docs/my-website/docs/proxy/guardrails/lakera_ai.md @@ -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 diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 738827b7ada..dbda524ca04 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -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: diff --git a/tests/guardrails_tests/test_lakera_v2.py b/tests/guardrails_tests/test_lakera_v2.py index 2a8731d5ecd..aad9929809c 100644 --- a/tests/guardrails_tests/test_lakera_v2.py +++ b/tests/guardrails_tests/test_lakera_v2.py @@ -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"] + diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py new file mode 100644 index 00000000000..f6e7b7841e2 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py @@ -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"]