[Feat] Add Lakera v2 Guardrail Support (#10880)

* feat: add types for lakera v2

* feat: add call v2 guard for lakera

* feat: add call v2 guard for lakera

* feat: add LAKERA_V2

* feat: add LAKERA_V2 params

* feat: add initialize_lakera_v2

* fix: lakera pii masking

* test: lakera pii masking

* fix: lakera pii masking with tracing

* fix: lakera pii masking with tracing

* fix: fix linting errors

* fix: lakera ai docs

docs/my-website/docs/proxy/guardrails/lakera_ai.md

@@ -8,7 +8,8 @@ import TabItem from '@theme/TabItem';
 ### 1. Define Guardrails on your LiteLLM config.yaml 
 
 Define your guardrails under the `guardrails` section
-```yaml
+
+```yaml showLineNumbers title="litellm config.yaml"
 model_list:
   - model_name: gpt-3.5-turbo
     litellm_params:
@@ -18,13 +19,13 @@ model_list:
 guardrails:
   - guardrail_name: "lakera-guard"
     litellm_params:
-      guardrail: lakera  # supported values: "aporia", "bedrock", "lakera"
+      guardrail: lakera_v2  # supported values: "aporia", "bedrock", "lakera"
       mode: "during_call"
       api_key: os.environ/LAKERA_API_KEY
       api_base: os.environ/LAKERA_API_BASE
   - guardrail_name: "lakera-pre-guard"
     litellm_params:
-      guardrail: lakera  # supported values: "aporia", "bedrock", "lakera"
+      guardrail: lakera_v2  # supported values: "aporia", "bedrock", "lakera"
       mode: "pre_call"
       api_key: os.environ/LAKERA_API_KEY
       api_base: os.environ/LAKERA_API_BASE
@@ -53,7 +54,7 @@ litellm --config config.yaml --detailed_debug
 
 Expect this to fail since since `ishaan@berri.ai` in the request is PII
 
-```shell
+```shell showLineNumbers title="Curl Request"
 curl -i http://localhost:4000/v1/chat/completions \
   -H "Content-Type: application/json" \
   -H "Authorization: Bearer sk-npnwjPQciVRok5yNZgKmFQ" \
@@ -108,7 +109,7 @@ Expected response on failure
 
 <TabItem label="Successful Call " value = "allowed">
 
-```shell
+```shell showLineNumbers title="Curl Request"
 curl -i http://localhost:4000/v1/chat/completions \
   -H "Content-Type: application/json" \
   -H "Authorization: Bearer sk-npnwjPQciVRok5yNZgKmFQ" \
@@ -125,31 +126,3 @@ curl -i http://localhost:4000/v1/chat/completions \
 
 
 </Tabs>
-
-## Advanced 
-### Set category-based thresholds.
-
-Lakera has 2 categories for prompt_injection attacks:
-- jailbreak
-- prompt_injection
-
-```yaml
-model_list:
-  - model_name: fake-openai-endpoint
-    litellm_params:
-      model: openai/fake
-      api_key: fake-key
-      api_base: https://exampleopenaiendpoint-production.up.railway.app/
-
-guardrails:
-  - guardrail_name: "lakera-guard"
-    litellm_params:
-      guardrail: lakera  # supported values: "aporia", "bedrock", "lakera"
-      mode: "during_call"
-      api_key: os.environ/LAKERA_API_KEY
-      api_base: os.environ/LAKERA_API_BASE
-      category_thresholds:
-        prompt_injection: 0.1
-        jailbreak: 0.1
-  
-```

litellm/exceptions.py

@@ -809,6 +809,13 @@ class LiteLLMUnknownProvider(BadRequestError):
         return self.message
 
 
+class GuardrailRaisedException(Exception):
+    def __init__(self, guardrail_name: Optional[str] = None, message: str = ""):
+        self.guardrail_name = guardrail_name
+        self.message = f"Guardrail raised an exception, Guardrail: {guardrail_name}, Message: {message}"
+        super().__init__(self.message)
+
+
 class BlockedPiiEntityError(Exception):
     def __init__(
         self,

litellm/integrations/custom_guardrail.py

@@ -211,6 +211,8 @@ class CustomGuardrail(CustomLogger):
             masked_entity_count=masked_entity_count,
         )
         if "metadata" in request_data:
+            if request_data["metadata"] is None:
+                request_data["metadata"] = {}
             request_data["metadata"]["standard_logging_guardrail_information"] = slg
         elif "litellm_metadata" in request_data:
             request_data["litellm_metadata"][
@@ -294,6 +296,24 @@ class CustomGuardrail(CustomLogger):
         )
         raise e
 
+    def mask_content_in_string(
+        self,
+        content_string: str,
+        mask_string: str,
+        start_index: int,
+        end_index: int,
+    ) -> str:
+        """
+        Mask the content in the string between the start and end indices.
+        """
+
+        # Do nothing if the start or end are not valid
+        if not (0 <= start_index < end_index <= len(content_string)):
+            return content_string
+
+        # Mask the content
+        return content_string[:start_index] + mask_string + content_string[end_index:]
+
 
 def log_guardrail_information(func):
     """

litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py

@@ -0,0 +1,326 @@
+import copy
+import os
+from datetime import datetime
+from typing import Dict, List, Literal, Optional, Tuple, Union
+
+import litellm
+from litellm._logging import verbose_proxy_logger
+from litellm.exceptions import GuardrailRaisedException
+from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.llms.custom_httpx.http_handler import (
+    get_async_httpx_client,
+    httpxSpecialProvider,
+)
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.secret_managers.main import get_secret_str
+from litellm.types.guardrails import GuardrailEventHooks
+from litellm.types.llms.openai import AllMessageValues
+from litellm.types.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import (
+    LakeraAIRequest,
+    LakeraAIResponse,
+)
+
+
+class LakeraAIGuardrail(CustomGuardrail):
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        api_base: Optional[str] = None,
+        project_id: Optional[str] = None,
+        payload: Optional[bool] = True,
+        breakdown: Optional[bool] = True,
+        metadata: Optional[Dict] = None,
+        dev_info: Optional[bool] = True,
+        **kwargs,
+    ):
+        """
+        Initialize the LakeraAIGuardrail class.
+
+        This calls: https://api.lakera.ai/v2/guard
+
+        Args:
+            api_key: Optional[str] = None,
+            api_base: Optional[str] = None,
+            project_id: Optional[str] = None,
+            payload: Optional[bool] = True,
+            breakdown: Optional[bool] = True,
+            metadata: Optional[Dict] = None,
+            dev_info: Optional[bool] = True,
+        """
+        self.async_handler = get_async_httpx_client(
+            llm_provider=httpxSpecialProvider.GuardrailCallback
+        )
+        self.lakera_api_key = api_key or os.environ["LAKERA_API_KEY"]
+        self.project_id = project_id
+        self.api_base = (
+            api_base or get_secret_str("LAKERA_API_BASE") or "https://api.lakera.ai"
+        )
+        self.payload: Optional[bool] = payload
+        self.breakdown: Optional[bool] = breakdown
+        self.metadata: Optional[Dict] = metadata
+        self.dev_info: Optional[bool] = dev_info
+        super().__init__(**kwargs)
+
+    async def call_v2_guard(
+        self, messages: List[AllMessageValues]
+    ) -> Tuple[LakeraAIResponse, Dict]:
+        """
+        Call the Lakera AI v2 guard API.
+        """
+        status: Literal["success", "failure"] = "success"
+        exception_str: str = ""
+        start_time: datetime = datetime.now()
+        lakera_response: Optional[LakeraAIResponse] = None
+        request: Dict = {}
+        masked_entity_count: Dict = {}
+        try:
+            request = dict(
+                LakeraAIRequest(
+                    messages=messages,
+                    project_id=self.project_id,
+                    payload=self.payload,
+                    breakdown=self.breakdown,
+                    metadata=self.metadata,
+                    dev_info=self.dev_info,
+                )
+            )
+            verbose_proxy_logger.debug("Lakera AI v2 guard request: %s", request)
+            response = await self.async_handler.post(
+                url=f"{self.api_base}/v2/guard",
+                headers={"Authorization": f"Bearer {self.lakera_api_key}"},
+                json=request,
+            )
+            verbose_proxy_logger.debug(
+                "Lakera AI v2 guard response: %s", response.json()
+            )
+            lakera_response = LakeraAIResponse(**response.json())
+            return lakera_response, masked_entity_count
+        except Exception as e:
+            status = "failure"
+            exception_str = str(e)
+            raise e
+        finally:
+            ####################################################
+            # Create Guardrail Trace for logging on Langfuse, Datadog, etc.
+            ####################################################
+            guardrail_json_response: Union[Exception, str, dict, List[dict]] = {}
+            if status == "success":
+                copy_lakera_response_dict = (
+                    dict(copy.deepcopy(lakera_response)) if lakera_response else {}
+                )
+                # payload contains PII, we don't want to log it
+                copy_lakera_response_dict.pop("payload")
+                guardrail_json_response = copy_lakera_response_dict
+            else:
+                guardrail_json_response = exception_str
+            self.add_standard_logging_guardrail_information_to_request_data(
+                guardrail_json_response=guardrail_json_response,
+                guardrail_status=status,
+                request_data=dict(request) or {},
+                start_time=start_time.timestamp(),
+                end_time=datetime.now().timestamp(),
+                duration=(datetime.now() - start_time).total_seconds(),
+                masked_entity_count=masked_entity_count,
+            )
+
+    def _mask_pii_in_messages(
+        self,
+        messages: List[AllMessageValues],
+        lakera_response: Optional[LakeraAIResponse],
+        masked_entity_count: Dict,
+    ) -> List[AllMessageValues]:
+        """
+        Return a copy of messages with any detected PII replaced by
+        “[MASKED <TYPE>]” tokens.
+        """
+        payload = lakera_response.get("payload") if lakera_response else None
+        if not payload:
+            return messages
+
+        # Copy so we don’t edit the originals
+        masked = [msg.copy() for msg in messages]
+
+        # For each message, find its detections on the fly
+        for idx, msg in enumerate(masked):
+            content = msg.get("content", "")
+            if not content:
+                continue
+
+            # For v1, we only support masking content strings
+            if not isinstance(content, str):
+                continue
+
+            # Filter only detections for this message
+            detected_modifications = [d for d in payload if d.get("message_id") == idx]
+            if not detected_modifications:
+                continue
+
+            for modification in detected_modifications:
+                start, end = modification.get("start", 0), modification.get("end", 0)
+
+                # Extract the type (e.g. 'credit_card' → 'CREDIT_CARD')
+                detector_type = modification.get("detector_type", "")
+                if not detector_type:
+                    continue
+
+                typ = detector_type.split("/")[-1].upper() or "PII"
+                mask = f"[MASKED {typ}]"
+                if start is not None and end is not None:
+                    content = self.mask_content_in_string(
+                        content_string=content,
+                        mask_string=mask,
+                        start_index=start,
+                        end_index=end,
+                    )
+                    masked_entity_count[typ] = masked_entity_count.get(typ, 0) + 1
+
+            msg["content"] = content
+        return masked
+
+    async def async_pre_call_hook(
+        self,
+        user_api_key_dict: UserAPIKeyAuth,
+        cache: litellm.DualCache,
+        data: Dict,
+        call_type: Literal[
+            "completion",
+            "text_completion",
+            "embeddings",
+            "image_generation",
+            "moderation",
+            "audio_transcription",
+            "pass_through_endpoint",
+            "rerank",
+        ],
+    ) -> Optional[Union[Exception, str, Dict]]:
+        from litellm.proxy.common_utils.callback_utils import (
+            add_guardrail_to_applied_guardrails_header,
+        )
+
+        event_type: GuardrailEventHooks = GuardrailEventHooks.during_call
+        if self.should_run_guardrail(data=data, event_type=event_type) is not True:
+            return data
+
+        new_messages: Optional[List[AllMessageValues]] = data.get("messages")
+        if new_messages is None:
+            verbose_proxy_logger.warning(
+                "Lakera AI: not running guardrail. No messages in data"
+            )
+            return data
+
+        #########################################################
+        ########## 1. Make the Lakera AI v2 guard API request ##########
+        #########################################################
+        lakera_guardrail_response, masked_entity_count = await self.call_v2_guard(
+            messages=new_messages
+        )
+
+        #########################################################
+        ########## 2. Handle flagged content ##########
+        #########################################################
+        if lakera_guardrail_response.get("flagged") is True:
+            # If only PII violations exist, mask the PII
+            if self._is_only_pii_violation(lakera_guardrail_response):
+                data["messages"] = self._mask_pii_in_messages(
+                    messages=new_messages,
+                    lakera_response=lakera_guardrail_response,
+                    masked_entity_count=masked_entity_count,
+                )
+                verbose_proxy_logger.info(
+                    "Lakera AI: Masked PII in messages instead of blocking request"
+                )
+            else:
+                # If there are other violations or not set to mask PII, raise exception
+                raise GuardrailRaisedException(
+                    guardrail_name=self.guardrail_name,
+                    message="Lakera AI flagged this request. Please review the request and try again.",
+                )
+
+        #########################################################
+        ########## 3. Add the guardrail to the applied guardrails header ##########
+        #########################################################
+        add_guardrail_to_applied_guardrails_header(
+            request_data=data, guardrail_name=self.guardrail_name
+        )
+
+        return data
+
+    async def async_moderation_hook(
+        self,
+        data: dict,
+        user_api_key_dict: UserAPIKeyAuth,
+        call_type: Literal[
+            "completion",
+            "embeddings",
+            "image_generation",
+            "moderation",
+            "audio_transcription",
+            "responses",
+        ],
+    ):
+        from litellm.proxy.common_utils.callback_utils import (
+            add_guardrail_to_applied_guardrails_header,
+        )
+
+        event_type: GuardrailEventHooks = GuardrailEventHooks.during_call
+        if self.should_run_guardrail(data=data, event_type=event_type) is not True:
+            return
+
+        new_messages: Optional[List[AllMessageValues]] = data.get("messages")
+        if new_messages is None:
+            verbose_proxy_logger.warning(
+                "Lakera AI: not running guardrail. No messages in data"
+            )
+            return
+
+        #########################################################
+        ########## 1. Make the Lakera AI v2 guard API request ##########
+        #########################################################
+        lakera_guardrail_response, masked_entity_count = await self.call_v2_guard(
+            messages=new_messages
+        )
+
+        #########################################################
+        ########## 2. Handle flagged content ##########
+        #########################################################
+        if lakera_guardrail_response.get("flagged") is True:
+            # If only PII violations exist, mask the PII
+            if self._is_only_pii_violation(lakera_guardrail_response):
+                data["messages"] = self._mask_pii_in_messages(
+                    messages=new_messages,
+                    lakera_response=lakera_guardrail_response,
+                    masked_entity_count=masked_entity_count,
+                )
+                verbose_proxy_logger.info(
+                    "Lakera AI: Masked PII in messages instead of blocking request"
+                )
+            else:
+                # If there are other violations or not set to mask PII, raise exception
+                raise GuardrailRaisedException(
+                    guardrail_name=self.guardrail_name,
+                    message="Lakera AI flagged this request. Please review the request and try again.",
+                )
+
+        #########################################################
+        ########## 3. Add the guardrail to the applied guardrails header ##########
+        #########################################################
+        add_guardrail_to_applied_guardrails_header(
+            request_data=data, guardrail_name=self.guardrail_name
+        )
+
+        return data
+
+    def _is_only_pii_violation(
+        self, lakera_response: Optional[LakeraAIResponse]
+    ) -> bool:
+        """
+        Returns True if there are only PII violations in the response.
+        """
+        if not lakera_response:
+            return False
+
+        for item in lakera_response.get("payload", []) or []:
+            detector_type = item.get("detector_type", "") or ""
+            if not detector_type.startswith("pii/"):
+                return False
+        return True

litellm/proxy/guardrails/guardrail_initializers.py

@@ -60,6 +60,24 @@ def initialize_lakera(litellm_params: LitellmParams, guardrail: Guardrail):
     litellm.logging_callback_manager.add_litellm_callback(_lakera_callback)
 
 
+def initialize_lakera_v2(litellm_params: LitellmParams, guardrail: Guardrail):
+    from litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import LakeraAIGuardrail
+
+    _lakera_v2_callback = LakeraAIGuardrail(
+        api_base=litellm_params.api_base,
+        api_key=litellm_params.api_key,
+        guardrail_name=guardrail.get("guardrail_name", ""),
+        event_hook=litellm_params.mode,
+        default_on=litellm_params.default_on,
+        project_id=litellm_params.project_id,
+        payload=litellm_params.payload,
+        breakdown=litellm_params.breakdown,
+        metadata=litellm_params.metadata,
+        dev_info=litellm_params.dev_info,
+    )
+    litellm.logging_callback_manager.add_litellm_callback(_lakera_v2_callback)
+
+
 def initialize_aim(litellm_params: LitellmParams, guardrail: Guardrail):
     from litellm.proxy.guardrails.guardrail_hooks.aim import AimGuardrail
 

litellm/types/guardrails.py

@@ -24,6 +24,7 @@ class SupportedGuardrailIntegrations(Enum):
     BEDROCK = "bedrock"
     GURDRAILS_AI = "guardrails_ai"
     LAKERA = "lakera"
+    LAKERA_V2 = "lakera_v2"
     PRESIDIO = "presidio"
     HIDE_SECRETS = "hide-secrets"
     AIM = "aim"
@@ -278,9 +279,37 @@ class BedrockGuardrailConfigModel(BaseModel):
     )
 
 
+class LakeraV2GuardrailConfigModel(BaseModel):
+    """Configuration parameters for the Lakera AI v2 guardrail"""
+
+    api_key: Optional[str] = Field(
+        default=None, description="API key for the Lakera AI service"
+    )
+    api_base: Optional[str] = Field(
+        default=None, description="Base URL for the Lakera AI API"
+    )
+    project_id: Optional[str] = Field(
+        default=None, description="Project ID for the Lakera AI project"
+    )
+    payload: Optional[bool] = Field(
+        default=True, description="Whether to include payload in the response"
+    )
+    breakdown: Optional[bool] = Field(
+        default=True, description="Whether to include breakdown in the response"
+    )
+    metadata: Optional[Dict] = Field(
+        default=None, description="Additional metadata to include in the request"
+    )
+    dev_info: Optional[bool] = Field(
+        default=True,
+        description="Whether to include developer information in the response",
+    )
+
+
 class LitellmParams(
     PresidioConfigModel,
     BedrockGuardrailConfigModel,
+    LakeraV2GuardrailConfigModel,
 ):
     guardrail: str = Field(description="The type of guardrail integration to use")
     mode: str = Field(

litellm/types/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py

@@ -0,0 +1,42 @@
+from typing import Dict, List, Optional, TypedDict
+
+from litellm.types.llms.openai import AllMessageValues
+
+
+class LakeraAIRequest(TypedDict, total=False):
+    messages: List[AllMessageValues]
+    project_id: Optional[str]
+    payload: Optional[bool]
+    breakdown: Optional[bool]
+    metadata: Optional[Dict]
+    dev_info: Optional[bool]
+
+
+class LakeraAIPayloadItem(TypedDict, total=False):
+    start: Optional[int]
+    end: Optional[int]
+    text: Optional[str]
+    detector_type: Optional[str]
+    labels: Optional[List[str]]
+
+
+class LakeraAIBreakdownItem(TypedDict, total=False):
+    project_id: Optional[str]
+    policy_id: Optional[str]
+    detector_id: Optional[str]
+    detector_type: Optional[str]
+    detected: Optional[bool]
+
+
+class LakeraAIDevInfo(TypedDict, total=False):
+    git_revision: Optional[str]
+    git_timestamp: Optional[str]
+    model_version: Optional[str]
+    version: Optional[str]
+
+
+class LakeraAIResponse(TypedDict, total=False):
+    flagged: Optional[bool]
+    payload: Optional[List[LakeraAIPayloadItem]]
+    breakdown: Optional[List[LakeraAIBreakdownItem]]
+    dev_info: Optional[LakeraAIDevInfo]

tests/guardrails_tests/test_lakera_v2.py

@@ -0,0 +1,56 @@
+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
+from litellm.types.utils import CallTypes as LitellmCallTypes
+
+
+@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=os.environ.get("LAKERA_API_KEY"),
+    )
+    
+    # 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"]
+    assert "4111-1111-1111-1111" not in user_message
+    assert "test@example.com" not in user_message
+