fix(lakera_ai.py): fix hardcoded prompt_injection string in lakera integration

enterprise/enterprise_hooks/lakera_ai.py

@@ -10,7 +10,7 @@ import sys, os
 sys.path.insert(
     0, os.path.abspath("../..")
 )  # Adds the parent directory to the system path
-from typing import Literal, List, Dict
+from typing import Literal, List, Dict, Optional
 import litellm, sys
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.integrations.custom_logger import CustomLogger
@@ -64,11 +64,24 @@ class _ENTERPRISE_lakeraAI_Moderation(CustomLogger):
             return
         text = ""
         if "messages" in data and isinstance(data["messages"], list):
-            enabled_roles = litellm.guardrail_name_config_map[
+            prompt_injection_obj: Optional[GuardrailItem] = (
-                "prompt_injection"
+                litellm.guardrail_name_config_map.get("prompt_injection")
-            ].enabled_roles
+            )
+            if prompt_injection_obj is not None:
+                enabled_roles = prompt_injection_obj.enabled_roles
+            else:
+                enabled_roles = None
+
             if enabled_roles is None:
                 enabled_roles = default_roles
+
+            stringified_roles: List[str] = []
+            if enabled_roles is not None:  # convert to list of str
+                for role in enabled_roles:
+                    if isinstance(role, Role):
+                        stringified_roles.append(role.value)
+                    elif isinstance(role, str):
+                        stringified_roles.append(role)
             lakera_input_dict: Dict = {
                 role: None for role in INPUT_POSITIONING_MAP.keys()
             }
@@ -76,7 +89,7 @@ class _ENTERPRISE_lakeraAI_Moderation(CustomLogger):
             tool_call_messages: List = []
             for message in data["messages"]:
                 role = message.get("role")
-                if role in enabled_roles:
+                if role in stringified_roles:
                     if "tool_calls" in message:
                         tool_call_messages = [
                             *tool_call_messages,

litellm/proxy/_new_secret_config.yaml

@@ -4,7 +4,4 @@ model_list:
       model: "*"
 
 litellm_settings:
-  guardrails:
+  callbacks: ["lakera_prompt_injection"]
-    - prompt_injection:  # your custom name for guardrail
-        callbacks: [lakera_prompt_injection] # litellm callbacks to use
-        default_on: true # will run on all llm requests when true