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

fix linting

Browse source
This commit is contained in:
Ishaan Jaff 2025-03-12 18:44:51 -07:00
parent a28ac57233
commit ed68ad7775
6 changed files with 19 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

1
enterprise/enterprise_hooks/google_text_moderation.py
View file

@ -94,6 +94,7 @@ class _ENTERPRISE_GoogleTextModeration(CustomLogger):
"image_generation",
"moderation",
"audio_transcription",
"responses",
],
):
"""

1
enterprise/enterprise_hooks/llama_guard.py
View file

@ -107,6 +107,7 @@ class _ENTERPRISE_LlamaGuard(CustomLogger):
"image_generation",
"moderation",
"audio_transcription",
"responses",
],
):
"""

1
enterprise/enterprise_hooks/llm_guard.py
View file

@ -126,6 +126,7 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
"image_generation",
"moderation",
"audio_transcription",
"responses",
],
):
"""

1
litellm/proxy/example_config_yaml/custom_guardrail.py
View file

@ -66,6 +66,7 @@ class myCustomGuardrail(CustomGuardrail):
"image_generation",
"moderation",
"audio_transcription",
"responses",
],
):
"""

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

@ -25,8 +25,12 @@ 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 = (
@ -34,7 +38,9 @@ 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(
@ -68,6 +74,7 @@ class AimGuardrail(CustomGuardrail):
"image_generation",
"moderation",
"audio_transcription",
"responses",
],
) -> Union[Exception, str, dict, None]:
verbose_proxy_logger.debug("Inside AIM Moderation Hook")
@ -77,9 +84,10 @@ class AimGuardrail(CustomGuardrail):
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}", "x-aim-litellm-hook": hook} | (
{"x-aim-user-email": user_email} if user_email else {}
)
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(
f"{self.api_base}/detect/openai",
headers=headers,

1
litellm/proxy/guardrails/guardrail_hooks/custom_guardrail.py
View file

@ -70,6 +70,7 @@ class myCustomGuardrail(CustomGuardrail):
"image_generation",
"moderation",
"audio_transcription",
"responses",
],
):
"""