Litellm dev 01 14 2025 p1 (#7771)

* First-class Aim Guardrails support (#7738)

* initial aim support

* add tests

* docs(langsmith_integration.md): cleanup

* style: cleanup unused imports

---------

Co-authored-by: Tomer Bin <117278227+hxtomer@users.noreply.github.com>
This commit is contained in:
Krish Dholakia 2025-01-14 16:18:21 -08:00 committed by GitHub
parent 293bb4c522
commit 29663c2db5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 408 additions and 71 deletions

View file

@ -161,6 +161,18 @@ def init_guardrails_v2( # noqa: PLR0915
category_thresholds=litellm_params.get("category_thresholds"),
)
litellm.callbacks.append(_lakera_callback) # type: ignore
elif litellm_params["guardrail"] == SupportedGuardrailIntegrations.AIM.value:
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"],
)
litellm.callbacks.append(_aim_callback) # type: ignore
elif (
litellm_params["guardrail"] == SupportedGuardrailIntegrations.PRESIDIO.value
):