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

1from typing import TYPE_CHECKING, Final 

2 

3from litellm.types.guardrails import SupportedGuardrailIntegrations 

4 

5from .prompt_security import PromptSecurityGuardrail 

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.prompt_security import ( 

14 PromptSecurityGuardrail, 

15 ) 

16 

17 _prompt_security_callback: Final = PromptSecurityGuardrail( 

18 api_base=litellm_params.api_base, 

19 api_key=litellm_params.api_key, 

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

21 event_hook=litellm_params.mode, 

22 default_on=litellm_params.default_on, 

23 streaming_transform_mode=getattr(litellm_params, "streaming_transform_mode", None), 

24 file_sanitization_fail_open=getattr(litellm_params, "file_sanitization_fail_open", None), 

25 block_on_file_modify=getattr(litellm_params, "block_on_file_modify", None), 

26 ) 

27 litellm.logging_callback_manager.add_litellm_callback(_prompt_security_callback) 

28 

29 return _prompt_security_callback 

30 

31 

32guardrail_initializer_registry: Final = { 

33 SupportedGuardrailIntegrations.PROMPT_SECURITY.value: initialize_guardrail, 

34} 

35 

36 

37guardrail_class_registry: Final = { 

38 SupportedGuardrailIntegrations.PROMPT_SECURITY.value: PromptSecurityGuardrail, 

39}