mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 19:24:27 +00:00
ed283bc5b4
All checks were successful
Read Version from pyproject.toml / read-version (push) Successful in 12s
* test_default_on_guardrail * update debug on custom guardrail * refactor guardrails init * guardrail registry * allow switching guardrails default_on * fix circle import issue * fix bedrock applying guardrails where content is a list * fix unused import * docs default on guardrail * docs fix per api key
114 lines
4.2 KiB
Python
114 lines
4.2 KiB
Python
# litellm/proxy/guardrails/guardrail_initializers.py
|
|
import litellm
|
|
from litellm.types.guardrails import *
|
|
|
|
|
|
def initialize_aporia(litellm_params, guardrail):
|
|
from litellm.proxy.guardrails.guardrail_hooks.aporia_ai import AporiaGuardrail
|
|
|
|
_aporia_callback = AporiaGuardrail(
|
|
api_base=litellm_params["api_base"],
|
|
api_key=litellm_params["api_key"],
|
|
guardrail_name=guardrail["guardrail_name"],
|
|
event_hook=litellm_params["mode"],
|
|
default_on=litellm_params["default_on"],
|
|
)
|
|
litellm.callbacks.append(_aporia_callback)
|
|
|
|
|
|
def initialize_bedrock(litellm_params, guardrail):
|
|
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"],
|
|
default_on=litellm_params["default_on"],
|
|
)
|
|
litellm.callbacks.append(_bedrock_callback)
|
|
|
|
|
|
def initialize_lakera(litellm_params, guardrail):
|
|
from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import lakeraAI_Moderation
|
|
|
|
_lakera_callback = lakeraAI_Moderation(
|
|
api_base=litellm_params["api_base"],
|
|
api_key=litellm_params["api_key"],
|
|
guardrail_name=guardrail["guardrail_name"],
|
|
event_hook=litellm_params["mode"],
|
|
category_thresholds=litellm_params.get("category_thresholds"),
|
|
default_on=litellm_params["default_on"],
|
|
)
|
|
litellm.callbacks.append(_lakera_callback)
|
|
|
|
|
|
def initialize_aim(litellm_params, guardrail):
|
|
from litellm.proxy.guardrails.guardrail_hooks.aim import AimGuardrail
|
|
|
|
_aim_callback = AimGuardrail(
|
|
api_base=litellm_params["api_base"],
|
|
api_key=litellm_params["api_key"],
|
|
guardrail_name=guardrail["guardrail_name"],
|
|
event_hook=litellm_params["mode"],
|
|
default_on=litellm_params["default_on"],
|
|
)
|
|
litellm.callbacks.append(_aim_callback)
|
|
|
|
|
|
def initialize_presidio(litellm_params, guardrail):
|
|
from litellm.proxy.guardrails.guardrail_hooks.presidio import (
|
|
_OPTIONAL_PresidioPIIMasking,
|
|
)
|
|
|
|
_presidio_callback = _OPTIONAL_PresidioPIIMasking(
|
|
guardrail_name=guardrail["guardrail_name"],
|
|
event_hook=litellm_params["mode"],
|
|
output_parse_pii=litellm_params["output_parse_pii"],
|
|
presidio_ad_hoc_recognizers=litellm_params["presidio_ad_hoc_recognizers"],
|
|
mock_redacted_text=litellm_params.get("mock_redacted_text") or None,
|
|
default_on=litellm_params["default_on"],
|
|
)
|
|
litellm.callbacks.append(_presidio_callback)
|
|
|
|
if litellm_params["output_parse_pii"]:
|
|
_success_callback = _OPTIONAL_PresidioPIIMasking(
|
|
output_parse_pii=True,
|
|
guardrail_name=guardrail["guardrail_name"],
|
|
event_hook=GuardrailEventHooks.post_call.value,
|
|
presidio_ad_hoc_recognizers=litellm_params["presidio_ad_hoc_recognizers"],
|
|
default_on=litellm_params["default_on"],
|
|
)
|
|
litellm.callbacks.append(_success_callback)
|
|
|
|
|
|
def initialize_hide_secrets(litellm_params, guardrail):
|
|
from enterprise.enterprise_hooks.secret_detection import _ENTERPRISE_SecretDetection
|
|
|
|
_secret_detection_object = _ENTERPRISE_SecretDetection(
|
|
detect_secrets_config=litellm_params.get("detect_secrets_config"),
|
|
event_hook=litellm_params["mode"],
|
|
guardrail_name=guardrail["guardrail_name"],
|
|
default_on=litellm_params["default_on"],
|
|
)
|
|
litellm.callbacks.append(_secret_detection_object)
|
|
|
|
|
|
def initialize_guardrails_ai(litellm_params, guardrail):
|
|
from litellm.proxy.guardrails.guardrail_hooks.guardrails_ai import GuardrailsAI
|
|
|
|
_guard_name = litellm_params.get("guard_name")
|
|
if not _guard_name:
|
|
raise Exception(
|
|
"GuardrailsAIException - Please pass the Guardrails AI guard name via 'litellm_params::guard_name'"
|
|
)
|
|
|
|
_guardrails_ai_callback = GuardrailsAI(
|
|
api_base=litellm_params.get("api_base"),
|
|
guard_name=_guard_name,
|
|
guardrail_name=SupportedGuardrailIntegrations.GURDRAILS_AI.value,
|
|
default_on=litellm_params["default_on"],
|
|
)
|
|
litellm.callbacks.append(_guardrails_ai_callback)
|