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

1from typing import TYPE_CHECKING, Final 

2 

3from litellm.types.guardrails import SupportedGuardrailIntegrations 

4 

5from .noma import NomaGuardrail 

6from .noma_v2 import NomaV2Guardrail 

7 

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

9 from litellm.types.guardrails import Guardrail, LitellmParams 

10 

11 

12def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): 

13 use_v2 = getattr(litellm_params, "use_v2", False) 

14 if isinstance(use_v2, str): 

15 use_v2 = use_v2.lower() == "true" 

16 if use_v2: 

17 return initialize_guardrail_v2(litellm_params=litellm_params, guardrail=guardrail) 

18 

19 import litellm 

20 

21 _noma_callback: Final = NomaGuardrail( 

22 guardrail_name=guardrail.get("guardrail_name", ""), 

23 api_key=litellm_params.api_key, 

24 api_base=litellm_params.api_base, 

25 application_id=litellm_params.application_id, 

26 monitor_mode=litellm_params.monitor_mode, 

27 block_failures=litellm_params.block_failures, 

28 anonymize_input=litellm_params.anonymize_input, 

29 event_hook=litellm_params.mode, 

30 default_on=litellm_params.default_on, 

31 ) 

32 litellm.logging_callback_manager.add_litellm_callback(_noma_callback) 

33 

34 return _noma_callback 

35 

36 

37def initialize_guardrail_v2(litellm_params: "LitellmParams", guardrail: "Guardrail"): 

38 import litellm 

39 

40 _noma_v2_callback: Final = NomaV2Guardrail( 

41 guardrail_name=guardrail.get("guardrail_name", ""), 

42 api_key=litellm_params.api_key, 

43 api_base=litellm_params.api_base, 

44 application_id=litellm_params.application_id, 

45 monitor_mode=litellm_params.monitor_mode, 

46 block_failures=litellm_params.block_failures, 

47 event_hook=litellm_params.mode, 

48 default_on=litellm_params.default_on, 

49 ) 

50 litellm.logging_callback_manager.add_litellm_callback(_noma_v2_callback) 

51 

52 return _noma_v2_callback 

53 

54 

55guardrail_initializer_registry: Final = { 

56 SupportedGuardrailIntegrations.NOMA.value: initialize_guardrail, 

57 SupportedGuardrailIntegrations.NOMA_V2.value: initialize_guardrail_v2, 

58} 

59 

60 

61guardrail_class_registry: Final = { 

62 SupportedGuardrailIntegrations.NOMA.value: NomaGuardrail, 

63 SupportedGuardrailIntegrations.NOMA_V2.value: NomaV2Guardrail, 

64}