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

1import os 

2import sys 

3from typing import Final 

4 

5import litellm 

6from litellm._logging import verbose_proxy_logger 

7from litellm.proxy.proxy_server import LiteLLM_TeamTable, UserAPIKeyAuth 

8from litellm.types.guardrails import * 

9 

10sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path 

11 

12 

13def can_modify_guardrails(team_obj: LiteLLM_TeamTable | None) -> bool: 

14 if team_obj is None: 14 ↛ 17line 14 didn't jump to line 17 because the condition on line 14 was always true

15 return True 

16 

17 team_metadata: Final = team_obj.metadata or {} 

18 

19 if team_metadata.get("guardrails", None) is not None and isinstance(team_metadata.get("guardrails"), dict): 

20 if team_metadata.get("guardrails", {}).get("modify_guardrails", None) is False: 

21 return False 

22 

23 return True 

24 

25 

26async def should_proceed_based_on_metadata(data: dict, guardrail_name: str) -> bool: 

27 """ 

28 checks if this guardrail should be applied to this call 

29 """ 

30 if "metadata" in data and isinstance(data["metadata"], dict): 

31 if "guardrails" in data["metadata"]: 

32 # expect users to pass 

33 # guardrails: { prompt_injection: true, rail_2: false } 

34 request_guardrails: Final = data["metadata"]["guardrails"] 

35 verbose_proxy_logger.debug( 

36 "Guardrails %s passed in request - checking which to apply", 

37 request_guardrails, 

38 ) 

39 

40 requested_callback_names: Final = [] 

41 

42 # v1 implementation of this 

43 if isinstance(request_guardrails, dict): 

44 # get guardrail configs from `init_guardrails.py` 

45 # for all requested guardrails -> get their associated callbacks 

46 for _guardrail_name, should_run in request_guardrails.items(): 

47 if should_run is False: 

48 verbose_proxy_logger.debug( 

49 "Guardrail %s skipped because request set to False", 

50 _guardrail_name, 

51 ) 

52 continue 

53 

54 # lookup the guardrail in guardrail_name_config_map 

55 guardrail_item: GuardrailItem = litellm.guardrail_name_config_map[_guardrail_name] 

56 

57 guardrail_callbacks = guardrail_item.callbacks 

58 requested_callback_names.extend(guardrail_callbacks) 

59 

60 verbose_proxy_logger.debug("requested_callback_names %s", requested_callback_names) 

61 if guardrail_name in requested_callback_names: 

62 return True 

63 

64 # Do no proceeed if - "metadata": { "guardrails": { "lakera_prompt_injection": false } } 

65 return False 

66 

67 return True 

68 

69 

70async def should_proceed_based_on_api_key(user_api_key_dict: UserAPIKeyAuth, guardrail_name: str) -> bool: 

71 """ 

72 checks if this guardrail should be applied to this call 

73 """ 

74 if user_api_key_dict.permissions is not None: 

75 # { prompt_injection: true, rail_2: false } 

76 verbose_proxy_logger.debug( 

77 "Guardrails valid for API Key= %s - checking which to apply", 

78 user_api_key_dict.permissions, 

79 ) 

80 

81 if not isinstance(user_api_key_dict.permissions, dict): 

82 verbose_proxy_logger.error( 

83 "API Key permissions must be a dict - %s running guardrail %s", 

84 user_api_key_dict, 

85 guardrail_name, 

86 ) 

87 return True 

88 

89 for _guardrail_name, should_run in user_api_key_dict.permissions.items(): 

90 if should_run is False: 

91 verbose_proxy_logger.debug( 

92 "Guardrail %s skipped because request set to False", 

93 _guardrail_name, 

94 ) 

95 continue 

96 

97 # lookup the guardrail in guardrail_name_config_map 

98 guardrail_item: GuardrailItem = litellm.guardrail_name_config_map[_guardrail_name] 

99 

100 guardrail_callbacks = guardrail_item.callbacks 

101 if guardrail_name in guardrail_callbacks: 

102 return True 

103 

104 # Do not proceeed if - "metadata": { "guardrails": { "lakera_prompt_injection": false } } 

105 return False 

106 return True