Add restart command and fix queue

Authenticator.py

@@ -85,9 +85,9 @@ def __init__(self, settings):
         self.proxy_list = []
         self.check_loop_running = False
 
-    def start_check_loop(self):
+    def start_chatbot(self):
         """
-        Starts background thread
+        Initializes chatbot and starts background proxy checker thread if needed
         :return:
         """
         # Official API
@@ -145,6 +145,27 @@ def start_check_loop(self):
             logging.error('Wrong chatgpt_api_type!')
             raise Exception('Wrong chatgpt_api_type')
 
+    def stop_chatbot(self):
+        """
+        Stops background handler and removes chatbot
+        :return:
+        """
+        logging.info('Stopping chatbot...')
+        # Clear loop flag
+        self.check_loop_running = False
+        self.chatbot_working = False
+
+        # Sleep some time
+        time.sleep(10)
+
+        # Remove old chatbot
+        try:
+            if self.chatbot is not None:
+                del self.chatbot
+                self.chatbot = None
+        except Exception as e:
+            logging.warning('Error clearing chatbot! ' + str(e))
+
     def proxy_get(self):
         """
         Retrieves proxy from auto_proxy_from url into self.proxy_list
@@ -260,6 +281,10 @@ def proxy_checker_loop(self):
                         break
                     time.sleep(1)
 
+                # Exit if thread stopped
+                if not self.check_loop_running:
+                    break
+
             # Check is not successful
             else:
                 # Get proxy
@@ -288,6 +313,11 @@ def proxy_checker_loop(self):
                 default_config = self.get_chatbot_config()
 
                 while True:
+                    # Exit if thread stopped
+                    if not self.check_loop_running:
+                        kill_all_processes(processes_and_times)
+                        break
+
                     # Create and start processes
                     while len(self.proxy_list) > 0 \
                             and len(processes_and_times) \

BotHandler.py

@@ -34,6 +34,7 @@
 BOT_COMMAND_QUEUE = 'queue'
 BOT_COMMAND_GPT = 'gpt'
 BOT_COMMAND_DRAW = 'draw'
+BOT_COMMAND_RESTART = 'restart'
 
 # List of markdown chars to escape with \\
 MARKDOWN_ESCAPE = ['_', '*', '[', ']', '(', ')', '~', '>', '#', '+', '-', '=', '|', '{', '}', '.', '!']
@@ -44,6 +45,8 @@ def __init__(self, settings, messages, ai_handler: AIHandler):
         self.settings = settings
         self.messages = messages
         self.ai_handler = ai_handler
+        self.application = None
+        self.event_loop = None
 
         # Response loop running flag
         self.response_loop_running = False
@@ -71,18 +74,22 @@ def bot_start(self):
         while True:
             try:
                 # Build bot
-                application = ApplicationBuilder().token(self.settings['telegram']['api_key']) \
-                    .write_timeout(30).read_timeout(30).build()
-                application.add_handler(CommandHandler(BOT_COMMAND_START, self.bot_command_start))
-                application.add_handler(CommandHandler(BOT_COMMAND_HELP, self.bot_command_help))
-                application.add_handler(CommandHandler(BOT_COMMAND_QUEUE, self.bot_command_queue))
-                application.add_handler(CommandHandler(BOT_COMMAND_GPT, self.bot_command_gpt))
-                application.add_handler(CommandHandler(BOT_COMMAND_DRAW, self.bot_command_draw))
-                application.add_handler(MessageHandler(filters.TEXT & (~filters.COMMAND), self.bot_read_message))
+                builder = ApplicationBuilder().token(self.settings['telegram']['api_key'])
+                builder.write_timeout(30)
+                builder.read_timeout(30)
+                self.application = builder.build()
+                self.application.add_handler(CommandHandler(BOT_COMMAND_START, self.bot_command_start))
+                self.application.add_handler(CommandHandler(BOT_COMMAND_HELP, self.bot_command_help))
+                self.application.add_handler(CommandHandler(BOT_COMMAND_QUEUE, self.bot_command_queue))
+                self.application.add_handler(CommandHandler(BOT_COMMAND_GPT, self.bot_command_gpt))
+                self.application.add_handler(CommandHandler(BOT_COMMAND_DRAW, self.bot_command_draw))
+                self.application.add_handler(CommandHandler(BOT_COMMAND_RESTART, self.bot_command_restart))
+                self.application.add_handler(MessageHandler(filters.TEXT & (~filters.COMMAND), self.bot_read_message))
 
                 # Start bot
-                asyncio.set_event_loop(asyncio.new_event_loop())
-                asyncio.run(application.run_polling())
+                self.event_loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(self.event_loop)
+                self.event_loop.run_until_complete(self.application.run_polling())
             except Exception as e:
                 logging.error('Telegram bot error! ' + str(e))
             logging.info('Restarting bot polling after 5 seconds...')
@@ -159,6 +166,47 @@ async def bot_read_message(self, update: Update, context: ContextTypes.DEFAULT_T
             except Exception as e:
                 logging.error('Error sending message! ' + str(e))
 
+    async def bot_command_restart(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
+        """
+        /restart command
+        :param update:
+        :param context:
+        :return:
+        """
+        user = update.message.from_user
+        chat_id = update.effective_chat.id
+        logging.info('/restart command from user ' + str(user.full_name) + ' request: ' + ' '.join(context.args))
+
+        logging.info('Restarting...')
+        try:
+            await context.bot.send_message(chat_id=chat_id, text=str(self.messages['restarting']).replace('\\n', '\n'))
+        except Exception as e:
+            logging.error('Error sending message! ' + str(e))
+
+        # Restart chatbot
+        self.ai_handler.authenticator.stop_chatbot()
+        self.ai_handler.authenticator.start_chatbot()
+        time.sleep(1)
+
+        # Restart telegram bot
+        if self.event_loop is not None:
+            self.event_loop.stop()
+            try:
+                self.event_loop.close()
+            except:
+                pass
+
+        # Sleep some time again
+        time.sleep(10)
+
+        # Done
+        logging.info('Restarting done')
+        try:
+            await context.bot.send_message(chat_id=chat_id,
+                                           text=str(self.messages['restarting_done']).replace('\\n', '\n'))
+        except Exception as e:
+            logging.error('Error sending message! ' + str(e))
+
     async def bot_command_draw(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
         """
         /draw command
@@ -235,27 +283,25 @@ async def bot_command_queue(self, update: Update, context: ContextTypes.DEFAULT_
             except Exception as e:
                 logging.error('Error sending message! ' + str(e))
         else:
-            i = 0
             message = ''
 
+            # Current request
+            if processing_container is not None:
+                message += '1. ' + processing_container.user_name + ', '
+                message += RequestResponseContainer.REQUEST_NAMES[processing_container.request_type]
+                message += ': ' + processing_container.request + '\n\n'
+
             # From queue
             if not self.requests_queue.empty():
                 for i in range(self.requests_queue.qsize()):
                     container = self.requests_queue.queue[i]
                     text_request = container.request
                     text_from = container.user_name
-                    message += str(i + 1) + '. ' + text_from + ', '
+                    queue_index = (i + 1) if processing_container is not None else i
+                    message += str(queue_index + 1) + '. ' + text_from + ', '
                     message += RequestResponseContainer.REQUEST_NAMES[container.request_type]
                     message += ': ' + text_request + '\n\n'
 
-            # Current request
-            if processing_container is not None:
-                if len(message) > 0:
-                    i += 1
-                message += str(i + 1) + '. ' + processing_container.user_name + ', '
-                message += RequestResponseContainer.REQUEST_NAMES[processing_container.request_type]
-                message += ': ' + processing_container.request + '\n\n'
-
             # Send queue stats
             try:
                 await context.bot.send_message(chat_id=chat_id, text=str(self.messages['queue_stats'])
@@ -320,7 +366,7 @@ async def send_reply(self, chat_id: int, message: str, reply_to_message_id: int,
                     message = message.replace(escape_char, '\\' + escape_char)
 
                 try:
-                    await telegram.Bot(self.settings['telegram']['api_key'])\
+                    await telegram.Bot(self.settings['telegram']['api_key']) \
                         .sendMessage(chat_id=chat_id,
                                      text=message,
                                      reply_to_message_id=reply_to_message_id,
@@ -332,15 +378,15 @@ async def send_reply(self, chat_id: int, message: str, reply_to_message_id: int,
             except Exception as e:
                 logging.info(e)
                 try:
-                    await telegram.Bot(self.settings['telegram']['api_key'])\
+                    await telegram.Bot(self.settings['telegram']['api_key']) \
                         .sendMessage(chat_id=chat_id,
                                      text=message.replace('\\n', '\n'),
                                      reply_to_message_id=reply_to_message_id)
                 except Exception as e:
                     logging.error('Error sending message! ' + str(e))
         else:
             try:
-                await telegram.Bot(self.settings['telegram']['api_key'])\
+                await telegram.Bot(self.settings['telegram']['api_key']) \
                     .sendMessage(chat_id=chat_id,
                                  text=message.replace('\\n', '\n'),
                                  reply_to_message_id=reply_to_message_id)

main.py

@@ -18,7 +18,6 @@
 import argparse
 import datetime
 import json
-import locale
 import logging
 import os
 import signal
@@ -131,7 +130,7 @@ def main():
     logging_setup()
 
     # Connect interrupt signal
-    signal.signal(signal.SIGINT, exit_)
+    #signal.signal(signal.SIGINT, exit_)
 
     # Parse arguments and load settings and messages
     args = parse_args()
@@ -146,8 +145,8 @@ def main():
     # Set requests_queue to ai_handler
     ai_handler.requests_queue = bot_handler.requests_queue
 
-    # Start checker loop
-    authenticator.start_check_loop()
+    # Initialize chatbot and start checker loop
+    authenticator.start_chatbot()
 
     # Start AIHandler
     ai_handler.thread_start()

messages.json

@@ -9,5 +9,7 @@
   "queue_stats": "Queue:\\n\\n{0}",
   "queue_accepted": "Added to the queue. Position: {0}/{1}",
   "gpt_error": "Error: {0}\\n\\nMake another request or try again later =(",
-  "reset": "Reset command sent"
+  "reset": "Reset command sent",
+  "restarting": "Restarting in progress...\\nPlease wait",
+  "restarting_done": "The restart is completed"
 }

requirements.txt

@@ -1,6 +1,6 @@
 psutil>=5.9.1
 telegram~=0.0.1
-python-telegram-bot~=20.0
+python-telegram-bot~=20.1
 revChatGPT>=2.3.14
 openai>=0.26.4
 tiktoken>=0.2.0