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

1from typing import TYPE_CHECKING, Final 

2 

3from litellm.types.guardrails import SupportedGuardrailIntegrations 

4 

5from .panw_prisma_airs import PanwPrismaAirsHandler 

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 guardrail_name: Final = guardrail.get("guardrail_name") 

15 

16 # Note: api_key and profile_name can be None - handler will use env vars or API key's linked profile 

17 if not guardrail_name: 

18 raise ValueError("PANW Prisma AIRS: guardrail_name is required") 

19 

20 _panw_callback: Final = PanwPrismaAirsHandler( 

21 **{ 

22 **litellm_params.model_dump(exclude_unset=True), 

23 "guardrail_name": guardrail_name, 

24 "event_hook": litellm_params.mode, 

25 "default_on": litellm_params.default_on or False, 

26 } 

27 ) 

28 litellm.logging_callback_manager.add_litellm_callback(_panw_callback) 

29 

30 return _panw_callback 

31 

32 

33guardrail_initializer_registry: Final = { 

34 SupportedGuardrailIntegrations.PANW_PRISMA_AIRS.value: initialize_guardrail, 

35} 

36 

37 

38guardrail_class_registry: Final = { 

39 SupportedGuardrailIntegrations.PANW_PRISMA_AIRS.value: PanwPrismaAirsHandler, 

40}