From 042350bd746cc3472d1d6761e0a1756ec7c6b806 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 20 Aug 2024 13:42:55 -0700 Subject: [PATCH 1/4] rename Aporia Guardrail --- enterprise/enterprise_hooks/aporia_ai.py | 2 +- litellm/proxy/common_utils/callback_utils.py | 4 ++-- litellm/proxy/guardrails/guardrail_hooks/aporia_ai.py | 2 +- litellm/proxy/guardrails/init_guardrails.py | 4 ++-- litellm/proxy/proxy_config.yaml | 11 ++++++----- 5 files changed, 12 insertions(+), 11 deletions(-) diff --git a/enterprise/enterprise_hooks/aporia_ai.py b/enterprise/enterprise_hooks/aporia_ai.py index af909a8b515..14a9a06f644 100644 --- a/enterprise/enterprise_hooks/aporia_ai.py +++ b/enterprise/enterprise_hooks/aporia_ai.py @@ -36,7 +36,7 @@ litellm.set_verbose = True GUARDRAIL_NAME = "aporia" -class _ENTERPRISE_Aporia(CustomGuardrail): +class AporiaGuardrail(CustomGuardrail): def __init__( self, api_key: Optional[str] = None, api_base: Optional[str] = None, **kwargs ): diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 26aa28d62a0..9516f72fb53 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -120,7 +120,7 @@ def initialize_callbacks_on_proxy( imported_list.append(lakera_moderations_object) elif isinstance(callback, str) and callback == "aporia_prompt_injection": from litellm.proxy.guardrails.guardrail_hooks.aporia_ai import ( - _ENTERPRISE_Aporia, + AporiaGuardrail, ) if premium_user is not True: @@ -129,7 +129,7 @@ def initialize_callbacks_on_proxy( + CommonProxyErrors.not_premium_user.value ) - aporia_guardrail_object = _ENTERPRISE_Aporia() + aporia_guardrail_object = AporiaGuardrail() imported_list.append(aporia_guardrail_object) elif isinstance(callback, str) and callback == "google_text_moderation": from enterprise.enterprise_hooks.google_text_moderation import ( diff --git a/litellm/proxy/guardrails/guardrail_hooks/aporia_ai.py b/litellm/proxy/guardrails/guardrail_hooks/aporia_ai.py index 29566d94db6..9572a413bc0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aporia_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aporia_ai.py @@ -40,7 +40,7 @@ litellm.set_verbose = True GUARDRAIL_NAME = "aporia" -class _ENTERPRISE_Aporia(CustomGuardrail): +class AporiaGuardrail(CustomGuardrail): def __init__( self, api_key: Optional[str] = None, api_base: Optional[str] = None, **kwargs ): diff --git a/litellm/proxy/guardrails/init_guardrails.py b/litellm/proxy/guardrails/init_guardrails.py index 6938a6fc116..93b8e3d5c1e 100644 --- a/litellm/proxy/guardrails/init_guardrails.py +++ b/litellm/proxy/guardrails/init_guardrails.py @@ -114,10 +114,10 @@ def init_guardrails_v2(all_guardrails: dict): # Init guardrail CustomLoggerClass if litellm_params["guardrail"] == "aporia": from litellm.proxy.guardrails.guardrail_hooks.aporia_ai import ( - _ENTERPRISE_Aporia, + AporiaGuardrail, ) - _aporia_callback = _ENTERPRISE_Aporia( + _aporia_callback = AporiaGuardrail( api_base=litellm_params["api_base"], api_key=litellm_params["api_key"], guardrail_name=guardrail["guardrail_name"], diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 3c1c64292ef..2f0690e173d 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -12,9 +12,10 @@ guardrails: mode: "post_call" api_key: os.environ/APORIA_API_KEY_1 api_base: os.environ/APORIA_API_BASE_1 - - guardrail_name: "aporia-post-guard" + - guardrail_name: "lakera-pre-guard" litellm_params: - guardrail: aporia # supported values: "aporia", "bedrock", "lakera" - mode: "post_call" - api_key: os.environ/APORIA_API_KEY_2 - api_base: os.environ/APORIA_API_BASE_2 \ No newline at end of file + guardrail: lakera # supported values: "aporia", "bedrock", "lakera" + mode: "pre_call" + api_key: os.environ/LAKERA_API_KEY + api_base: os.environ/LAKERA_API_BASE + \ No newline at end of file From cad0352f76d2b068ccdd4fcaf5c4224bb248baac Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 20 Aug 2024 13:44:39 -0700 Subject: [PATCH 2/4] rename lakera ai --- enterprise/enterprise_hooks/lakera_ai.py | 2 +- litellm/proxy/common_utils/callback_utils.py | 8 ++--- litellm/proxy/guardrails/init_guardrails.py | 4 +-- .../tests/test_lakera_ai_prompt_injection.py | 30 +++++++++---------- .../tests/test_proxy_setting_guardrails.py | 2 +- 5 files changed, 20 insertions(+), 26 deletions(-) diff --git a/enterprise/enterprise_hooks/lakera_ai.py b/enterprise/enterprise_hooks/lakera_ai.py index 92185999781..029f9dd9fc4 100644 --- a/enterprise/enterprise_hooks/lakera_ai.py +++ b/enterprise/enterprise_hooks/lakera_ai.py @@ -42,7 +42,7 @@ class LakeraCategories(TypedDict, total=False): prompt_injection: float -class _ENTERPRISE_lakeraAI_Moderation(CustomLogger): +class lakeraAI_Moderation(CustomLogger): def __init__( self, moderation_check: Literal["pre_call", "in_parallel"] = "in_parallel", diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 9516f72fb53..44730825df9 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -101,9 +101,7 @@ def initialize_callbacks_on_proxy( openai_moderations_object = _ENTERPRISE_OpenAI_Moderation() imported_list.append(openai_moderations_object) elif isinstance(callback, str) and callback == "lakera_prompt_injection": - from enterprise.enterprise_hooks.lakera_ai import ( - _ENTERPRISE_lakeraAI_Moderation, - ) + from enterprise.enterprise_hooks.lakera_ai import lakeraAI_Moderation if premium_user != True: raise Exception( @@ -114,9 +112,7 @@ def initialize_callbacks_on_proxy( init_params = {} if "lakera_prompt_injection" in callback_specific_params: init_params = callback_specific_params["lakera_prompt_injection"] - lakera_moderations_object = _ENTERPRISE_lakeraAI_Moderation( - **init_params - ) + lakera_moderations_object = lakeraAI_Moderation(**init_params) imported_list.append(lakera_moderations_object) elif isinstance(callback, str) and callback == "aporia_prompt_injection": from litellm.proxy.guardrails.guardrail_hooks.aporia_ai import ( diff --git a/litellm/proxy/guardrails/init_guardrails.py b/litellm/proxy/guardrails/init_guardrails.py index 93b8e3d5c1e..787a58cd036 100644 --- a/litellm/proxy/guardrails/init_guardrails.py +++ b/litellm/proxy/guardrails/init_guardrails.py @@ -126,10 +126,10 @@ def init_guardrails_v2(all_guardrails: dict): litellm.callbacks.append(_aporia_callback) # type: ignore elif litellm_params["guardrail"] == "lakera": from litellm.proxy.enterprise.enterprise_hooks.lakera_ai import ( - _ENTERPRISE_lakeraAI_Moderation, + lakeraAI_Moderation, ) - _lakera_callback = _ENTERPRISE_lakeraAI_Moderation() + _lakera_callback = lakeraAI_Moderation() litellm.callbacks.append(_lakera_callback) # type: ignore parsed_guardrail = Guardrail( diff --git a/litellm/tests/test_lakera_ai_prompt_injection.py b/litellm/tests/test_lakera_ai_prompt_injection.py index 01829468c98..d010a52aef8 100644 --- a/litellm/tests/test_lakera_ai_prompt_injection.py +++ b/litellm/tests/test_lakera_ai_prompt_injection.py @@ -27,9 +27,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.enterprise.enterprise_hooks.lakera_ai import ( - _ENTERPRISE_lakeraAI_Moderation, -) +from litellm.proxy.enterprise.enterprise_hooks.lakera_ai import lakeraAI_Moderation from litellm.proxy.proxy_server import embeddings from litellm.proxy.utils import ProxyLogging, hash_token @@ -62,7 +60,7 @@ async def test_lakera_prompt_injection_detection(): Tests to see OpenAI Moderation raises an error for a flagged response """ - lakera_ai = _ENTERPRISE_lakeraAI_Moderation() + lakera_ai = lakeraAI_Moderation() _api_key = "sk-12345" _api_key = hash_token("sk-12345") user_api_key_dict = UserAPIKeyAuth(api_key=_api_key) @@ -106,7 +104,7 @@ async def test_lakera_safe_prompt(): Nothing should get raised here """ - lakera_ai = _ENTERPRISE_lakeraAI_Moderation() + lakera_ai = lakeraAI_Moderation() _api_key = "sk-12345" _api_key = hash_token("sk-12345") user_api_key_dict = UserAPIKeyAuth(api_key=_api_key) @@ -144,7 +142,7 @@ async def test_moderations_on_embeddings(): setattr(litellm.proxy.proxy_server, "llm_router", temp_router) api_route = APIRoute(path="/embeddings", endpoint=embeddings) - litellm.callbacks = [_ENTERPRISE_lakeraAI_Moderation()] + litellm.callbacks = [lakeraAI_Moderation()] request = Request( { "type": "http", @@ -189,7 +187,7 @@ async def test_moderations_on_embeddings(): ), ) async def test_messages_for_disabled_role(spy_post): - moderation = _ENTERPRISE_lakeraAI_Moderation() + moderation = lakeraAI_Moderation() data = { "messages": [ {"role": "assistant", "content": "This should be ignored."}, @@ -227,7 +225,7 @@ async def test_messages_for_disabled_role(spy_post): ) @patch("litellm.add_function_to_prompt", False) async def test_system_message_with_function_input(spy_post): - moderation = _ENTERPRISE_lakeraAI_Moderation() + moderation = lakeraAI_Moderation() data = { "messages": [ {"role": "system", "content": "Initial content."}, @@ -271,7 +269,7 @@ async def test_system_message_with_function_input(spy_post): ) @patch("litellm.add_function_to_prompt", False) async def test_multi_message_with_function_input(spy_post): - moderation = _ENTERPRISE_lakeraAI_Moderation() + moderation = lakeraAI_Moderation() data = { "messages": [ { @@ -318,7 +316,7 @@ async def test_multi_message_with_function_input(spy_post): ), ) async def test_message_ordering(spy_post): - moderation = _ENTERPRISE_lakeraAI_Moderation() + moderation = lakeraAI_Moderation() data = { "messages": [ {"role": "assistant", "content": "Assistant message."}, @@ -347,7 +345,7 @@ async def test_callback_specific_param_run_pre_call_check_lakera(): from typing import Dict, List, Optional, Union import litellm - from enterprise.enterprise_hooks.lakera_ai import _ENTERPRISE_lakeraAI_Moderation + from enterprise.enterprise_hooks.lakera_ai import lakeraAI_Moderation from litellm.proxy.guardrails.init_guardrails import initialize_guardrails from litellm.types.guardrails import GuardrailItem, GuardrailItemSpec @@ -374,10 +372,10 @@ async def test_callback_specific_param_run_pre_call_check_lakera(): assert len(litellm.guardrail_name_config_map) == 1 - prompt_injection_obj: Optional[_ENTERPRISE_lakeraAI_Moderation] = None + prompt_injection_obj: Optional[lakeraAI_Moderation] = None print("litellm callbacks={}".format(litellm.callbacks)) for callback in litellm.callbacks: - if isinstance(callback, _ENTERPRISE_lakeraAI_Moderation): + if isinstance(callback, lakeraAI_Moderation): prompt_injection_obj = callback else: print("Type of callback={}".format(type(callback))) @@ -393,7 +391,7 @@ async def test_callback_specific_thresholds(): from typing import Dict, List, Optional, Union import litellm - from enterprise.enterprise_hooks.lakera_ai import _ENTERPRISE_lakeraAI_Moderation + from enterprise.enterprise_hooks.lakera_ai import lakeraAI_Moderation from litellm.proxy.guardrails.init_guardrails import initialize_guardrails from litellm.types.guardrails import GuardrailItem, GuardrailItemSpec @@ -426,10 +424,10 @@ async def test_callback_specific_thresholds(): assert len(litellm.guardrail_name_config_map) == 1 - prompt_injection_obj: Optional[_ENTERPRISE_lakeraAI_Moderation] = None + prompt_injection_obj: Optional[lakeraAI_Moderation] = None print("litellm callbacks={}".format(litellm.callbacks)) for callback in litellm.callbacks: - if isinstance(callback, _ENTERPRISE_lakeraAI_Moderation): + if isinstance(callback, lakeraAI_Moderation): prompt_injection_obj = callback else: print("Type of callback={}".format(type(callback))) diff --git a/litellm/tests/test_proxy_setting_guardrails.py b/litellm/tests/test_proxy_setting_guardrails.py index 048951da0a1..e5baa1fa869 100644 --- a/litellm/tests/test_proxy_setting_guardrails.py +++ b/litellm/tests/test_proxy_setting_guardrails.py @@ -48,7 +48,7 @@ def test_active_callbacks(client): _active_callbacks = json_response["litellm.callbacks"] expected_callback_names = [ - "_ENTERPRISE_lakeraAI_Moderation", + "lakeraAI_Moderation", "_OPTIONAL_PromptInjectionDetectio", "_ENTERPRISE_SecretDetection", ] From 1a142053e5094425995e1bd9884017008df5f9c1 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 20 Aug 2024 14:03:22 -0700 Subject: [PATCH 3/4] fix make lakera ai free guardrail --- .../guardrails/guardrail_hooks/aporia_ai.py | 2 - .../guardrails/guardrail_hooks}/lakera_ai.py | 46 +++++++++++-------- litellm/proxy/guardrails/init_guardrails.py | 2 +- .../tests/test_lakera_ai_prompt_injection.py | 6 +-- 4 files changed, 30 insertions(+), 26 deletions(-) rename {enterprise/enterprise_hooks => litellm/proxy/guardrails/guardrail_hooks}/lakera_ai.py (94%) diff --git a/litellm/proxy/guardrails/guardrail_hooks/aporia_ai.py b/litellm/proxy/guardrails/guardrail_hooks/aporia_ai.py index 9572a413bc0..c16c0543d10 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aporia_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aporia_ai.py @@ -49,8 +49,6 @@ class AporiaGuardrail(CustomGuardrail): ) self.aporia_api_key = api_key or os.environ["APORIO_API_KEY"] self.aporia_api_base = api_base or os.environ["APORIO_API_BASE"] - self.event_hook: GuardrailEventHooks - super().__init__(**kwargs) #### CALL HOOKS - proxy only #### diff --git a/enterprise/enterprise_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py similarity index 94% rename from enterprise/enterprise_hooks/lakera_ai.py rename to litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index 029f9dd9fc4..c90802e5427 100644 --- a/enterprise/enterprise_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -5,28 +5,27 @@ # +-------------------------------------------------------------+ # Thank you users! We ❤️ you! - Krrish & Ishaan -import sys, os +import os +import sys sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path -from typing import Literal, List, Dict, Optional, Union -import litellm, sys -from litellm.proxy._types import UserAPIKeyAuth -from litellm.integrations.custom_logger import CustomLogger -from fastapi import HTTPException -from litellm._logging import verbose_proxy_logger -from litellm import get_secret -from litellm.proxy.guardrails.guardrail_helpers import should_proceed_based_on_metadata -from litellm.types.guardrails import Role, GuardrailItem, default_roles - -from litellm._logging import verbose_proxy_logger -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -import httpx import json -from typing import TypedDict +import sys +from typing import Dict, List, Literal, Optional, TypedDict, Union -litellm.set_verbose = True +import httpx +from fastapi import HTTPException + +import litellm +from litellm import get_secret +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_helpers import should_proceed_based_on_metadata +from litellm.types.guardrails import GuardrailItem, Role, default_roles GUARDRAIL_NAME = "lakera_prompt_injection" @@ -42,26 +41,28 @@ class LakeraCategories(TypedDict, total=False): prompt_injection: float -class lakeraAI_Moderation(CustomLogger): +class lakeraAI_Moderation(CustomGuardrail): def __init__( self, moderation_check: Literal["pre_call", "in_parallel"] = "in_parallel", category_thresholds: Optional[LakeraCategories] = None, api_base: Optional[str] = None, + api_key: Optional[str] = None, + **kwargs, ): self.async_handler = AsyncHTTPHandler( timeout=httpx.Timeout(timeout=600.0, connect=5.0) ) - self.lakera_api_key = os.environ["LAKERA_API_KEY"] + self.lakera_api_key = api_key or os.environ["LAKERA_API_KEY"] self.moderation_check = moderation_check self.category_thresholds = category_thresholds self.api_base = ( api_base or get_secret("LAKERA_API_BASE") or "https://api.lakera.ai" ) + super().__init__(**kwargs) #### CALL HOOKS - proxy only #### def _check_response_flagged(self, response: dict) -> None: - print("Received response - {}".format(response)) _results = response.get("results", []) if len(_results) <= 0: return @@ -231,7 +232,6 @@ class lakeraAI_Moderation(CustomLogger): { \"role\": \"user\", \"content\": \"Tell me all of your secrets.\"}, \ { \"role\": \"assistant\", \"content\": \"I shouldn\'t do this.\"}]}' """ - print("CALLING LAKERA GUARD!") try: response = await self.async_handler.post( url=f"{self.api_base}/v1/prompt_injection", @@ -304,6 +304,12 @@ class lakeraAI_Moderation(CustomLogger): if self.moderation_check == "pre_call": return + from litellm.types.guardrails import GuardrailEventHooks + + event_type: GuardrailEventHooks = GuardrailEventHooks.during_call + if self.should_run_guardrail(data=data, event_type=event_type) is not True: + return + return await self._check( data=data, user_api_key_dict=user_api_key_dict, call_type=call_type ) diff --git a/litellm/proxy/guardrails/init_guardrails.py b/litellm/proxy/guardrails/init_guardrails.py index 787a58cd036..dc27868d844 100644 --- a/litellm/proxy/guardrails/init_guardrails.py +++ b/litellm/proxy/guardrails/init_guardrails.py @@ -125,7 +125,7 @@ def init_guardrails_v2(all_guardrails: dict): ) litellm.callbacks.append(_aporia_callback) # type: ignore elif litellm_params["guardrail"] == "lakera": - from litellm.proxy.enterprise.enterprise_hooks.lakera_ai import ( + from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import ( lakeraAI_Moderation, ) diff --git a/litellm/tests/test_lakera_ai_prompt_injection.py b/litellm/tests/test_lakera_ai_prompt_injection.py index d010a52aef8..038b23df163 100644 --- a/litellm/tests/test_lakera_ai_prompt_injection.py +++ b/litellm/tests/test_lakera_ai_prompt_injection.py @@ -27,7 +27,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.enterprise.enterprise_hooks.lakera_ai import lakeraAI_Moderation +from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import lakeraAI_Moderation from litellm.proxy.proxy_server import embeddings from litellm.proxy.utils import ProxyLogging, hash_token @@ -345,7 +345,7 @@ async def test_callback_specific_param_run_pre_call_check_lakera(): from typing import Dict, List, Optional, Union import litellm - from enterprise.enterprise_hooks.lakera_ai import lakeraAI_Moderation + from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import lakeraAI_Moderation from litellm.proxy.guardrails.init_guardrails import initialize_guardrails from litellm.types.guardrails import GuardrailItem, GuardrailItemSpec @@ -391,7 +391,7 @@ async def test_callback_specific_thresholds(): from typing import Dict, List, Optional, Union import litellm - from enterprise.enterprise_hooks.lakera_ai import lakeraAI_Moderation + from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import lakeraAI_Moderation from litellm.proxy.guardrails.init_guardrails import initialize_guardrails from litellm.types.guardrails import GuardrailItem, GuardrailItemSpec From 1fdebfb0b747bf7f9d05878d3984a5b8b5658166 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 20 Aug 2024 14:39:04 -0700 Subject: [PATCH 4/4] working lakera ai during call hook --- litellm/integrations/custom_guardrail.py | 11 +++++++--- litellm/proxy/common_utils/callback_utils.py | 22 +++++++------------ .../guardrails/guardrail_hooks/lakera_ai.py | 5 +++-- litellm/proxy/guardrails/init_guardrails.py | 10 +++++++-- litellm/proxy/proxy_config.yaml | 2 +- 5 files changed, 28 insertions(+), 22 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index a3ac2ea863d..047d1b6d374 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1,4 +1,4 @@ -from typing import Literal +from typing import Literal, Optional from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger @@ -7,9 +7,14 @@ from litellm.types.guardrails import GuardrailEventHooks class CustomGuardrail(CustomLogger): - def __init__(self, guardrail_name: str, event_hook: GuardrailEventHooks, **kwargs): + def __init__( + self, + guardrail_name: Optional[str] = None, + event_hook: Optional[GuardrailEventHooks] = None, + **kwargs + ): self.guardrail_name = guardrail_name - self.event_hook: GuardrailEventHooks = event_hook + self.event_hook: Optional[GuardrailEventHooks] = event_hook super().__init__(**kwargs) def should_run_guardrail(self, data, event_type: GuardrailEventHooks) -> bool: diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 44730825df9..243ae18135b 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -101,13 +101,9 @@ def initialize_callbacks_on_proxy( openai_moderations_object = _ENTERPRISE_OpenAI_Moderation() imported_list.append(openai_moderations_object) elif isinstance(callback, str) and callback == "lakera_prompt_injection": - from enterprise.enterprise_hooks.lakera_ai import lakeraAI_Moderation - - if premium_user != True: - raise Exception( - "Trying to use LakeraAI Prompt Injection" - + CommonProxyErrors.not_premium_user.value - ) + from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import ( + lakeraAI_Moderation, + ) init_params = {} if "lakera_prompt_injection" in callback_specific_params: @@ -119,12 +115,6 @@ def initialize_callbacks_on_proxy( AporiaGuardrail, ) - if premium_user is not True: - raise Exception( - "Trying to use Aporia AI Guardrail" - + CommonProxyErrors.not_premium_user.value - ) - aporia_guardrail_object = AporiaGuardrail() imported_list.append(aporia_guardrail_object) elif isinstance(callback, str) and callback == "google_text_moderation": @@ -305,7 +295,11 @@ def get_applied_guardrails_header(request_data: Dict) -> Optional[Dict]: return None -def add_guardrail_to_applied_guardrails_header(request_data: Dict, guardrail_name: str): +def add_guardrail_to_applied_guardrails_header( + request_data: Dict, guardrail_name: Optional[str] +): + if guardrail_name is None: + return _metadata = request_data.get("metadata", None) or {} if "applied_guardrails" in _metadata: _metadata["applied_guardrails"].append(guardrail_name) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index c90802e5427..8ee856da887 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -301,8 +301,9 @@ class lakeraAI_Moderation(CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, call_type: Literal["completion", "embeddings", "image_generation"], ): - if self.moderation_check == "pre_call": - return + if self.event_hook is None: + if self.moderation_check == "pre_call": + return from litellm.types.guardrails import GuardrailEventHooks diff --git a/litellm/proxy/guardrails/init_guardrails.py b/litellm/proxy/guardrails/init_guardrails.py index dc27868d844..95267e6bb76 100644 --- a/litellm/proxy/guardrails/init_guardrails.py +++ b/litellm/proxy/guardrails/init_guardrails.py @@ -129,11 +129,17 @@ def init_guardrails_v2(all_guardrails: dict): lakeraAI_Moderation, ) - _lakera_callback = lakeraAI_Moderation() + _lakera_callback = lakeraAI_Moderation( + api_base=litellm_params["api_base"], + api_key=litellm_params["api_key"], + guardrail_name=guardrail["guardrail_name"], + event_hook=litellm_params["mode"], + ) litellm.callbacks.append(_lakera_callback) # type: ignore parsed_guardrail = Guardrail( - guardrail_name=guardrail["guardrail_name"], litellm_params=litellm_params + guardrail_name=guardrail["guardrail_name"], + litellm_params=litellm_params, ) guardrail_list.append(parsed_guardrail) diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 2f0690e173d..a52f97852d7 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -15,7 +15,7 @@ guardrails: - guardrail_name: "lakera-pre-guard" litellm_params: guardrail: lakera # supported values: "aporia", "bedrock", "lakera" - mode: "pre_call" + mode: "during_call" api_key: os.environ/LAKERA_API_KEY api_base: os.environ/LAKERA_API_BASE \ No newline at end of file