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

1from typing import TYPE_CHECKING, Final 

2 

3from litellm.types.guardrails import SupportedGuardrailIntegrations 

4 

5from .model_armor import ModelArmorGuardrail 

6 

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

8 from litellm.types.guardrails import Guardrail, LitellmParams 

9 

10 

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

12 import litellm 

13 from litellm.proxy.guardrails.guardrail_hooks.model_armor import ( 

14 ModelArmorGuardrail, 

15 ) 

16 

17 _model_armor_callback: Final = ModelArmorGuardrail( 

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

19 event_hook=litellm_params.mode, 

20 template_id=litellm_params.template_id, 

21 project_id=litellm_params.project_id, 

22 location=litellm_params.location, 

23 credentials=litellm_params.credentials, 

24 api_endpoint=litellm_params.api_endpoint, 

25 default_on=litellm_params.default_on, 

26 mask_request_content=litellm_params.mask_request_content, 

27 mask_response_content=litellm_params.mask_response_content, 

28 fail_on_error=litellm_params.fail_on_error, 

29 skip_unscannable_attachments=litellm_params.skip_unscannable_attachments, 

30 sanitize_error_detail=litellm_params.sanitize_error_detail, 

31 ) 

32 litellm.logging_callback_manager.add_litellm_callback(_model_armor_callback) 

33 

34 return _model_armor_callback 

35 

36 

37guardrail_initializer_registry: Final = { 

38 SupportedGuardrailIntegrations.MODEL_ARMOR.value: initialize_guardrail, 

39} 

40 

41 

42guardrail_class_registry: Final = { 

43 SupportedGuardrailIntegrations.MODEL_ARMOR.value: ModelArmorGuardrail, 

44}