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 the ability to init a custom guardrail

Browse source
This commit is contained in:
Ishaan Jaff 2024-08-23 08:21:07 -07:00
parent be853d93da
commit 65e222b40f
1 changed files with 20 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -1,3 +1,4 @@
import importlib
import traceback
from typing import Dict, List, Literal
@ -161,6 +162,25 @@ def init_guardrails_v2(all_guardrails: dict):
category_thresholds=litellm_params.get("category_thresholds"),
)
litellm.callbacks.append(_lakera_callback) # type: ignore
elif (
isinstance(litellm_params["guardrail"], str)
and "." in litellm_params["guardrail"]
):
# Custom guardrail
_guardrail = litellm_params["guardrail"]
_file_name, _class_name = _guardrail.split(".")
verbose_proxy_logger.debug(
"Initializing custom guardrail: %s, file_name: %s, class_name: %s",
_guardrail,
_file_name,
_class_name,
)
_guardrail_class = getattr(importlib.import_module(_file_name), _class_name)
_guardrail_callback = _guardrail_class(
guardrail_name=guardrail["guardrail_name"],
event_hook=litellm_params["mode"],
)
litellm.callbacks.append(_guardrail_callback) # type: ignore
parsed_guardrail = Guardrail(
guardrail_name=guardrail["guardrail_name"],
@ -169,6 +189,5 @@ def init_guardrails_v2(all_guardrails: dict):
guardrail_list.append(parsed_guardrail)
guardrail_name = guardrail["guardrail_name"]
# pretty print guardrail_list in green
print(f"\nGuardrail List:{guardrail_list}\n") # noqa