mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 11:14:04 +00:00
0d027b22fd
* allow using os.environ for slack urls * use env vars for webhook urls * fix types for get_secret * fix linting * fix linting * fix linting * linting fixes * linting fix * docs alerting slack * fix get data
42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
from typing import Literal, Optional
|
|
|
|
from litellm._logging import verbose_logger
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
|
|
class CustomGuardrail(CustomLogger):
|
|
|
|
def __init__(
|
|
self,
|
|
guardrail_name: Optional[str] = None,
|
|
event_hook: Optional[GuardrailEventHooks] = None,
|
|
**kwargs
|
|
):
|
|
self.guardrail_name = guardrail_name
|
|
self.event_hook: Optional[GuardrailEventHooks] = event_hook
|
|
super().__init__(**kwargs)
|
|
|
|
# older v1 implementation - not used, just kept for backward compatibility
|
|
self.moderation_check: Literal["pre_call", "in_parallel"] = kwargs.get(
|
|
"moderation_check", "pre_call"
|
|
)
|
|
|
|
def should_run_guardrail(self, data, event_type: GuardrailEventHooks) -> bool:
|
|
metadata = data.get("metadata") or {}
|
|
requested_guardrails = metadata.get("guardrails") or []
|
|
verbose_logger.debug(
|
|
"inside should_run_guardrail for guardrail=%s event_type= %s guardrail_supported_event_hooks= %s requested_guardrails= %s",
|
|
self.guardrail_name,
|
|
event_type,
|
|
self.event_hook,
|
|
requested_guardrails,
|
|
)
|
|
|
|
if self.guardrail_name not in requested_guardrails:
|
|
return False
|
|
|
|
if self.event_hook != event_type:
|
|
return False
|
|
|
|
return True
|