1、优化多账号负载均衡时的选择 2、添加一系列功能

bot.py

@@ -28,6 +28,7 @@
 import chatbot
 from config import Config
 from utils.text_to_img import to_image
+import conversation_manager
 from manager.ratelimit import RateLimitManager
 import time
 from revChatGPT.V1 import Error as V1Error
@@ -71,11 +72,32 @@ async def create_timeout_task(target: Union[Friend, Group], source: Source):
 
 
 async def handle_message(target: Union[Friend, Group], session_id: str, message: str, source: Source) -> str:
+    number = session_id.split('-')[1]
     if not message.strip():
         return config.response.placeholder
 
     timeout_task = None
 
+    # 如果消息包含help命令（config.trigger.help_command所定义的内容），则回滚会话
+    if message.strip() in config.trigger.help_command:
+        return config.response.help_command.format(max_sessions=config.max_record.max_sessions)
+
+    # 如果消息包含 会话列表 命令（config.trigger.talk_list_command所定义的内容），则输出会话列表
+    if message.strip() in config.trigger.talk_list_command:
+        return "会话列表如下：\n" + '\n'.join(f"{i+1}. {x}" for i, x in enumerate(conversation_manager.get_user_sessions(number)))
+
+    # 如果消息包含会话上限x命令，则进入会话x
+    max_record_search = re.search(config.trigger.max_talk_sessions, message)
+    if max_record_search:
+        conversation_manager.update_max_sessions(number, int(max_record_search.group(1)))
+        if int(max_record_search.group(1)) > 10:
+                max_record = 10
+        elif int(max_record_search.group(1)) < 1:
+                max_record = 1
+        else:
+                max_record = int(max_record_search.group(1))
+        return f"会话上限已设置为{max_record}！"
+
     session, is_new_session = chatbot.get_chat_session(session_id)
 
     # 回滚
@@ -170,14 +192,19 @@ async def friend_message_listener(app: Ariadne, friend: Friend, source: Source,
     if rate_usage >= 1:
         response = config.ratelimit.exceed
     else:
-        response = await handle_message(friend, f"friend-{friend.id}", chain.display, source)
+        response = await handle_message(friend, f"fd-{friend.id}-", chain.display, source)
         if rate_usage >= config.ratelimit.warning_rate:
             limit = rateLimitManager.get_limit('好友', friend.id)
             usage = rateLimitManager.get_usage('好友', friend.id)
             current_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(time.time()))
             response = response + '\n' + config.ratelimit.warning_msg.format(usage=usage['count'], limit=limit['rate'],
                                                                              current_time=current_time)
-    await respond(friend, source, response)
+    if len(response) <= 3000:
+        await respond(friend, source, response)
+    else:
+        chunks = [response[i:i+3000] for i in range(0, len(response), 3000)]
+        for chunk in chunks:
+            await respond(friend, source, chunk)
 
 
 GroupTrigger = Annotated[MessageChain, MentionMe(config.trigger.require_mention != "at"), DetectPrefix(
@@ -193,7 +220,7 @@ async def group_message_listener(group: Group, source: Source, chain: GroupTrigg
     if rate_usage >= 1:
         return config.ratelimit.exceed
     else:
-        response = await handle_message(group, f"group-{group.id}", chain.display, source)
+        response = await handle_message(group, f"gp-{group.id}-", chain.display, source)
         if rate_usage >= config.ratelimit.warning_rate:
             limit = rateLimitManager.get_limit('群组', group.id)
             usage = rateLimitManager.get_usage('群组', group.id)

chatbot.py

@@ -4,6 +4,9 @@
 from manager.bot import BotManager, BotInfo
 import atexit
 from loguru import logger
+import conversation_manager
+import re
+from tinydb import TinyDB, Query
 import revChatGPT.V1 as V1
 
 config = Config.load_config()
@@ -24,10 +27,12 @@ def setup():
 class ChatSession:
     chatbot: BotInfo = None
     session_id: str
+    number: str
 
     def __init__(self, session_id):
         self.session_id = session_id
-        self.prev_conversation_id = []
+        self.number = session_id.split('-')[1]
+        self.prev_conversation_id = None
         self.prev_parent_id = []
         self.parent_id = None
         self.conversation_id = None
@@ -56,47 +61,188 @@ def reset_conversation(self):
             self.chatbot.bot.delete_conversation(self.conversation_id)
         self.conversation_id = None
         self.parent_id = None
-        self.prev_conversation_id = []
+        self.prev_conversation_id = None
         self.prev_parent_id = []
         self.chatbot = botManager.pick()
 
     def rollback_conversation(self) -> bool:
         if len(self.prev_parent_id) <= 0:
             return False
-        self.conversation_id = self.prev_conversation_id.pop()
+        # self.conversation_id = self.prev_conversation_id.pop()
         self.parent_id = self.prev_parent_id.pop()
+        # 回滚一次对话
+        conversation_manager.rollback_last_parent_id(self.number,
+                                                     self.conversation_id)
         return True
 
+    # 解析读取历史会话得到的字符串
+    def extract_conversations(self, result):
+
+        output = ""
+        message_mapping = result['mapping']
+        current_node_id = result['current_node']
+        current_node = message_mapping[current_node_id]
+        conversation = []
+
+        while current_node['parent'] is not None:
+            current_message = current_node['message']
+            if current_message is not None:
+                author_role = current_message['author']['role']
+                message_content = current_message['content']['parts'][0]
+                if author_role == 'user':
+                    conversation.append("you:" + message_content)
+                elif author_role == 'assistant':
+                    conversation.append("bot:" + message_content)
+
+            current_node_id = current_node['parent']
+            current_node = message_mapping[current_node_id]
+
+        for line in reversed(conversation):
+            output += line + "\n"
+
+        return output
+
     async def get_chat_response(self, message) -> str:
-        self.prev_conversation_id.append(self.conversation_id)
+
+        self.prev_conversation_id = self.conversation_id
         self.prev_parent_id.append(self.parent_id)
+        logger.info(f"当前id:{self.conversation_id}, 父节点id:{self.parent_id}")
+
+        # 如果消息包含进入会话x命令，则进入会话x
+        into_talk_search = re.search(config.trigger.goto_talk, message)
+        if into_talk_search:
+            conversation_parents = conversation_manager.get_last_parent_id(self.number, int(into_talk_search.group(1)))
+            if conversation_parents:
+                self.conversation_id = conversation_parents[0]
+                self.parent_id = conversation_parents[1][-1]
+                self.prev_conversation_id = conversation_parents[0]
+                self.prev_parent_id = conversation_parents[1][0:-1]
+                self.chatbot = botManager.pick_id(conversation_parents[2])
+                return f"进入会话{into_talk_search.group(1)}成功！"
+            else:
+                return f"进入会话{into_talk_search.group(1)}失败！停留在当前会话"
+
+        # 如果消息包含删除会话x命令，则删除对应会话
+        delete_talk_search = re.search(config.trigger.delete_talk, message)
+        if delete_talk_search:
+            result = conversation_manager.delete_session_record(self.number, int(delete_talk_search.group(1)))
+            if result == 2:  # 如果删除的恰巧是最后一个，也就是当前会话
+                self.reset_conversation()
+            if result:
+                return f"删除会话{delete_talk_search.group(1)}成功！"
+            else:
+                return f"删除会话{delete_talk_search.group(1)}失败！会话不存在"
+
+        # 如果消息包含读取会话x命令，则读取对应会话
+        read_talk_search = re.search(config.trigger.read_talk, message)
+        if read_talk_search:
+            conversation_parents = conversation_manager.get_last_parent_id(self.number, int(read_talk_search.group(1)))
+            if conversation_parents:
+                chatbot_save = self.chatbot
+                self.chatbot = botManager.pick_id(conversation_parents[2])
+                result = self.chatbot.bot.get_msg_history(conversation_parents[0], encoding='utf-8')
+                result = self.extract_conversations(result)
+                self.chatbot = chatbot_save
+                logger.info(
+                    f"会话{read_talk_search.group(1)}的聊天记录如下：\n{result}")
+                return f"读取会话成功！\n会话{read_talk_search.group(1)}的聊天记录如下：\n{result}"
+            else:
+                return f"读取会话{read_talk_search.group(1)}失败！停留在当前会话"
+
+        # 如果消息包含清空会话命令（config.trigger.rollback_command所定义的内容），则清空会话
+        if message.strip() in config.trigger.clear_talk_command:
+            self.reset_conversation()
+            if conversation_manager.clear_user_sessions(self.number):
+                return "清空会话成功！"
+            else:
+                return "清空会话失败！你是不是还没对话过。"
+
+        # 如果消息包含 会话名：*** 命令，则删除对应会话
+        rename_talk_search = re.search(config.trigger.rename_talk, message)
+        if rename_talk_search:
+            if self.conversation_id:
+                conversation_manager.update_session(self.number, self.conversation_id, rename_talk_search.group(1))
+                self.chatbot.bot.change_title(self.conversation_id, (
+                            self.session_id.encode('unicode-escape') + rename_talk_search.group(1).encode(
+                        'unicode-escape')).decode('utf-8'))
+                return f"会话名已改为{rename_talk_search.group(1)}"
+            else:
+                return "会话还没开始，不能设置会话名！"
 
         bot = self.chatbot.bot
+        botManager.update_bot_time(self.chatbot.id)
         bot.conversation_id = self.conversation_id
         bot.parent_id = self.parent_id
+        logger.info(
+            f"当前id:{self.conversation_id}, 父节点id:{self.parent_id}")
 
         loop = asyncio.get_event_loop()
         resp = await loop.run_in_executor(None, self.chatbot.ask, message, self.conversation_id, self.parent_id)
 
-        if self.conversation_id is None and self.chatbot.account.title_pattern:
-            self.chatbot.bot.change_title(resp["conversation_id"],
-                                          self.chatbot.account.title_pattern.format(session_id=self.session_id))
+        flag = False
+        if self.conversation_id is None:
+            flag = True
 
         self.conversation_id = resp["conversation_id"]
         self.parent_id = resp["parent_id"]
 
+        # 添加对话记录
+        conversation_manager.add_session_record(self.number, self.conversation_id, message[:20],
+                                                self.chatbot.account_id,
+                                                self.parent_id)
+
+        if flag:
+            self.chatbot.bot.change_title(self.conversation_id, f"{self.session_id}{message[:20].encode('utf-8')}")
+        # self.chatbot.bot.change_title(resp["conversation_id"],self.chatbot.account.title_pattern.format(session_id=f"{self.session_id}"+str(message[:20].encode("utf-8"))))##########################
+
+        logger.info(
+            f"当前id:{self.conversation_id}, 父节点id:{self.parent_id}")
+
         return resp["message"]
 
 
 __sessions = {}
 
 
 def get_chat_session(session_id: str) -> Tuple[ChatSession, bool]:
+    number = session_id.split('-')[1]
     new_session = False
-    if session_id not in __sessions:
-        __sessions[session_id] = ChatSession(session_id)
-        new_session = True
-    return __sessions[session_id], new_session
+
+    if number not in __sessions:  #有可能是重启容器了（读取旧聊天最后一个），也有可能是新的用户（创建新聊天）
+
+        # 创建一个新的聊天会话
+        __sessions[number] = ChatSession(session_id)
+
+        # 创建一个新的聊天会话
+        __sessions[number] = ChatSession(session_id)
+
+        # 打开数据库
+        db = TinyDB('data/session_records.json')
+
+        # 获取数据表
+        table = db.table('session_records')
+
+        # 读取数据
+        session_records = table.all()
+        if len(session_records) > 0:
+            session_records = session_records[0]
+        else:
+            session_records = {}
+
+        # 读取session_records字典
+        # with open('session_records.json', 'r') as f:
+        #     session_records = json.loads(f.read())
+        if number in session_records:  #重启容器了
+            if len(session_records[number]["sessions"]) > 0:
+                __sessions[number].conversation_id = session_records[number]["sessions"][-1]["conversation_id"]
+                __sessions[number].parent_id = session_records[number]["sessions"][-1]["parent_ids"][-1]
+                __sessions[number].prev_conversation_id = session_records[number]["sessions"][-1]["conversation_id"]
+                __sessions[number].prev_parent_id = session_records[number]["sessions"][-1]["parent_ids"][0:-1]
+                __sessions[number].chatbot = botManager.pick_id(session_records[number]["sessions"][-1]["account_id"])
+        else:  # 新用户开始聊天
+            new_session = True
+
+    return __sessions[number], new_session
 
 
 def conversation_remover():

config.py

@@ -47,13 +47,15 @@ class OpenAIAuthBase(BaseModel):
     auto_remove_old_conversations: bool = False
     """自动删除旧的对话"""
 
+
     class Config(BaseConfig):
         extra = Extra.allow
 
 
 class OpenAIEmailAuth(OpenAIAuthBase):
-    email: str
+    email: str = ''
     """OpenAI 注册邮箱"""
+
     password: str
     """OpenAI 密码"""
     isMicrosoftLogin: bool = False
@@ -90,6 +92,9 @@ class TextToImage(BaseModel):
     """纵坐标"""
     wkhtmltoimage: Union[str, None] = None
 
+class MaxRecord(BaseModel):
+    max_sessions: int = 5
+    """会话数量上限"""
 
 class Trigger(BaseModel):
     prefix: List[str] = [""]
@@ -100,6 +105,23 @@ class Trigger(BaseModel):
     """重置会话的命令"""
     rollback_command: List[str] = ["回滚会话"]
     """回滚会话的命令"""
+    help_command: List[str] = ["help"]
+    """请求帮助的命令"""
+    talk_list_command: List[str] = ["会话列表"]
+    """会话列表的命令"""
+    clear_talk_command: List[str] = ["清空会话"]
+    """清空会话的命令"""
+    max_talk_sessions: str = r"会话上限\s*(\d+)"
+
+    goto_talk: str = r"进入会话\s*(\d+)"
+
+    delete_talk: str = r"删除会话\s*(\d+)"
+
+    read_talk: str = r"读取会话\s*(\d+)"
+
+    rename_talk: str = r"会话名\s*[:：]\s*([^:：\s]+)"
+
+
 
 
 class Response(BaseModel):
@@ -153,6 +175,21 @@ class Response(BaseModel):
     queued_notice: str = "消息已收到！当前我还有{queue_size}条消息要回复，请您稍等。"
     """新消息进入队列时，发送的通知。 queue_size 是当前排队的消息数"""
 
+    help_command: str = (
+        "功能列表：\n"
+        "样例：指令 - 指令的功能\n"
+        "帮助 - 显示功能列表\n"
+        "重置会话 - 离开当前会话（保留），开启新的会话\n"
+        "回滚会话 - 回滚一次对话\n"
+        "会话名：*** - 设置当前会话名(开始会话之后才能设置)\n" 
+        "会话列表 - 显示自己现有的会话列表\n"
+        "进入会话x - 进入会话列表中第x个会话\n"
+        "读取会话x - 读取会话列表中第x的会话的所有会话记录\n"
+        "删除会话x - 删除会话列表中第x个会话\n"
+        "清空会话 - 清空自己所有的会话\n"
+        "会话上限x - 设置自己会话数量上限x，最大为{max_sessions}"
+    )
+
 
 class System(BaseModel):
     accept_group_invite: bool = False
@@ -186,6 +223,7 @@ class Config(BaseModel):
     response: Response = Response()
     system: System = System()
     presets: Preset = Preset()
+    max_record: MaxRecord = MaxRecord()
     ratelimit: Ratelimit = Ratelimit()
 
     def scan_presets(self):