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

1from typing import TYPE_CHECKING, Final 

2 

3from litellm.types.guardrails import ( 

4 GuardrailEventHooks, 

5 Mode, 

6 SupportedGuardrailIntegrations, 

7) 

8 

9from .repelloai import RepelloAIGuardrail 

10 

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

12 from litellm.types.guardrails import Guardrail, LitellmParams 

13 

14 

15def _event_hook_from_mode( 

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

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

18 if isinstance(mode, Mode): 

19 return mode 

20 if isinstance(mode, list): 

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

22 return GuardrailEventHooks(mode) 

23 

24 

25def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> RepelloAIGuardrail: 

26 import litellm 

27 

28 _repelloai_callback: Final = RepelloAIGuardrail( 

29 guardrail_name=guardrail["guardrail_name"], 

30 api_key=litellm_params.api_key, 

31 api_base=litellm_params.api_base, 

32 asset_id=litellm_params.asset_id, 

33 unreachable_fallback=litellm_params.unreachable_fallback, 

34 event_hook=_event_hook_from_mode(litellm_params.mode), 

35 default_on=litellm_params.default_on or False, 

36 ) 

37 litellm.logging_callback_manager.add_litellm_callback(_repelloai_callback) 

38 

39 return _repelloai_callback 

40 

41 

42guardrail_initializer_registry: Final = { 

43 SupportedGuardrailIntegrations.REPELLOAI.value: initialize_guardrail, 

44} 

45 

46 

47guardrail_class_registry: Final = { 

48 SupportedGuardrailIntegrations.REPELLOAI.value: RepelloAIGuardrail, 

49}