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

1from typing import Optional, Dict 

2 

3import unify 

4import pandas as pd 

5 

6from mindsdb.integrations.libs.base import BaseMLEngine 

7 

8from mindsdb.utilities import log 

9 

10from mindsdb.integrations.utilities.handler_utils import get_api_key 

11 

12 

13logger = log.getLogger(__name__) 

14 

15 

16class UnifyHandler(BaseMLEngine): 

17 """ 

18 Integration with the Unifyai Python Library 

19 """ 

20 name = 'unify' 

21 

22 def create(self, target: str, df: Optional[pd.DataFrame] = None, args: Optional[Dict] = None) -> None: 

23 if 'using' not in args: 

24 raise Exception("Unify engine requires a USING clause! Refer to its documentation for more details.") 

25 

26 self.generative = True 

27 self.model_storage.json_set('args', args) 

28 

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

30 

31 args = self.model_storage.json_get('args') 

32 

33 input_keys = list(args.keys()) 

34 

35 logger.info(f"Input keys: {input_keys}!") 

36 

37 if 'model' not in args['using']: 

38 raise Exception("Unify requires an model parameter in the USING clause! Refer to its documentation for more details.") 

39 model = args['using']['model'] 

40 

41 if 'provider' not in args['using']: 

42 raise Exception("Unify requires a provider parameter in the USING clause! Refer to its documentation for more details.") 

43 provider = args['using']['provider'] 

44 

45 self.endpoint = model + '@' + provider 

46 self.api_key = get_api_key('unify', args["using"], self.engine_storage, strict=False) 

47 available_endpoints = unify.utils.list_endpoints(api_key=self.api_key) 

48 if self.endpoint not in available_endpoints: 

49 raise Exception("The model, provider or their combination is not supported by Unify! The supported endpoints are: " + str(available_endpoints)) 

50 

51 question_column = args['using']['column'] 

52 if question_column not in df.columns: 

53 raise RuntimeError(f'Column "{question_column}" not found in input data') 

54 

55 result_df = pd.DataFrame() 

56 result_df['predictions'] = df[question_column].apply(self.predict_answer) 

57 result_df = result_df.rename(columns={'predictions': args['target']}) 

58 

59 return result_df 

60 

61 def predict_answer(self, text): 

62 

63 client = unify.Unify(self.endpoint, api_key=self.api_key) 

64 

65 response = client.generate(text) 

66 

67 return response