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

1from typing import TYPE_CHECKING, Final 

2 

3from litellm.types.guardrails import SupportedGuardrailIntegrations 

4 

5from .zscaler_ai_guard import ZscalerAIGuard 

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 

14 _zscaler_ai_guard_callback: Final = ZscalerAIGuard( 

15 api_base=litellm_params.api_base, 

16 api_key=litellm_params.api_key, 

17 policy_id=litellm_params.policy_id, 

18 send_user_api_key_alias=litellm_params.send_user_api_key_alias, 

19 send_user_api_key_user_id=litellm_params.send_user_api_key_user_id, 

20 send_user_api_key_team_id=litellm_params.send_user_api_key_team_id, 

21 timeout=litellm_params.timeout, 

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

23 event_hook=litellm_params.mode, 

24 default_on=litellm_params.default_on, 

25 ) 

26 litellm.logging_callback_manager.add_litellm_callback(_zscaler_ai_guard_callback) 

27 

28 return _zscaler_ai_guard_callback 

29 

30 

31guardrail_initializer_registry: Final = { 

32 SupportedGuardrailIntegrations.ZSCALER_AI_GUARD.value: initialize_guardrail, 

33} 

34 

35 

36guardrail_class_registry: Final = { 

37 SupportedGuardrailIntegrations.ZSCALER_AI_GUARD.value: ZscalerAIGuard, 

38}