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

1"""Block Code Execution guardrail: blocks or masks fenced code blocks by language.""" 

2 

3from typing import TYPE_CHECKING, Final, Literal, cast 

4 

5from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations 

6 

7from .block_code_execution import BlockCodeExecutionGuardrail 

8 

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

10 from litellm.types.guardrails import Guardrail, LitellmParams 

11 

12# Default: run on both request and response (and during_call is supported too) 

13DEFAULT_EVENT_HOOKS: Final = [ 

14 GuardrailEventHooks.pre_call.value, 

15 GuardrailEventHooks.post_call.value, 

16] 

17 

18 

19def _get_param( 

20 litellm_params: "LitellmParams", 

21 guardrail: "Guardrail", 

22 key: str, 

23 default: object = None, 

24) -> object: 

25 """Get a param from litellm_params, with fallback to raw guardrail litellm_params (for extra fields not on LitellmParams).""" 

26 value: Final = getattr(litellm_params, key, default) 

27 if value is not None: 

28 return value 

29 raw: Final = guardrail.get("litellm_params") 

30 if isinstance(raw, dict) and key in raw: 

31 return raw[key] 

32 return default 

33 

34 

35def initialize_guardrail( 

36 litellm_params: "LitellmParams", 

37 guardrail: "Guardrail", 

38) -> BlockCodeExecutionGuardrail: 

39 """Initialize the Block Code Execution guardrail from config.""" 

40 import litellm 

41 

42 guardrail_name: Final = guardrail.get("guardrail_name") 

43 if not guardrail_name: 

44 raise ValueError("Block Code Execution guardrail requires a guardrail_name") 

45 

46 blocked_languages: Final[list[str] | None] = cast( 

47 list[str] | None, 

48 _get_param(litellm_params, guardrail, "blocked_languages"), 

49 ) 

50 action: Final = cast( 

51 Literal["block", "mask"], 

52 _get_param(litellm_params, guardrail, "action", "block"), 

53 ) 

54 confidence_threshold: Final = float( 

55 cast( 

56 int | float | str, 

57 _get_param(litellm_params, guardrail, "confidence_threshold", 0.5), 

58 ) 

59 ) 

60 detect_execution_intent: Final = bool(_get_param(litellm_params, guardrail, "detect_execution_intent", True)) 

61 mode: Final = _get_param(litellm_params, guardrail, "mode") 

62 event_hook: Final = cast( 

63 Literal["pre_call", "post_call", "during_call"] | list[str] | None, 

64 mode if mode is not None else DEFAULT_EVENT_HOOKS, 

65 ) 

66 

67 instance: Final = BlockCodeExecutionGuardrail( 

68 guardrail_name=guardrail_name, 

69 blocked_languages=blocked_languages, 

70 action=action, 

71 confidence_threshold=confidence_threshold, 

72 detect_execution_intent=detect_execution_intent, 

73 event_hook=event_hook, 

74 default_on=bool(_get_param(litellm_params, guardrail, "default_on", False)), 

75 ) 

76 litellm.logging_callback_manager.add_litellm_callback(instance) 

77 return instance 

78 

79 

80guardrail_initializer_registry: Final = { 

81 SupportedGuardrailIntegrations.BLOCK_CODE_EXECUTION.value: initialize_guardrail, 

82} 

83 

84guardrail_class_registry: Final = { 

85 SupportedGuardrailIntegrations.BLOCK_CODE_EXECUTION.value: BlockCodeExecutionGuardrail, 

86} 

87 

88__all__ = [ 

89 "BlockCodeExecutionGuardrail", 

90 "initialize_guardrail", 

91]