Coverage for mindsdb / integrations / handlers / minds_endpoint_handler / minds_endpoint_handler.py: 0%

1import os 

2import pandas as pd 

3import openai 

4from openai import OpenAI, NotFoundError, AuthenticationError 

5from typing import Dict, Optional 

6from mindsdb.integrations.handlers.openai_handler import Handler as OpenAIHandler 

7from mindsdb.integrations.utilities.handler_utils import get_api_key 

8from mindsdb.integrations.handlers.minds_endpoint_handler.settings import minds_endpoint_handler_config 

9from mindsdb.utilities import log 

10 

11logger = log.getLogger(__name__) 

12 

13 

14class MindsEndpointHandler(OpenAIHandler): 

15 """ 

16 This handler handles connection to the Minds Endpoint. 

17 """ 

18 

19 name = 'minds_endpoint' 

20 

21 def __init__(self, *args, **kwargs): 

22 super().__init__(*args, **kwargs) 

23 self.api_base = minds_endpoint_handler_config.BASE_URL 

24 self.default_model = 'gpt-3.5-turbo' 

25 self.default_mode = 'default' 

26 

27 @staticmethod 

28 def _check_client_connection(client: OpenAI): 

29 """ 

30 Check the Minds Endpoint engine client connection by listing models. 

31 

32 Args: 

33 client (OpenAI): OpenAI client configured with the Minds Endpoint API credentials. 

34 

35 Raises: 

36 Exception: If the client connection (API key) is invalid. 

37 

38 Returns: 

39 None 

40 """ 

41 try: 

42 client.models.list() 

43 except NotFoundError: 

44 pass 

45 except AuthenticationError as e: 

46 if e.body['code'] == 401: 

47 raise Exception('Invalid api key') 

48 raise Exception(f'Something went wrong: {e}') 

49 

50 def create_engine(self, connection_args): 

51 """ 

52 Validate the Minds Endpoint API credentials on engine creation. 

53 

54 Args: 

55 connection_args (dict): Connection arguments. 

56 

57 Raises: 

58 Exception: If the handler is not configured with valid API credentials. 

59 

60 Returns: 

61 None 

62 """ 

63 connection_args = {k.lower(): v for k, v in connection_args.items()} 

64 api_key = connection_args.get('minds_endpoint_api_key') 

65 if api_key is not None: 

66 org = connection_args.get('api_organization') 

67 api_base = connection_args.get('api_base') or os.environ.get('minds_endpoint_BASE', minds_endpoint_handler_config.BASE_URL) 

68 client = self._get_client(api_key=api_key, base_url=api_base, org=org) 

69 MindsEndpointHandler._check_client_connection(client) 

70 

71 @staticmethod 

72 def create_validation(target, args=None, **kwargs): 

73 """ 

74 Validate the Minds Endpoint API credentials on model creation. 

75 

76 Args: 

77 target (str): Target column, not required for LLMs. 

78 args (dict): Handler arguments. 

79 kwargs (dict): Handler keyword arguments. 

80 

81 Raises: 

82 Exception: If the handler is not configured with valid API credentials. 

83 

84 Returns: 

85 None 

86 """ 

87 if 'using' not in args: 

88 raise Exception( 

89 "Minds Endpoint engine requires a USING clause! Refer to its documentation for more details." 

90 ) 

91 else: 

92 args = args['using'] 

93 

94 engine_storage = kwargs['handler_storage'] 

95 connection_args = engine_storage.get_connection_args() 

96 api_key = get_api_key('minds_endpoint', args, engine_storage=engine_storage) 

97 api_base = connection_args.get('api_base') or args.get('api_base') or os.environ.get('minds_endpoint_BASE', minds_endpoint_handler_config.BASE_URL) 

98 org = args.get('api_organization') 

99 client = OpenAIHandler._get_client(api_key=api_key, base_url=api_base, org=org) 

100 MindsEndpointHandler._check_client_connection(client) 

101 

102 def predict(self, df: pd.DataFrame, args: Optional[Dict] = None) -> pd.DataFrame: 

103 """ 

104 Call the Minds Endpoint engine to predict the next token. 

105 

106 Args: 

107 df (pd.DataFrame): Input data. 

108 args (dict): Handler arguments. 

109 

110 Returns: 

111 pd.DataFrame: Predicted data 

112 """ 

113 api_key = get_api_key('minds_endpoint', args, self.engine_storage) 

114 supported_models = self._get_supported_models(api_key, self.api_base) 

115 self.chat_completion_models = [model.id for model in supported_models] 

116 return super().predict(df, args) 

117 

118 @staticmethod 

119 def _get_supported_models(api_key, base_url, org=None): 

120 """ 

121 Get the list of supported models for the Minds Endpoint engine. 

122 

123 Args: 

124 api_key (str): API key. 

125 base_url (str): Base URL. 

126 org (str): Organization name. 

127 

128 Returns: 

129 List: List of supported models. 

130 """ 

131 client = openai.OpenAI(api_key=api_key, base_url=base_url, organization=org) 

132 return client.models.list()