diff --git a/bot/bot.py b/bot/bot.py index f4510a6d..de97f841 100644 --- a/bot/bot.py +++ b/bot/bot.py @@ -42,6 +42,7 @@ HELP_MESSAGE = """Commands: ⚪ /retry – Regenerate last bot answer ⚪ /new – Start new dialog +⚪ /remove_all_history – Remove all history ⚪ /mode – Select chat mode ⚪ /settings – Show settings ⚪ /balance – Show balance @@ -178,7 +179,7 @@ async def retry_handle(update: Update, context: CallbackContext): await message_handle(update, context, message=last_dialog_message["user"], use_new_dialog_timeout=False) -async def message_handle(update: Update, context: CallbackContext, message=None, use_new_dialog_timeout=True): +async def message_handle(update: Update, context: CallbackContext, message=None, use_new_dialog_timeout=False): # check if bot was mentioned (for group chats) if not await is_bot_mentioned(update, context): return @@ -493,7 +494,7 @@ async def set_chat_mode_handle(update: Update, context: CallbackContext): chat_mode = query.data.split("|")[1] db.set_user_attribute(user_id, "current_chat_mode", chat_mode) - db.start_new_dialog(user_id) + db.try_resume_dialog(user_id, chat_mode) await context.bot.send_message( update.callback_query.message.chat.id, @@ -504,6 +505,8 @@ async def set_chat_mode_handle(update: Update, context: CallbackContext): def get_settings_menu(user_id: int): current_model = db.get_user_attribute(user_id, "current_model") + if current_model not in config.models['available_text_models']: + current_model = config.models['available_text_models'][0] text = config.models["info"][current_model]["description"] text += "\n\n" @@ -548,7 +551,6 @@ async def set_settings_handle(update: Update, context: CallbackContext): _, model_key = query.data.split("|") db.set_user_attribute(user_id, "current_model", model_key) - db.start_new_dialog(user_id) text, reply_markup = get_settings_menu(user_id) try: @@ -636,11 +638,21 @@ async def error_handle(update: Update, context: CallbackContext) -> None: except: await context.bot.send_message(update.effective_chat.id, "Some error in error handler") +async def remove_all_history_handle(update: Update, context: CallbackContext): + await register_user_if_not_exists(update, context, update.message.from_user) + user_id = update.message.from_user.id + db.remove_all_dialogs(user_id) + await update.message.reply_text("Cleared history ✅") + db.start_new_dialog(user_id) + await update.message.reply_text("Starting new dialog ✅") + + async def post_init(application: Application): await application.bot.set_my_commands([ BotCommand("/new", "Start new dialog"), BotCommand("/mode", "Select chat mode"), BotCommand("/retry", "Re-generate response for previous query"), + BotCommand("/remove_all_history", "Remove all history"), BotCommand("/balance", "Show balance"), BotCommand("/settings", "Show settings"), BotCommand("/help", "Show help message"), @@ -686,6 +698,7 @@ def run_bot() -> None: application.add_handler(CallbackQueryHandler(set_settings_handle, pattern="^set_settings")) application.add_handler(CommandHandler("balance", show_balance_handle, filters=user_filter)) + application.add_handler(CommandHandler("remove_all_history", remove_all_history_handle, filters=user_filter)) application.add_error_handler(error_handle) diff --git a/bot/database.py b/bot/database.py index b6bafe35..1c5e62eb 100644 --- a/bot/database.py +++ b/bot/database.py @@ -59,6 +59,8 @@ def add_new_user( def start_new_dialog(self, user_id: int): self.check_if_user_exists(user_id, raise_exception=True) + self.__remove_current_dialog(user_id) + dialog_id = str(uuid.uuid4()) dialog_dict = { "_id": dialog_id, @@ -126,3 +128,81 @@ def set_dialog_messages(self, user_id: int, dialog_messages: list, dialog_id: Op {"_id": dialog_id, "user_id": user_id}, {"$set": {"messages": dialog_messages}} ) + + def __remove_current_dialog(self, user_id: int, ) -> None: + """ + Removes the current dialog for the user + + Args: + user_id (int): user identifier + + Notes: + User's current_dialog_id is set to None, make sure you call `start_new_session` after this. + + If you want to completely remove all the dialogs, use `remove_all_dialogs` instead. + + Side Effects: + User's current_dialog_id is set to None. + """ + self.check_if_user_exists(user_id, raise_exception=True) + + dialog_id = self.get_user_attribute(user_id, "current_dialog_id") + self.set_user_attribute(user_id, "current_dialog_id", None) + + if dialog_id is None: + return + self.dialog_collection.delete_one({"_id": dialog_id}) + + def remove_all_dialogs(self, user_id: int): + """ + Removes all dialogs for the user + + Args: + user_id (int): user identifier + + Notes: + User's current_dialog_id is set to None, make sure you call `start_new_session` after this. + + """ + result = self.dialog_collection.delete_many({"user_id": user_id}) + print(f"Deleted: {result.deleted_count} dialogs") + self.set_user_attribute(user_id, "current_dialog_id", None) + + + def try_resume_dialog(self, user_id: int, current_chat_mode: str = None) -> str: + """ + Try to resume an existing dialog for the user + + Args: + user_id (int): user identifier + current_chat_mode (str, optional): the selected chat mode Defaults to None. + + Returns: + str: currrent_dialog_id + + Notes: + This function is called when the user switch between modes. If cannot find existing dialog, create a new one. + """ + self.check_if_user_exists(user_id, raise_exception=True) + if current_chat_mode is None: + current_chat_mode = self.get_user_attribute(user_id, "chat_mode") + # + """ + Get dialog_id of the most recent dialog + + Notes: + `sort` and `limit` are used to support old project with tons of unremoved dialogs in the DB. + """ + dialog_ids = self.dialog_collection.find( + filter={"user_id": user_id, "chat_mode": current_chat_mode}, + projection={"_id": 1}, + ).sort("start_time", -1).limit(1) + + if dialog_ids is not None: + return self.start_new_dialog(user_id) + latest_relevant_dialog_id = dialog_ids[0] + cid = latest_relevant_dialog_id['_id'] + self.set_user_attribute( + user_id, "current_dialog_id", cid + ) + return cid \ No newline at end of file