diff --git a/enterprise/enterprise_hooks/lakera_ai.py b/enterprise/enterprise_hooks/lakera_ai.py index 75e346cdb..40136f741 100644 --- a/enterprise/enterprise_hooks/lakera_ai.py +++ b/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" - ].enabled_roles + prompt_injection_obj: Optional[GuardrailItem] = ( + litellm.guardrail_name_config_map.get("prompt_injection") + ) + 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, diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 6df2e0825..1bf073513 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -4,7 +4,4 @@ model_list: model: "*" litellm_settings: - guardrails: - - 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 \ No newline at end of file + callbacks: ["lakera_prompt_injection"] \ No newline at end of file