Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py: 40%

1from __future__ import annotations 

2 

3from typing import TYPE_CHECKING, Final 

4 

5from litellm.types.guardrails import ( 

6 GuardrailEventHooks, 

7 Mode, 

8 SupportedGuardrailIntegrations, 

9) 

10 

11from .headroom import HeadroomGuardrail 

12 

13if TYPE_CHECKING: 13 ↛ 14line 13 didn't jump to line 14 because the condition on line 13 was never true

14 from litellm.types.guardrails import Guardrail, LitellmParams 

15 

16 

17def _coerce_event_hook( 

18 mode: str | list[str] | Mode, 

19) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode: 

20 if isinstance(mode, Mode): 

21 return mode 

22 if isinstance(mode, list): 

23 return [GuardrailEventHooks(item) for item in mode] 

24 return GuardrailEventHooks(mode) 

25 

26 

27def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> HeadroomGuardrail: 

28 import litellm 

29 

30 _callback: Final = HeadroomGuardrail( 

31 api_base=litellm_params.api_base, 

32 api_key=litellm_params.api_key, 

33 model=litellm_params.model, 

34 guardrail_name=guardrail["guardrail_name"], 

35 event_hook=_coerce_event_hook(litellm_params.mode), 

36 default_on=litellm_params.default_on or False, 

37 unreachable_fallback=litellm_params.unreachable_fallback, 

38 timeout=litellm_params.timeout, 

39 ccr_retrieval=litellm_params.ccr_retrieval, 

40 ) 

41 litellm.logging_callback_manager.add_litellm_callback( # pyright: ignore[reportUnknownMemberType] 

42 _callback 

43 ) 

44 return _callback 

45 

46 

47guardrail_initializer_registry: Final = { 

48 SupportedGuardrailIntegrations.HEADROOM.value: initialize_guardrail, 

49} 

50 

51guardrail_class_registry: Final = { 

52 SupportedGuardrailIntegrations.HEADROOM.value: HeadroomGuardrail, 

53}