Satisf ruff checks

This commit is contained in:
AlphaRex-pixel 2026-08-15 16:51:14 +05:30
parent 28bd0b80b3
commit f5e8549c73
2 changed files with 14 additions and 3 deletions

View file

@ -9,7 +9,7 @@ if TYPE_CHECKING:
from litellm.types.guardrails import Guardrail, LitellmParams
def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"):
def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> TealTigerGuardrail:
import litellm
_tealtiger_callback: Final = TealTigerGuardrail(

View file

@ -20,6 +20,7 @@ from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.types.guardrails import GuardrailEventHooks, Mode
from litellm.types.utils import GenericGuardrailAPIInputs
from .engine import Action, PolicyMode, TealEngine
@ -41,13 +42,23 @@ class TealTigerGuardrail(CustomGuardrail):
self,
policies: Sequence[Mapping[str, object]] | None = None,
policy_mode: str = "ENFORCE",
**kwargs, # kwargs-ok: forwards unknown/future CustomGuardrail.__init__ params (guardrail_name, event_hook, default_on, ...) to the base class
guardrail_name: str | None = None,
event_hook: GuardrailEventHooks # mutable-ok: matches CustomGuardrail's actual signature
| list[GuardrailEventHooks]
| Mode
| None = None,
default_on: bool = False,
) -> None:
# Only these 3 base-class params are ever passed by this repo's own
# registry (see __init__.py's initialize_guardrail) — declared
# explicitly, rather than a **kwargs passthrough, so each argument's
# type is checkable at the call to super().__init__() below instead
# of erasing everything to object.
self.engine: Final = TealEngine(
policies=policies or DEFAULT_POLICIES,
mode=PolicyMode(policy_mode),
)
super().__init__(**kwargs)
super().__init__(guardrail_name=guardrail_name, event_hook=event_hook, default_on=default_on)
@log_guardrail_information
async def apply_guardrail(