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

1"""Custom code guardrail integration for LiteLLM. 

2 

3This module allows users to write custom guardrail logic using Python-like code 

4that runs in a sandboxed environment with access to LiteLLM-provided primitives. 

5 

6Pre-built custom code for common guardrails (e.g. response rejection detection) 

7is available in response_rejection_code.py. 

8""" 

9 

10from typing import TYPE_CHECKING, Final 

11 

12from litellm.types.guardrails import SupportedGuardrailIntegrations 

13 

14from .custom_code_guardrail import CustomCodeGuardrail 

15from .response_rejection_code import ( 

16 DEFAULT_REJECTION_PHRASES, 

17 RESPONSE_REJECTION_GUARDRAIL_CODE, 

18) 

19 

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

21 from litellm.types.guardrails import Guardrail, LitellmParams 

22 

23 

24def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> CustomCodeGuardrail: 

25 """ 

26 Initialize a custom code guardrail. 

27 

28 Args: 

29 litellm_params: Configuration parameters including the custom code 

30 guardrail: The guardrail configuration dict 

31 

32 Returns: 

33 CustomCodeGuardrail instance 

34 """ 

35 import litellm 

36 

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

38 if not guardrail_name: 

39 raise ValueError("Custom code guardrail requires a guardrail_name") 

40 

41 # Get the custom code from litellm_params 

42 custom_code: Final = getattr(litellm_params, "custom_code", None) 

43 if not custom_code: 

44 raise ValueError("Custom code guardrail requires 'custom_code' in litellm_params") 

45 

46 custom_code_guardrail: Final = CustomCodeGuardrail( 

47 guardrail_name=guardrail_name, 

48 custom_code=custom_code, 

49 event_hook=litellm_params.mode, 

50 default_on=litellm_params.default_on, 

51 ) 

52 

53 litellm.logging_callback_manager.add_litellm_callback(custom_code_guardrail) 

54 return custom_code_guardrail 

55 

56 

57guardrail_initializer_registry: Final = { 

58 SupportedGuardrailIntegrations.CUSTOM_CODE.value: initialize_guardrail, 

59} 

60 

61guardrail_class_registry: Final = { 

62 SupportedGuardrailIntegrations.CUSTOM_CODE.value: CustomCodeGuardrail, 

63} 

64 

65__all__ = [ 

66 "DEFAULT_REJECTION_PHRASES", 

67 "RESPONSE_REJECTION_GUARDRAIL_CODE", 

68 "CustomCodeGuardrail", 

69 "initialize_guardrail", 

70]