diff --git a/docs/my-website/docs/proxy/guardrails/lakera_ai.md b/docs/my-website/docs/proxy/guardrails/lakera_ai.md index ba1ca0b218..e66329dcb0 100644 --- a/docs/my-website/docs/proxy/guardrails/lakera_ai.md +++ b/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 -```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 \ - -## 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 - -``` \ No newline at end of file diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 5beb2041c3..9f3411143a 100644 --- a/litellm/exceptions.py +++ b/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, diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 71c683cc9f..0bc9481563 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/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): """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py new file mode 100644 index 0000000000..01b6b9db4d --- /dev/null +++ b/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 ]” 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 diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 68b0dadc96..55d746f1b6 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/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 diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 5e04f2e24b..baa1cd9c54 100644 --- a/litellm/types/guardrails.py +++ b/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( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/types/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py new file mode 100644 index 0000000000..94c385b6ce --- /dev/null +++ b/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] diff --git a/tests/guardrails_tests/test_lakera_v2.py b/tests/guardrails_tests/test_lakera_v2.py new file mode 100644 index 0000000000..19c4424bee --- /dev/null +++ b/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 +