phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 11:14:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

working lakera ai during call hook

Browse source
This commit is contained in:
Ishaan Jaff 2024-08-20 14:39:04 -07:00
parent 9d809e8404
commit cdbd245c3d
5 changed files with 28 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

22
litellm/proxy/common_utils/callback_utils.py
View file

@ -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)