fix linting

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

View file

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

View file

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

View file

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

View file

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

View file

@ -25,8 +25,12 @@ class AimGuardrailMissingSecrets(Exception):
class AimGuardrail(CustomGuardrail): class AimGuardrail(CustomGuardrail):
def __init__(self, api_key: Optional[str] = None, api_base: Optional[str] = None, **kwargs): def __init__(
self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) 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") self.api_key = api_key or os.environ.get("AIM_API_KEY")
if not self.api_key: if not self.api_key:
msg = ( msg = (
@ -34,7 +38,9 @@ class AimGuardrail(CustomGuardrail):
"pass it as a parameter to the guardrail in the config file" "pass it as a parameter to the guardrail in the config file"
) )
raise AimGuardrailMissingSecrets(msg) 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) super().__init__(**kwargs)
async def async_pre_call_hook( async def async_pre_call_hook(
@ -68,6 +74,7 @@ class AimGuardrail(CustomGuardrail):
"image_generation", "image_generation",
"moderation", "moderation",
"audio_transcription", "audio_transcription",
"responses",
], ],
) -> Union[Exception, str, dict, None]: ) -> Union[Exception, str, dict, None]:
verbose_proxy_logger.debug("Inside AIM Moderation Hook") 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: async def call_aim_guardrail(self, data: dict, hook: str) -> None:
user_email = data.get("metadata", {}).get("headers", {}).get("x-aim-user-email") user_email = data.get("metadata", {}).get("headers", {}).get("x-aim-user-email")
headers = {"Authorization": f"Bearer {self.api_key}", "x-aim-litellm-hook": hook} | ( headers = {
{"x-aim-user-email": user_email} if user_email else {} "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( response = await self.async_handler.post(
f"{self.api_base}/detect/openai", f"{self.api_base}/detect/openai",
headers=headers, headers=headers,

View file

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