phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 18:54:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(utils.py): handle key error in msg validation (#8325)

Browse source 
* fix(utils.py): handle key error in msg validation

* Support running Aim Guard during LLM call (#7918)

* support running Aim Guard during LLM call

* Rename header

* adjust docs and fix type annotations

* fix(timeout.md): doc fix for openai example on dynamic timeouts

---------

Co-authored-by: Tomer Bin <117278227+hxtomer@users.noreply.github.com>
This commit is contained in:
Krish Dholakia 2025-02-06 18:13:46 -08:00 • committed by GitHub
parent fac1d2ccef
commit f031926b82
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 56 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

40
litellm/proxy/guardrails/guardrail_hooks/aim.py
View file

@ -6,7 +6,7 @@
# +-------------------------------------------------------------+
import os
from typing import Literal, Optional
from typing import Literal, Optional, Union
from fastapi import HTTPException
@ -25,12 +25,8 @@ class AimGuardrailMissingSecrets(Exception):
class AimGuardrail(CustomGuardrail):
def __init__(
self, api_key: Optional[str] = None, api_base: Optional[str] = None, **kwargs
):
self.async_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback
)
def __init__(self, api_key: Optional[str] = None, api_base: Optional[str] = None, **kwargs):
self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
self.api_key = api_key or os.environ.get("AIM_API_KEY")
if not self.api_key:
msg = (
@ -38,9 +34,7 @@ class AimGuardrail(CustomGuardrail):
"pass it as a parameter to the guardrail in the config file"
)
raise AimGuardrailMissingSecrets(msg)
self.api_base = (
api_base or os.environ.get("AIM_API_BASE") or "https://api.aim.security"
)
self.api_base = api_base or os.environ.get("AIM_API_BASE") or "https://api.aim.security"
super().__init__(**kwargs)
async def async_pre_call_hook(
@ -58,11 +52,32 @@ class AimGuardrail(CustomGuardrail):
"pass_through_endpoint",
"rerank",
],
) -> Exception | str | dict | None:
) -> Union[Exception, str, dict, None]:
verbose_proxy_logger.debug("Inside AIM Pre-Call Hook")
await self.call_aim_guardrail(data, hook="pre_call")
return data
async def async_moderation_hook(
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
],
) -> Union[Exception, str, dict, None]:
verbose_proxy_logger.debug("Inside AIM Moderation Hook")
await self.call_aim_guardrail(data, hook="moderation")
return data
async def call_aim_guardrail(self, data: dict, hook: str) -> None:
user_email = data.get("metadata", {}).get("headers", {}).get("x-aim-user-email")
headers = {"Authorization": f"Bearer {self.api_key}"} | (
headers = {"Authorization": f"Bearer {self.api_key}", "x-aim-litellm-hook": hook} | (
{"x-aim-user-email": user_email} if user_email else {}
)
response = await self.async_handler.post(
@ -80,4 +95,3 @@ class AimGuardrail(CustomGuardrail):
)
if detected:
raise HTTPException(status_code=400, detail=res["detection_message"])
return data