scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

fix lakera ai tests

Browse source
This commit is contained in:
Ishaan Jaff 2024-08-20 17:28:18 -07:00
parent f130ff2587
commit 5bd94f64da
2 changed files with 17 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py
View file

@ -295,14 +295,16 @@ class lakeraAI_Moderation(CustomGuardrail):
if self.event_hook is None:
if self.moderation_check == "in_parallel":
return None
else:
# v2 guardrails implementation
if (
self.should_run_guardrail(
data=data, event_type=GuardrailEventHooks.pre_call
)
is not True
):
return None
if (
self.should_run_guardrail(
data=data, event_type=GuardrailEventHooks.pre_call
)
is not True
):
return None
return await self._check(
data=data, user_api_key_dict=user_api_key_dict, call_type=call_type
@ -317,12 +319,13 @@ class lakeraAI_Moderation(CustomGuardrail):
if self.event_hook is None:
if self.moderation_check == "pre_call":
return
else:
# V2 Guardrails implementation
from litellm.types.guardrails import GuardrailEventHooks
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
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

2
litellm/tests/test_lakera_ai_prompt_injection.py
View file

@ -85,6 +85,8 @@ async def test_lakera_prompt_injection_detection():
# Assert that the laker ai response is in the exception raise
assert "lakera_ai_response" in http_exception.detail
assert "Violated content safety policy" in str(http_exception)
except Exception as e:
print("got exception running lakera ai test", str(e))
@patch(