Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/guardrails/init_guardrails.py: 11%

1from typing import Any, Final, cast 

2 

3import litellm 

4from litellm import Router 

5from litellm._logging import verbose_proxy_logger 

6from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy 

7 

8# v2 implementation 

9from litellm.types.guardrails import Guardrail, GuardrailItem, GuardrailItemSpec 

10 

11all_guardrails: list[GuardrailItem] = [] 

12 

13""" 

14Map guardrail_name: <pre_call>, <post_call>, during_call 

15 

16""" 

17 

18 

19def init_guardrails_v2( 

20 all_guardrails: list[dict], 

21 config_file_path: str | None = None, 

22 llm_router: Router | None = None, 

23): 

24 from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER 

25 

26 guardrail_list: Final[list[Guardrail]] = [] 

27 

28 for guardrail in all_guardrails: 

29 try: 

30 initialized_guardrail = IN_MEMORY_GUARDRAIL_HANDLER.initialize_guardrail( 

31 guardrail=cast(Guardrail, guardrail), 

32 config_file_path=config_file_path, 

33 llm_router=llm_router, 

34 source="config", 

35 ) 

36 except (ValueError, TypeError) as init_error: 

37 verbose_proxy_logger.error( 

38 "Skipping guardrail '%s': invalid configuration, proxy is starting WITHOUT this guardrail: %s", 

39 guardrail.get("guardrail_name"), 

40 init_error, 

41 ) 

42 continue 

43 if initialized_guardrail: 

44 guardrail_list.append(initialized_guardrail) 

45 

46 # verbose_proxy_logger.debug(f"\nGuardrail List:{guardrail_list}\n") 

47 

48 # Populate router's guardrail_list for load balancing support 

49 _populate_router_guardrail_list(guardrail_list=guardrail_list) 

50 

51 

52def _populate_router_guardrail_list(guardrail_list: list[Guardrail]) -> None: 

53 """ 

54 Populate the router's guardrail_list from initialized guardrails. 

55 

56 This enables load balancing across multiple guardrail deployments 

57 with the same guardrail_name. 

58 """ 

59 from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER 

60 from litellm.proxy.proxy_server import llm_router 

61 from litellm.types.router import GuardrailTypedDict 

62 

63 if llm_router is None: 

64 verbose_proxy_logger.debug("Router not initialized yet, skipping guardrail_list population") 

65 return 

66 

67 router_guardrail_list: Final[list[GuardrailTypedDict]] = [] 

68 

69 for guardrail in guardrail_list: 

70 guardrail_id = guardrail.get("guardrail_id") 

71 guardrail_name = guardrail.get("guardrail_name") 

72 litellm_params: Any = guardrail.get("litellm_params", {}) 

73 

74 # Get the callback instance from the registry 

75 callback = None 

76 if guardrail_id: 

77 callback = IN_MEMORY_GUARDRAIL_HANDLER.guardrail_id_to_custom_guardrail.get(guardrail_id) 

78 

79 # Build litellm_params dict for the router 

80 params_dict = litellm_params.model_dump() if hasattr(litellm_params, "model_dump") else dict(litellm_params) 

81 

82 router_guardrail: GuardrailTypedDict = GuardrailTypedDict( 

83 guardrail_name=guardrail_name or "", 

84 litellm_params={ 

85 "guardrail": params_dict.get("guardrail", ""), 

86 "mode": params_dict.get("mode", ""), 

87 "api_key": params_dict.get("api_key"), 

88 "api_base": params_dict.get("api_base"), 

89 }, 

90 callback=callback, 

91 id=guardrail_id, 

92 ) 

93 

94 router_guardrail_list.append(router_guardrail) 

95 

96 llm_router.guardrail_list = router_guardrail_list 

97 verbose_proxy_logger.debug("Populated router guardrail_list with %s guardrails", len(router_guardrail_list)) 

98 

99 

100### LEGACY IMPLEMENTATION ### 

101def initialize_guardrails( 

102 guardrails_config: list[dict[str, GuardrailItemSpec]], 

103 premium_user: bool, 

104 config_file_path: str, 

105 litellm_settings: dict, 

106) -> dict[str, GuardrailItem]: 

107 try: 

108 verbose_proxy_logger.debug("validating guardrails passed %s", guardrails_config) 

109 global all_guardrails 

110 for item in guardrails_config: 

111 """ 

112 one item looks like this: 

113 

114 {'prompt_injection': {'callbacks': ['lakera_prompt_injection', 'prompt_injection_api_2'], 'default_on': True, 'enabled_roles': ['user']}} 

115 """ 

116 for k, v in item.items(): 

117 guardrail_item = GuardrailItem(**v, guardrail_name=k) 

118 all_guardrails.append(guardrail_item) 

119 litellm.guardrail_name_config_map[k] = guardrail_item 

120 

121 # set appropriate callbacks if they are default on 

122 default_on_callbacks: Final = set() 

123 callback_specific_params: Final = {} 

124 for guardrail in all_guardrails: 

125 verbose_proxy_logger.debug(guardrail.guardrail_name) 

126 verbose_proxy_logger.debug(guardrail.default_on) 

127 

128 callback_specific_params.update(guardrail.callback_args) 

129 

130 if guardrail.default_on is True: 

131 # add these to litellm callbacks if they don't exist 

132 for callback in guardrail.callbacks: 

133 if callback not in litellm.callbacks: 

134 default_on_callbacks.add(callback) 

135 

136 if guardrail.logging_only is True: 

137 if callback == "presidio": 

138 callback_specific_params["presidio"] = {"logging_only": True} 

139 

140 default_on_callbacks_list: Final = list(default_on_callbacks) 

141 if len(default_on_callbacks_list) > 0: 

142 initialize_callbacks_on_proxy( 

143 value=default_on_callbacks_list, 

144 premium_user=premium_user, 

145 config_file_path=config_file_path, 

146 litellm_settings=litellm_settings, 

147 callback_specific_params=callback_specific_params, 

148 ) 

149 

150 return litellm.guardrail_name_config_map 

151 except Exception as e: 

152 verbose_proxy_logger.exception("error initializing guardrails %s", e) 

153 raise e