Provide user higher control over the chat history

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)
 

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