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

add bedrock guardrails support

Browse source
This commit is contained in:
Ishaan Jaff 2024-08-22 16:09:55 -07:00
parent 5f1a03e897
commit 2f01a22ef7
4 changed files with 296 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

18
litellm/proxy/guardrails/init_guardrails.py
View file

@ -96,8 +96,10 @@ def init_guardrails_v2(all_guardrails: dict):
litellm_params = LitellmParams(
guardrail=litellm_params_data["guardrail"],
mode=litellm_params_data["mode"],
api_key=litellm_params_data["api_key"],
api_base=litellm_params_data["api_base"],
api_key=litellm_params_data.get("api_key"),
api_base=litellm_params_data.get("api_base"),
guardrailIdentifier=litellm_params_data.get("guardrailIdentifier"),
guardrailVersion=litellm_params_data.get("guardrailVersion"),
)
if (
@ -134,6 +136,18 @@ def init_guardrails_v2(all_guardrails: dict):
event_hook=litellm_params["mode"],
)
litellm.callbacks.append(_aporia_callback) # type: ignore
if litellm_params["guardrail"] == "bedrock":
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
BedrockGuardrail,
)
_bedrock_callback = BedrockGuardrail(
guardrail_name=guardrail["guardrail_name"],
event_hook=litellm_params["mode"],
guardrailIdentifier=litellm_params["guardrailIdentifier"],
guardrailVersion=litellm_params["guardrailVersion"],
)
litellm.callbacks.append(_bedrock_callback) # type: ignore
elif litellm_params["guardrail"] == "lakera":
from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import (
lakeraAI_Moderation,