From 042350bd746cc3472d1d6761e0a1756ec7c6b806 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 20 Aug 2024 13:42:55 -0700 Subject: [PATCH] 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 af909a8b5..14a9a06f6 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 26aa28d62..9516f72fb 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 29566d94d..9572a413b 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 6938a6fc1..93b8e3d5c 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 3c1c64292..2f0690e17 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