Merge pull request #17 from F33RNI/next

Next

.gitignore

@@ -2,3 +2,4 @@
 /*.zip
 /test.py
 /logs/
+/chats.json

AIHandler.py

@@ -18,20 +18,23 @@
 import queue
 import threading
 import time
+import uuid
 
 import openai
 
 import Authenticator
 import RequestResponseContainer
+from JSONReaderWriter import load_json, save_json
 
 EMPTY_RESPONSE_ERROR_MESSAGE = 'Empty response or unhandled error!'
 ERROR_CHATGPT_DISABLED = 'ChatGPT module is disabled in settings.json'
 ERROR_DALLE_DISABLED = 'DALL-E module is disabled in settings.json'
 
 
 class AIHandler:
-    def __init__(self, settings, authenticator):
+    def __init__(self, settings, chats_file, authenticator):
         self.settings = settings
+        self.chats_file = chats_file
         self.authenticator = authenticator
 
         # Loop running flag
@@ -46,10 +49,6 @@ def __init__(self, settings, authenticator):
         # Requests queue
         self.requests_queue = None
 
-        # Conversation id and parent id to continue dialog
-        self.conversation_id = None
-        self.parent_id = None
-
         # Check settings
         if self.settings is not None:
             # Initialize queue
@@ -68,6 +67,63 @@ def thread_start(self):
         thread.start()
         logging.info('AIHandler thread: ' + thread.name)
 
+    def get_chat(self, chat_id: int):
+        """
+        Retrieves conversation_id and parent_id for given chat_id or None if not exists
+        :param chat_id:
+        :return: (conversation_id, parent_id)
+        """
+        logging.info('Loading conversation_id for chat_id ' + str(chat_id))
+        chats = load_json(self.chats_file)
+        if chats is not None and str(chat_id) in chats:
+            chat = chats[str(chat_id)]
+            conversation_id = None
+            parent_id = None
+            if 'conversation_id' in chat:
+                conversation_id = chat['conversation_id']
+            if 'parent_id' in chat:
+                parent_id = chat['parent_id']
+
+            return conversation_id, parent_id
+        else:
+            return None, None
+
+    def set_chat(self, chat_id: int, conversation_id=None, parent_id=None):
+        """
+        Saves conversation ID and parent ID or Nones to remove it
+        :param chat_id:
+        :param conversation_id:
+        :param parent_id:
+        :return:
+        """
+        logging.info('Saving conversation_id ' + str(conversation_id) + ' and parent_id '
+                     + str(parent_id) + ' for chat_id ' + str(chat_id))
+        chats = load_json(self.chats_file)
+        if chats is not None:
+            if str(chat_id) in chats:
+                # Save or delete conversation_id
+                if conversation_id is not None and len(conversation_id) > 0:
+                    chats[str(chat_id)]['conversation_id'] = conversation_id
+                elif 'conversation_id' in chats[str(chat_id)]:
+                    del chats[str(chat_id)]['conversation_id']
+
+                # Save or delete parent_id
+                if parent_id is not None and len(parent_id) > 0:
+                    chats[str(chat_id)]['parent_id'] = parent_id
+                elif 'parent_id' in chats[str(chat_id)]:
+                    del chats[str(chat_id)]['parent_id']
+
+            # New chat
+            else:
+                chats[str(chat_id)] = {}
+                if conversation_id is not None and len(conversation_id) > 0:
+                    chats[str(chat_id)]['conversation_id'] = conversation_id
+                if parent_id is not None and len(parent_id) > 0:
+                    chats[str(chat_id)]['parent_id'] = parent_id
+        else:
+            chats = {}
+        save_json(self.chats_file, chats)
+
     def gpt_loop(self):
         """
         Background loop for handling requests
@@ -99,25 +155,33 @@ def gpt_loop(self):
 
                     # API type 0
                     if self.authenticator.api_type == 0:
-                        # Initialize conversation id
-                        if len(str(self.settings['chatgpt_api_0']['existing_conversation_id'])) > 0:
-                            self.conversation_id = str(self.settings['chatgpt_api_0']['existing_conversation_id'])
-                            logging.info('Conversation id: ' + str(self.conversation_id))
-                        else:
-                            self.conversation_id = None
+                        # Get conversation_id
+                        conversation_id, parent_id = self.get_chat(container.chat_id)
 
                         # Get chatbot from Authenticator class
                         chatbot = self.authenticator.chatbot
 
+                        # Reset chat
+                        chatbot.reset()
+
                         # Ask
-                        for data in chatbot.ask_stream(str(container.request), conversation_id=self.conversation_id):
+                        for data in chatbot.ask_stream(str(container.request), conversation_id=conversation_id):
                             # Initialize response
                             if api_response is None:
                                 api_response = ''
 
                             # Append response
                             api_response += str(data)
 
+                        # Generate and save conversation ID
+                        try:
+                            if conversation_id is None:
+                                conversation_id = str(uuid.uuid4())
+                            chatbot.save_conversation(conversation_id)
+                            self.set_chat(container.chat_id, conversation_id, parent_id)
+                        except Exception as e:
+                            logging.warning('Error saving conversation! ' + str(e))
+
                         # Remove tags
                         api_response = api_response.replace('<|im_end|>', '').replace('<|im_start|>', '')
 
@@ -140,34 +204,36 @@ def gpt_loop(self):
                         # Lock chatbot
                         self.authenticator.chatbot_locked = True
 
+                        # Get conversation_id and parent_id
+                        conversation_id, parent_id = self.get_chat(container.chat_id)
+
                         # Log request
-                        logging.info('Asking: ' + str(container.request))
-
-                        # Initialize conversation_id and parent_id
-                        if self.conversation_id is None:
-                            if len(str(self.settings['chatgpt_api_1']['chatgpt_dialog']['conversation_id'])) > 0:
-                                self.conversation_id \
-                                    = str(self.settings['chatgpt_api_1']['chatgpt_dialog']['conversation_id'])
-                                logging.info('Initial conversation id: ' + str(self.conversation_id))
-                        if self.parent_id is None:
-                            if len(str(self.settings['chatgpt_api_1']['chatgpt_dialog']['parent_id'])) > 0:
-                                self.parent_id = str(self.settings['chatgpt_api_1']['chatgpt_dialog']['parent_id'])
-                                logging.info('Initial parent id: ' + str(self.parent_id))
+                        logging.info('Asking: ' + str(container.request)
+                                     + ', conversation_id: ' + str(conversation_id) + ', parent_id: ' + str(parent_id))
+
+                        # Reset chat
+                        chatbot.reset_chat()
 
                         # Ask
                         for data in chatbot.ask(str(container.request),
-                                                conversation_id=self.conversation_id,
-                                                parent_id=self.parent_id):
+                                                conversation_id=conversation_id,
+                                                parent_id=parent_id):
                             # Get last response
                             api_response = data['message']
 
                             # Store conversation_id
-                            if data['conversation_id'] is not None:
-                                self.conversation_id = data['conversation_id']
+                            if 'conversation_id' in data and data['conversation_id'] is not None:
+                                conversation_id = data['conversation_id']
+
+                            # Store parent_id
+                            if 'parent_id' in data and data['parent_id'] is not None:
+                                parent_id = data['parent_id']
 
                         # Log conversation id and parent id
-                        logging.info('Current conversation id: ' + str(self.conversation_id)
-                                     + '\tParent id: ' + str(self.parent_id))
+                        logging.info('Current conversation_id: ' + conversation_id + ', parent_id: ' + parent_id)
+
+                        # Save conversation id
+                        self.set_chat(container.chat_id, conversation_id, parent_id)
 
                     # Wrong api type
                     else:

Authenticator.py

@@ -254,7 +254,8 @@ def proxy_checker_loop(self):
                         logging.warning(str(e))
 
                         # Wait before next try
-                        wait_seconds = int(self.settings['chatgpt_dialog']['too_many_requests_wait_time_seconds'])
+                        wait_seconds = \
+                            int(self.settings['chatgpt_api_1']['proxy']['too_many_requests_wait_time_seconds'])
                         logging.warning('Waiting ' + str(wait_seconds) + ' seconds...')
                         self.chatbot_too_many_requests = True
                         time.sleep(wait_seconds)

BotHandler.py

@@ -26,24 +26,25 @@
 from telegram.ext import ApplicationBuilder, ContextTypes, CommandHandler, MessageHandler, filters
 
 import RequestResponseContainer
-from AIHandler import AIHandler
 from main import TELEGRAMUS_VERSION
 
 BOT_COMMAND_START = 'start'
 BOT_COMMAND_HELP = 'help'
 BOT_COMMAND_QUEUE = 'queue'
 BOT_COMMAND_GPT = 'gpt'
 BOT_COMMAND_DRAW = 'draw'
+BOT_COMMAND_CLEAR = 'clear'
 BOT_COMMAND_RESTART = 'restart'
 
 # List of markdown chars to escape with \\
 MARKDOWN_ESCAPE = ['_', '*', '[', ']', '(', ')', '~', '>', '#', '+', '-', '=', '|', '{', '}', '.', '!']
 
 
 class BotHandler:
-    def __init__(self, settings, messages, ai_handler: AIHandler):
+    def __init__(self, settings, messages, chats_file, ai_handler):
         self.settings = settings
         self.messages = messages
+        self.chats_file = chats_file
         self.ai_handler = ai_handler
         self.application = None
         self.event_loop = None
@@ -83,6 +84,7 @@ def bot_start(self):
                 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_CLEAR, self.bot_command_clear))
                 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))
 
@@ -166,6 +168,34 @@ 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_clear(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
+        """
+        /clear command
+        :param update:
+        :param context:
+        :return:
+        """
+        user = update.message.from_user
+        chat_id = update.effective_chat.id
+        logging.info('/clear command from user ' + str(user.full_name) + ' request: ' + ' '.join(context.args))
+
+        # Delete conversation
+        if self.ai_handler.authenticator.api_type == 1:
+            conversation_id, _ = self.ai_handler.get_chat(chat_id)
+            if conversation_id is not None and len(conversation_id) > 0:
+                try:
+                    self.ai_handler.authenticator.chatbot.delete_conversation(conversation_id)
+                except Exception as e:
+                    logging.warning('Error deleting conversation ' + str(conversation_id) + ' ' + str(e))
+
+        # Clear conversation ID and parent ID
+        self.ai_handler.set_chat(chat_id, None, None)
+
+        try:
+            await context.bot.send_message(chat_id=chat_id, text=str(self.messages['chat_reset']).replace('\\n', '\n'))
+        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

JSONReaderWriter.py

@@ -0,0 +1,61 @@
+"""
+ Copyright (C) 2022 Fern Lane, GPT-telegramus
+ Licensed under the GNU Affero General Public License, Version 3.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+       https://www.gnu.org/licenses/agpl-3.0.en.html
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+ IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR
+ OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
+ ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
+ OTHER DEALINGS IN THE SOFTWARE.
+"""
+import json
+import logging
+import os.path
+
+
+def load_json(file_name: str):
+    """
+    Loads json from file_name
+    :return: json if loaded or empty json if not
+    """
+    try:
+        if os.path.exists(file_name):
+            logging.info('Loading ' + file_name + '...')
+            messages_file = open(file_name, encoding='utf-8')
+            json_content = json.load(messages_file)
+            messages_file.close()
+            if json_content is not None and len(str(json_content)) > 0:
+                logging.info('Loaded json: ' + str(json_content))
+            else:
+                json_content = None
+                logging.error('Error loading json data from file ' + file_name)
+        else:
+            logging.warning('No ' + file_name + ' file! Returning empty json')
+            return {}
+    except Exception as e:
+        json_content = None
+        logging.error(e, exc_info=True)
+
+    if json_content is None:
+        json_content = {}
+
+    return json_content
+
+
+def save_json(file_name: str, content):
+    """
+    Saves
+    :param file_name: filename to save
+    :param content: JSON dictionary
+    :return:
+    """
+    logging.info('Saving to ' + file_name + '...')
+    file = open(file_name, 'w')
+    json.dump(content, file, indent=4)
+    file.close()