diff --git a/enterprise/enterprise_hooks/google_text_moderation.py b/enterprise/enterprise_hooks/google_text_moderation.py index af5ea35987..fe26a03207 100644 --- a/enterprise/enterprise_hooks/google_text_moderation.py +++ b/enterprise/enterprise_hooks/google_text_moderation.py @@ -94,6 +94,7 @@ class _ENTERPRISE_GoogleTextModeration(CustomLogger): "image_generation", "moderation", "audio_transcription", + "responses", ], ): """ diff --git a/enterprise/enterprise_hooks/llama_guard.py b/enterprise/enterprise_hooks/llama_guard.py index 8abbc996d3..2c53fafa5b 100644 --- a/enterprise/enterprise_hooks/llama_guard.py +++ b/enterprise/enterprise_hooks/llama_guard.py @@ -107,6 +107,7 @@ class _ENTERPRISE_LlamaGuard(CustomLogger): "image_generation", "moderation", "audio_transcription", + "responses", ], ): """ diff --git a/enterprise/enterprise_hooks/llm_guard.py b/enterprise/enterprise_hooks/llm_guard.py index 1b639b8a08..078b8e216e 100644 --- a/enterprise/enterprise_hooks/llm_guard.py +++ b/enterprise/enterprise_hooks/llm_guard.py @@ -126,6 +126,7 @@ class _ENTERPRISE_LLMGuard(CustomLogger): "image_generation", "moderation", "audio_transcription", + "responses", ], ): """ diff --git a/litellm/proxy/example_config_yaml/custom_guardrail.py b/litellm/proxy/example_config_yaml/custom_guardrail.py index abd5b672cb..5a5c784410 100644 --- a/litellm/proxy/example_config_yaml/custom_guardrail.py +++ b/litellm/proxy/example_config_yaml/custom_guardrail.py @@ -66,6 +66,7 @@ class myCustomGuardrail(CustomGuardrail): "image_generation", "moderation", "audio_transcription", + "responses", ], ): """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/aim.py b/litellm/proxy/guardrails/guardrail_hooks/aim.py index 91d19e277c..cdc5f00963 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aim.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aim.py @@ -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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_guardrail.py index a45343b37d..87860477f0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_guardrail.py @@ -70,6 +70,7 @@ class myCustomGuardrail(CustomGuardrail): "image_generation", "moderation", "audio_transcription", + "responses", ], ): """