Coverage for mindsdb / integrations / utilities / handler_utils.py: 52%

1import os 

2from typing import Dict 

3 

4from mindsdb.interfaces.storage.model_fs import HandlerStorage 

5from mindsdb.utilities.config import Config 

6 

7"""Contains utilities to be used by handlers.""" 

8 

9 

10def get_api_key( 

11 api_name: str, 

12 create_args: Dict[str, str], 

13 engine_storage: HandlerStorage = None, 

14 strict: bool = True, 

15): 

16 """Gets the API key needed to use an ML Handler. 

17 

18 Args: 

19 api_name (str): Name of the API (e.g. openai, anthropic) 

20 create_args (Dict[str, str]): Args user passed to the created model with USING keyword 

21 engine_storage (HandlerStorage): Engine storage for the ML handler 

22 strict (bool): Whether or not to require the API key 

23 

24 Returns: 

25 api_key (str): The API key 

26 

27 API_KEY preference order: 

28 1. provided at inference 

29 2. provided at model creation 

30 3. provided at engine creation 

31 4. api key env variable 

32 5. api_key setting in config.json 

33 """ 

34 # Special case for vLLM - always return dummy key 

35 if api_name == "vllm": 35 ↛ 36line 35 didn't jump to line 36 because the condition on line 35 was never true

36 return "EMPTY" 

37 

38 # 1 

39 if "using" in create_args and f"{api_name.lower()}_api_key" in create_args["using"]: 

40 api_key = create_args["using"][f"{api_name.lower()}_api_key"] 

41 if api_key: 41 ↛ 45line 41 didn't jump to line 45 because the condition on line 41 was always true

42 return api_key 

43 

44 # 1.5 - Check for generic api_key in using 

45 if "using" in create_args and "api_key" in create_args["using"]: 

46 api_key = create_args["using"]["api_key"] 

47 if api_key: 47 ↛ 51line 47 didn't jump to line 51 because the condition on line 47 was always true

48 return api_key 

49 

50 # 2 

51 if f"{api_name.lower()}_api_key" in create_args: 

52 api_key = create_args[f"{api_name.lower()}_api_key"] 

53 if api_key: 53 ↛ 57line 53 didn't jump to line 57 because the condition on line 53 was always true

54 return api_key 

55 

56 # 2.5 - Check for generic api_key 

57 if "api_key" in create_args: 

58 api_key = create_args["api_key"] 

59 if api_key: 59 ↛ 63line 59 didn't jump to line 63 because the condition on line 59 was always true

60 return api_key 

61 

62 # 3 - Check in params dictionary if it exists (for agents) 

63 if "params" in create_args and create_args["params"] is not None: 

64 if f"{api_name.lower()}_api_key" in create_args["params"]: 

65 api_key = create_args["params"][f"{api_name.lower()}_api_key"] 

66 if api_key: 66 ↛ 69line 66 didn't jump to line 69 because the condition on line 66 was always true

67 return api_key 

68 # 3.5 - Check for generic api_key in params 

69 if "api_key" in create_args["params"]: 69 ↛ 75line 69 didn't jump to line 75 because the condition on line 69 was always true

70 api_key = create_args["params"]["api_key"] 

71 if api_key: 71 ↛ 75line 71 didn't jump to line 75 because the condition on line 71 was always true

72 return api_key 

73 

74 # 4 

75 if engine_storage is not None: 75 ↛ 88line 75 didn't jump to line 88 because the condition on line 75 was always true

76 connection_args = engine_storage.get_connection_args() 

77 if f"{api_name.lower()}_api_key" in connection_args: 77 ↛ 82line 77 didn't jump to line 82 because the condition on line 77 was always true

78 api_key = connection_args[f"{api_name.lower()}_api_key"] 

79 if api_key: 79 ↛ 82line 79 didn't jump to line 82 because the condition on line 79 was always true

80 return api_key 

81 # 4.5 - Check for generic api_key in connection_args 

82 if "api_key" in connection_args: 

83 api_key = connection_args["api_key"] 

84 if api_key: 

85 return api_key 

86 

87 # 5 

88 api_key = os.getenv(f"{api_name.lower()}_api_key") 

89 if api_key: 

90 return api_key 

91 api_key = os.getenv(f"{api_name.upper()}_API_KEY") 

92 if api_key: 

93 return api_key 

94 

95 # 6 

96 config = Config() 

97 api_cfg = config.get(api_name, {}) 

98 if f"{api_name.lower()}_api_key" in api_cfg: 

99 api_key = api_cfg[f"{api_name.lower()}_api_key"] 

100 if api_key: 

101 return api_key 

102 

103 # 7 

104 if "api_keys" in create_args and api_name in create_args["api_keys"]: 

105 api_key = create_args["api_keys"][api_name] 

106 if api_key: 

107 return api_key 

108 

109 if strict: 

110 provider_upper = api_name.upper() 

111 api_key_env_var = f"{provider_upper}_API_KEY" 

112 api_key_arg = f"{api_name.lower()}_api_key" 

113 error_message = ( 

114 f"API key for {api_name} not found. Please provide it using one of the following methods:\n" 

115 f"1. Set the {api_key_env_var} environment variable\n" 

116 f"2. Provide it as '{api_key_arg}' parameter or 'api_key' parameter when creating an agent using the CREATE AGENT syntax\n" 

117 f" Example: CREATE AGENT my_agent USING model='gpt-4', provider='{api_name}', {api_key_arg}='your-api-key';\n" 

118 f" Or: CREATE AGENT my_agent USING model='gpt-4', provider='{api_name}', api_key='your-api-key';\n" 

119 ) 

120 raise Exception(error_message) 

121 return None