Coverage for mindsdb / interfaces / chatbot / chatbot_executor.py: 13%

1 

2from mindsdb.interfaces.agents.constants import USER_COLUMN, ASSISTANT_COLUMN 

3 

4from .model_executor import ModelExecutor 

5from .types import Function, BotException 

6 

7 

8class BotExecutor: 

9 def __init__(self, chat_task, chat_memory): 

10 self.chat_task = chat_task 

11 self.chat_memory = chat_memory 

12 

13 def _get_model(self, model_name): 

14 return ModelExecutor(self.chat_task, model_name) 

15 

16 def _prepare_available_functions(self): 

17 

18 # collecting functions 

19 functions = [] 

20 

21 back_db_name = self.chat_task.bot_params.get('backoffice_db') 

22 if back_db_name is not None: 

23 back_db = self.chat_task.session.integration_controller.get_data_handler(back_db_name) 

24 if hasattr(back_db, 'back_office_config'): 

25 back_db_config = back_db.back_office_config() 

26 

27 for name, description in back_db_config.get('tools', {}).items(): 

28 functions.append( 

29 Function( 

30 name=name, 

31 description=description, 

32 callback=getattr(back_db, name) 

33 )) 

34 return functions 

35 

36 def process(self): 

37 # restart of the bot clear previous history 

38 if self.chat_memory.get_mode() is None: 

39 self.chat_memory.hide_history(left_count=1) 

40 self.chat_memory.set_mode('main') 

41 

42 functions = self._prepare_available_functions() 

43 

44 model_executor = self._get_model(self.chat_task.base_model_name) 

45 model_output = model_executor.call(self.chat_memory.get_history(), functions) 

46 return model_output 

47 

48 

49class MultiModeBotExecutor(BotExecutor): 

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

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

52 

53 self._modes = self.chat_task.bot_params['modes'] 

54 

55 def _get_avail_modes_items(self): 

56 return [ 

57 f'- code: {key}, description: {value["info"]}' 

58 for key, value in self._modes.items() 

59 ] 

60 

61 def _make_select_mode_prompt(self): 

62 # select mode tool 

63 task_items = self._get_avail_modes_items() 

64 

65 tasks = '\n'.join(task_items) 

66 

67 prompt = f'You are a helpful assistant and you can help with various types of tasks.' \ 

68 f'\nAvailable types of tasks:' \ 

69 f'\n{tasks}' \ 

70 f'\nUser have to choose task and assistant MUST call select_task function after it' 

71 

72 return prompt 

73 

74 def enter_bot_mode(self, functions): 

75 # choose prompt or model depending on mode 

76 mode_name = self.chat_memory.get_mode() 

77 

78 allowed_tools = None 

79 

80 if mode_name is None: 

81 # mode in not selected, lets to go to select menu 

82 model_executor = self._get_model(self.chat_task.base_model_name) 

83 prompt = self._make_select_mode_prompt() 

84 

85 model_executor.prompt = prompt 

86 

87 else: 

88 # mode is selected 

89 mode = self._modes.get(mode_name) 

90 if mode is None: 

91 # wrong mode 

92 self.chat_memory.set_mode(None) 

93 raise BotException(f'Error to use mode: {mode_name}') 

94 

95 if 'model' in mode: 

96 # this is model 

97 model_executor = self._get_model(mode['model']) 

98 

99 elif 'prompt' in mode: 

100 # it is just a prompt. let's use a bot model and custom prompt 

101 model_executor = self._get_model(self.chat_task.base_model_name) 

102 model_executor.prompt = mode['prompt'] 

103 

104 else: 

105 raise BotException(f'Mode is not supported: {mode}') 

106 

107 allowed_tools = mode.get('allowed_tools') 

108 

109 if allowed_tools is not None: 

110 functions = [ 

111 fnc 

112 for fnc in functions 

113 if fnc.name in allowed_tools 

114 ] 

115 

116 return model_executor, functions 

117 

118 def _mode_switching_function(self, switched_to_mode): 

119 # add mode tool 

120 

121 def _select_task(mode_name): 

122 if mode_name == '': 

123 self.chat_memory.set_mode(None) 

124 switched_to_mode.append(None) 

125 return 'success' 

126 

127 avail_modes = list(self._modes.keys()) 

128 if mode_name not in avail_modes: 

129 return f'Error: task is not found. Available tasks: {", ".join(avail_modes)}' 

130 self.chat_memory.set_mode(mode_name) 

131 switched_to_mode.append(mode_name) 

132 return 'success' 

133 

134 return Function( 

135 name='select_task', 

136 callback=_select_task, 

137 description='Have to be used by assistant to select task. Input is task type.' 

138 ' If user want to unselect task input should be empty string.' 

139 ' Available tasks: ' + '; '.join(self._get_avail_modes_items()) 

140 ) 

141 

142 def process(self): 

143 # this list should be changed if mode was switched 

144 switched_to_mode = [] 

145 

146 functions_all = self._prepare_available_functions() 

147 

148 # Modes handling 

149 functions_all.append(self._mode_switching_function(switched_to_mode)) 

150 

151 model_executor, functions = self.enter_bot_mode(functions_all) 

152 

153 # workaround: don't show history if mode is not selected, otherwise bot doesn't decide to change mode 

154 if self.chat_memory.get_mode() is None: 

155 self.chat_memory.hide_history(left_count=1) 

156 

157 model_output = model_executor.call(self.chat_memory.get_history(), functions) 

158 

159 if len(switched_to_mode) > 0: 

160 # mode changed: 

161 # - clear previous history 

162 # - run once again 

163 

164 # start conversation only from last message 

165 self.chat_memory.hide_history(left_count=1) 

166 

167 model_executor, functions = self.enter_bot_mode(functions_all) 

168 

169 model_output = model_executor.call(self.chat_memory.get_history(), functions) 

170 

171 return model_output 

172 

173 

174class AgentExecutor: 

175 def __init__(self, chat_task, chat_memory): 

176 self.chat_task = chat_task 

177 self.chat_memory = chat_memory 

178 

179 def _chat_history_to_conversation(self, history): 

180 

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

182 

183 messages = [] 

184 

185 for message in history: 

186 text = message.text 

187 

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

189 # skip empty rows 

190 continue 

191 

192 if message.user != bot_username: 

193 # create new message row 

194 messages.append({USER_COLUMN: text, ASSISTANT_COLUMN: None}) 

195 else: 

196 if len(messages) == 0: 

197 # add empty row 

198 messages.append({USER_COLUMN: None, ASSISTANT_COLUMN: None}) 

199 

200 # update answer in previous column 

201 messages[-1][ASSISTANT_COLUMN] = text 

202 return messages 

203 

204 def process(self): 

205 # restart of the bot clear previous history 

206 if self.chat_memory.get_mode() is None: 

207 self.chat_memory.hide_history(left_count=1) 

208 self.chat_memory.set_mode('main') 

209 

210 agents_controller = self.chat_task.session.agents_controller 

211 project_name = self.chat_task.project_name 

212 

213 agent = agents_controller.get_agent_by_id( 

214 self.chat_task.agent_id, 

215 project_name=project_name 

216 ) 

217 

218 messages = self._chat_history_to_conversation(self.chat_memory.get_history()) 

219 predictions = agents_controller.get_completion( 

220 agent, 

221 messages=messages, 

222 project_name=project_name, 

223 ) 

224 model_output = predictions[ASSISTANT_COLUMN].iloc[-1] 

225 

226 return model_output