Coverage for mindsdb / interfaces / chatbot / model_executor.py: 16%

1from typing import List 

2import datetime as dt 

3import pandas as pd 

4 

5from mindsdb.interfaces.storage import db 

6 

7from .types import BotException 

8from .types import ChatBotMessage 

9 

10 

11class ModelExecutor: 

12 

13 def __init__(self, chat_task, model_name): 

14 self.chat_task = chat_task 

15 

16 model = chat_task.session.model_controller.get_model( 

17 model_name, 

18 project_name=chat_task.project_name 

19 ) 

20 model_record = db.Predictor.query.get(model['id']) 

21 integration_record = db.Integration.query.get(model_record.integration_id) 

22 

23 self.model_info = { 

24 'model_name': model_name, 

25 'mode': model_record.learn_args['using'].get('mode'), 

26 'user_column': model_record.learn_args['using'].get('user_column'), 

27 'bot_column': model_record.learn_args['using'].get('assistant_column'), 

28 'output': model_record.to_predict[0], 

29 'engine': integration_record.engine, 

30 } 

31 

32 # redefined prompt 

33 self.prompt = None 

34 

35 def call(self, history: List[ChatBotMessage], functions): 

36 model_info = self.model_info 

37 

38 if model_info['mode'] != 'conversational': 

39 raise BotException('Not supported') 

40 

41 messages, user_name = self._chat_history_to_conversation(history, model_info) 

42 if model_info['engine'] == 'langchain': 

43 

44 all_tools = [] 

45 for function in functions: 

46 all_tools.append({ 

47 'name': function.name, 

48 'func': function.callback, 

49 'description': function.description 

50 }) 

51 

52 context_list = [ 

53 f"- Today's date is {dt.datetime.now().strftime('%Y-%m-%d')}." 

54 f" It must be used to understand the input date from string like 'tomorrow', 'today', 'yesterday'" 

55 ] 

56 context = '\n'.join(context_list) 

57 

58 # call model 

59 params = {'tools': all_tools, 'context': context, 'prompt': self.prompt} 

60 

61 predictions = self.chat_task.project_datanode.predict( 

62 model_name=model_info['model_name'], 

63 df=pd.DataFrame(messages), 

64 params=params 

65 ) 

66 

67 else: 

68 predictions = self.chat_task.project_datanode.predict( 

69 model_name=model_info['model_name'], 

70 df=pd.DataFrame(messages), 

71 params={'prompt': self.prompt, 'user_info': {'user_name': user_name}} 

72 ) 

73 

74 output_col = model_info['output'] 

75 model_output = predictions.iloc[-1][output_col] 

76 return model_output 

77 

78 def _chat_history_to_conversation(self, history, model_info): 

79 

80 bot_username = self.chat_task.bot_params['bot_username'] 

81 user_name = None 

82 

83 question_col = model_info['user_column'] 

84 answer_col = model_info['bot_column'] 

85 

86 messages = [] 

87 

88 for message in history: 

89 text = message.text 

90 

91 if text is None or text.strip() == '': 

92 # skip empty rows 

93 continue 

94 

95 if message.user != bot_username: 

96 user_name = message.user 

97 

98 # create new message row 

99 messages.append({question_col: text, answer_col: None}) 

100 else: 

101 if len(messages) == 0: 

102 # add empty row 

103 messages.append({question_col: None, answer_col: None}) 

104 

105 # update answer in previous column 

106 messages[-1][answer_col] = text 

107 return messages, user_name