Coverage for mindsdb / interfaces / chatbot / chatbot_controller.py: 66%
140 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-21 00:36 +0000
1from typing import Dict, List
3from mindsdb.api.executor.controllers.session_controller import SessionController
4from mindsdb.interfaces.agents.agents_controller import AgentsController
5from mindsdb.interfaces.chatbot.chatbot_task import ChatBotTask
6from mindsdb.interfaces.database.projects import ProjectController
7from mindsdb.interfaces.storage import db
8from mindsdb.interfaces.model.functions import get_project_records
9from mindsdb.utilities.exception import EntityNotExistsError
10from mindsdb.utilities.context import context as ctx
11from mindsdb.utilities.config import config
14default_project = config.get('default_project')
17class ChatBotController:
18 '''Handles CRUD operations at the database level for Chatbots'''
20 OBJECT_TYPE = 'chatbot'
22 def __init__(self, project_controller: ProjectController = None, agents_controller: AgentsController = None):
23 if project_controller is None: 23 ↛ 25line 23 didn't jump to line 25 because the condition on line 23 was always true
24 project_controller = ProjectController()
25 if agents_controller is None: 25 ↛ 27line 25 didn't jump to line 27 because the condition on line 25 was always true
26 agents_controller = AgentsController()
27 self.project_controller = project_controller
28 self.agents_controller = agents_controller
30 def get_chatbot(self, chatbot_name: str, project_name: str = default_project) -> dict:
31 '''
32 Gets a chatbot by name.
34 Parameters:
35 chatbot_name (str): The name of the chatbot
36 project_name (str): The name of the containing project
38 Returns:
39 bot (db.ChatBots): The database chatbot object
40 '''
42 project = self.project_controller.get(name=project_name)
44 query = db.session.query(
45 db.ChatBots, db.Tasks
46 ).join(
47 db.Tasks, db.ChatBots.id == db.Tasks.object_id
48 ).filter(
49 db.ChatBots.name == chatbot_name,
50 db.ChatBots.project_id == project.id,
51 db.Tasks.object_type == self.OBJECT_TYPE,
52 db.Tasks.company_id == ctx.company_id,
53 )
55 return self._get_chatbot(query, project)
57 def get_chatbot_by_id(self, chatbot_id: int) -> dict:
58 '''
59 Gets a chatbot by id.
61 Parameters:
62 chatbot_id (int): The id of the chatbot
64 Returns:
65 bot (db.ChatBots): The database chatbot object
66 '''
68 query = db.session.query(
69 db.ChatBots, db.Tasks
70 ).join(
71 db.Tasks, db.ChatBots.id == db.Tasks.object_id
72 ).filter(
73 db.ChatBots.id == chatbot_id,
74 db.Tasks.object_type == self.OBJECT_TYPE,
75 db.Tasks.company_id == ctx.company_id,
76 )
78 return self._get_chatbot(query)
80 def _get_chatbot(self, query, project: db.Project = None) -> dict:
81 '''
82 Gets a chatbot by query.
84 Parameters:
85 query: The query to get the chatbot
87 Returns:
88 bot (db.ChatBots): The database chatbot object
89 '''
91 query_result = query.first()
92 if query_result is None:
93 return None
94 bot, task = query_result
96 # Include DB, Agent, and Task information in response.
97 session = SessionController()
98 database_names = {
99 i['id']: i['name']
100 for i in session.database_controller.get_list()
101 }
103 agent = self.agents_controller.get_agent_by_id(bot.agent_id)
104 agent_obj = agent.as_dict() if agent is not None else None
106 bot_obj = {
107 'id': bot.id,
108 'name': bot.name,
109 'project': project.name if project else self.project_controller.get(bot.project_id).name,
110 'agent': agent_obj,
111 'database_id': bot.database_id, # TODO remove in future
112 'database': database_names.get(bot.database_id, '?'),
113 'model_name': bot.model_name,
114 'params': bot.params,
115 'created_at': bot.created_at,
116 'is_running': task.active,
117 'last_error': task.last_error,
118 'webhook_token': bot.webhook_token,
119 }
121 return bot_obj
123 def get_chatbots(self, project_name: str = default_project) -> List[dict]:
124 '''
125 Gets all chatbots in a project.
127 Parameters:
128 project_name (str): The name of the containing project. If None, then return from all projects
130 Returns:
131 all_bots (List[db.ChatBots]): List of database chatbot object
132 '''
134 project_names = {}
135 for project in get_project_records():
136 if project_name is not None and project.name != project_name:
137 continue
138 project_names[project.id] = project.name
140 if project_name is not None and project_name not in project_names.values():
141 raise EntityNotExistsError(f'Project {project_name} not found')
143 query = db.session.query(
144 db.ChatBots, db.Tasks
145 ).join(
146 db.Tasks, db.ChatBots.id == db.Tasks.object_id
147 ).filter(
148 db.ChatBots.project_id.in_(list(project_names.keys())),
149 db.Tasks.object_type == self.OBJECT_TYPE,
150 db.Tasks.company_id == ctx.company_id,
151 )
153 session = SessionController()
154 database_names = {
155 i['id']: i['name']
156 for i in session.database_controller.get_list()
157 }
159 bots = []
160 for bot, task in query.all():
161 agent = self.agents_controller.get_agent_by_id(bot.agent_id)
162 agent_obj = agent.as_dict() if agent is not None else None
163 bots.append(
164 {
165 'id': bot.id,
166 'name': bot.name,
167 'project': project_names[bot.project_id],
168 'agent': agent_obj,
169 'database_id': bot.database_id, # TODO remove in future
170 'database': database_names.get(bot.database_id, '?'),
171 'model_name': bot.model_name,
172 'params': bot.params,
173 'created_at': bot.created_at,
174 'is_running': task.active,
175 'last_error': task.last_error,
176 'webhook_token': bot.webhook_token,
177 }
178 )
180 return bots
182 def add_chatbot(
183 self,
184 name: str,
185 project_name: str,
186 model_name: str = None,
187 agent_name: str = None,
188 database_id: int = None,
189 is_running: bool = True,
190 params: Dict[str, str] = {}) -> db.ChatBots:
191 '''
192 Adds a chatbot to the database.
194 Parameters:
195 name (str): The name of the new chatbot
196 project_name (str): The containing project
197 model_name (str): The name of the existing ML model the chatbot will use
198 agent_name (str): The name of the existing agent the chatbot will use
199 database_id (int): The ID of the existing database the chatbot will use
200 is_running (bool): Whether or not to start the chatbot right after creation
201 params: (Dict[str, str]): Parameters to use when running the chatbot
203 Returns:
204 bot (db.ChatBots): The created chatbot
205 '''
207 is_cloud = config.get('cloud', False)
208 if is_cloud and ctx.user_class == 0: 208 ↛ 209line 208 didn't jump to line 209 because the condition on line 208 was never true
209 raise Exception("You can't create chatbot")
211 if project_name is None: 211 ↛ 212line 211 didn't jump to line 212 because the condition on line 211 was never true
212 project_name = default_project
213 project = self.project_controller.get(name=project_name)
215 bot = self.get_chatbot(name, project_name)
217 if bot is not None: 217 ↛ 218line 217 didn't jump to line 218 because the condition on line 217 was never true
218 raise Exception(f'Chat bot already exists: {name}')
220 # check database
221 session_controller = SessionController()
222 db_record = session_controller.integration_controller.get_by_id(database_id)
223 if db_record is None: 223 ↛ 224line 223 didn't jump to line 224 because the condition on line 223 was never true
224 raise Exception(f"Database doesn't exist: {database_id}")
226 if model_name is None and agent_name is None: 226 ↛ 227line 226 didn't jump to line 227 because the condition on line 226 was never true
227 raise ValueError('Need to provide either "model_name" or "agent_name" when creating a chatbot')
228 if agent_name is not None: 228 ↛ 229line 228 didn't jump to line 229 because the condition on line 228 was never true
229 agent = self.agents_controller.get_agent(agent_name, project_name)
230 if agent is None:
231 raise ValueError(f"Agent with name doesn't exist: {agent_name}")
232 model_name = agent.model_name
233 agent_id = agent.id
234 else:
235 # Create a new agent with the given model name.
236 agent_id = None
238 bot = db.ChatBots(
239 name=name,
240 project_id=project.id,
241 agent_id=agent_id,
242 model_name=model_name,
243 database_id=database_id,
244 params=params,
245 )
246 db.session.add(bot)
247 db.session.flush()
249 task_record = db.Tasks(
250 company_id=ctx.company_id,
251 user_class=ctx.user_class,
253 object_type=self.OBJECT_TYPE,
254 object_id=bot.id,
255 active=is_running
256 )
257 db.session.add(task_record)
259 db.session.commit()
261 return bot
263 def update_chatbot(
264 self,
265 chatbot_name: str,
266 project_name: str = default_project,
267 name: str = None,
268 model_name: str = None,
269 agent_name: str = None,
270 database_id: int = None,
271 is_running: bool = None,
272 params: Dict[str, str] = None,
273 webhook_token: str = None) -> db.ChatBots:
274 '''
275 Updates a chatbot in the database, creating it if it doesn't already exist.
277 Parameters:
278 chatbot_name (str): The name of the new chatbot, or existing chatbot to update
279 project_name (str): The containing project
280 name (str): The updated name of the chatbot
281 model_name (str): The name of the existing ML model the chatbot will use
282 agent_name (str): The name of the existing agent the chatbot will use
283 database_id (int): The ID of the existing database the chatbot will use
284 is_running (bool): Whether or not the chatbot will run after update/creation
285 params: (Dict[str, str]): Parameters to use when running the chatbot
287 Returns:
288 bot (db.ChatBots): The created or updated chatbot
289 '''
291 existing_chatbot = self.get_chatbot(chatbot_name, project_name=project_name)
292 if existing_chatbot is None: 292 ↛ 293line 292 didn't jump to line 293 because the condition on line 292 was never true
293 raise Exception(f'Chat bot not found: {chatbot_name}')
295 existing_chatbot_rec = db.ChatBots.query.get(existing_chatbot['id'])
297 if name is not None and name != chatbot_name: 297 ↛ 299line 297 didn't jump to line 299 because the condition on line 297 was never true
298 # check new name
299 bot2 = self.get_chatbot(name, project_name=project_name)
300 if bot2 is not None:
301 raise Exception(f'Chat already exists: {name}')
303 existing_chatbot_rec.name = name
305 if agent_name is not None: 305 ↛ 306line 305 didn't jump to line 306 because the condition on line 305 was never true
306 agent = self.agents_controller.get_agent(agent_name, project_name)
307 if agent is None:
308 raise ValueError(f"Agent with name doesn't exist: {agent_name}")
309 existing_chatbot_rec.agent_id = agent.id
311 if model_name is not None: 311 ↛ 313line 311 didn't jump to line 313 because the condition on line 311 was never true
312 # TODO check model_name
313 existing_chatbot_rec.model_name = model_name
314 if database_id is not None: 314 ↛ 316line 314 didn't jump to line 316 because the condition on line 314 was never true
315 # TODO check database_id
316 existing_chatbot_rec.database_id = database_id
318 task = db.Tasks.query.filter(
319 db.Tasks.object_type == self.OBJECT_TYPE,
320 db.Tasks.object_id == existing_chatbot_rec.id,
321 db.Tasks.company_id == ctx.company_id,
322 ).first()
324 if task is not None: 324 ↛ 331line 324 didn't jump to line 331 because the condition on line 324 was always true
325 if is_running is not None: 325 ↛ 326line 325 didn't jump to line 326 because the condition on line 325 was never true
326 task.active = is_running
328 # reload task
329 task.reload = True
331 if params is not None: 331 ↛ 337line 331 didn't jump to line 337 because the condition on line 331 was always true
332 # Merge params on update
333 existing_params = existing_chatbot_rec.params or {}
334 params.update(existing_params)
335 existing_chatbot_rec.params = params
337 if webhook_token is not None: 337 ↛ 338line 337 didn't jump to line 338 because the condition on line 337 was never true
338 existing_chatbot_rec.webhook_token = webhook_token
340 db.session.commit()
342 return existing_chatbot_rec
344 def delete_chatbot(self, chatbot_name: str, project_name: str = default_project):
345 '''
346 Deletes a chatbot by name.
348 Parameters:
349 chatbot_name (str): The name of the chatbot to delete
350 project_name (str): The name of the containing project
351 '''
353 bot = self.get_chatbot(chatbot_name, project_name)
354 if bot is None: 354 ↛ 355line 354 didn't jump to line 355 because the condition on line 354 was never true
355 raise Exception(f"Chat bot doesn't exist: {chatbot_name}")
357 bot_rec = db.ChatBots.query.get(bot['id'])
359 task = db.Tasks.query.filter(
360 db.Tasks.object_type == self.OBJECT_TYPE,
361 db.Tasks.object_id == bot_rec.id,
362 db.Tasks.company_id == ctx.company_id,
363 ).first()
365 if task is not None: 365 ↛ 368line 365 didn't jump to line 368 because the condition on line 365 was always true
366 db.session.delete(task)
368 db.session.delete(bot_rec)
370 db.session.commit()
372 def on_webhook(self, webhook_token: str, request: dict, chat_bot_memory: dict):
373 """
374 Handles incoming webhook requests.
375 Finds the chat bot associated with the webhook token and passes the request to the chat bot task.
377 Args:
378 webhook_token (str): The token to uniquely identify the webhook.
379 request (dict): The incoming webhook request.
380 chat_bot_memory (dict): The memory of the various chat-bots mapped by their webhook tokens.
381 """
382 query = db.session.query(
383 db.ChatBots, db.Tasks
384 ).join(
385 db.Tasks, db.ChatBots.id == db.Tasks.object_id
386 ).filter(
387 db.ChatBots.webhook_token == webhook_token,
388 db.Tasks.object_type == self.OBJECT_TYPE,
389 db.Tasks.company_id == ctx.company_id,
390 )
391 result = query.first()
393 chat_bot, task = result if result is not None else (None, None)
395 if chat_bot is None:
396 raise Exception(f"No chat bot exists for webhook token: {webhook_token}")
398 if not task.active:
399 raise Exception(f"Chat bot is not running: {chat_bot.name}")
401 chat_bot_task = ChatBotTask(task_id=task.id, object_id=chat_bot.id)
403 if webhook_token in chat_bot_memory:
404 chat_bot_task.set_memory(chat_bot_memory[webhook_token])
405 else:
406 chat_bot_memory[webhook_token] = chat_bot_task.get_memory()
408 chat_bot_task.on_webhook(request)