feat: optimize consumer thread pool

app.py

@@ -3,6 +3,7 @@
 import os
 import signal
 import sys
+import time
 
 from channel import channel_factory
 from common import const
@@ -24,6 +25,21 @@ def func(_signo, _stack_frame):
  signal.signal(_signo, func)
 
 
+def start_channel(channel_name: str):
+ channel = channel_factory.create_channel(channel_name)
+ if channel_name in ["wx", "wxy", "terminal", "wechatmp", "wechatmp_service", "wechatcom_app", "wework",
+ const.FEISHU, const.DINGTALK]:
+ PluginManager().load_plugins()
+
+ if conf().get("use_linkai"):
+ try:
+ from common import linkai_client
+ threading.Thread(target=linkai_client.start, args=(channel,)).start()
+ except Exception as e:
+ pass
+ channel.startup()
+
+
 def run():
  try:
  # load config
@@ -41,22 +57,11 @@ def run():
 
  if channel_name == "wxy":
  os.environ["WECHATY_LOG"] = "warn"
- # os.environ['WECHATY_PUPPET_SERVICE_ENDPOINT'] = '127.0.0.1:9001'
-
- channel = channel_factory.create_channel(channel_name)
- if channel_name in ["wx", "wxy", "terminal", "wechatmp", "wechatmp_service", "wechatcom_app", "wework", const.FEISHU,const.DINGTALK]:
- PluginManager().load_plugins()
-
- if conf().get("use_linkai"):
- try:
- from common import linkai_client
- threading.Thread(target=linkai_client.start, args=(channel, )).start()
- except Exception as e:
- pass
 
- # startup channel
- channel.startup()
+ start_channel(channel_name)
 
+ while True:
+ time.sleep(1)
  except Exception as e:
  logger.error("App startup failed!")
  logger.exception(e)

bot/linkai/link_ai_bot.py

@@ -400,7 +400,7 @@ def _send_image(self, channel, context, image_urls):
  i += 1
  if url.endswith(".mp4"):
  reply_type = ReplyType.VIDEO_URL
- elif url.endswith(".pdf") or url.endswith(".doc") or url.endswith(".docx"):
+ elif url.endswith(".pdf") or url.endswith(".doc") or url.endswith(".docx") or url.endswith(".csv"):
  reply_type = ReplyType.FILE
  url = _download_file(url)
  if not url:

channel/chat_channel.py

@@ -4,6 +4,7 @@
 import time
 from asyncio import CancelledError
 from concurrent.futures import Future, ThreadPoolExecutor
+from concurrent import futures
 
 from bridge.context import *
 from bridge.reply import *
@@ -17,6 +18,8 @@
 except Exception as e:
  pass
 
+handler_pool = ThreadPoolExecutor(max_workers=8) # 处理消息的线程池
+
 
 # 抽象类, 它包含了与消息通道无关的通用处理逻辑
 class ChatChannel(Channel):
@@ -25,7 +28,6 @@ class ChatChannel(Channel):
  futures = {} # 记录每个session_id提交到线程池的future对象, 用于重置会话时把没执行的future取消掉，正在执行的不会被取消
  sessions = {} # 用于控制并发，每个session_id同时只能有一个context在处理
  lock = threading.Lock() # 用于控制对sessions的访问
- handler_pool = ThreadPoolExecutor(max_workers=8) # 处理消息的线程池
 
  def __init__(self):
  _thread = threading.Thread(target=self.consume)
@@ -339,7 +341,7 @@ def consume(self):
  if not context_queue.empty():
  context = context_queue.get()
  logger.debug("[WX] consume context: {}".format(context))
- future: Future = self.handler_pool.submit(self._handle, context)
+ future: Future = handler_pool.submit(self._handle, context)
  future.add_done_callback(self._thread_pool_callback(session_id, context=context))
  if session_id not in self.futures:
  self.futures[session_id] = []

channel/wechat/wechat_channel.py

@@ -15,6 +15,7 @@
 from bridge.context import *
 from bridge.reply import *
 from channel.chat_channel import ChatChannel
+from channel import chat_channel
 from channel.wechat.wechat_message import *
 from common.expired_dict import ExpiredDict
 from common.log import logger
@@ -112,30 +113,39 @@ def __init__(self):
  self.auto_login_times = 0
 
  def startup(self):
- itchat.instance.receivingRetryCount = 600 # 修改断线超时时间
- # login by scan QRCode
- hotReload = conf().get("hot_reload", False)
- status_path = os.path.join(get_appdata_dir(), "itchat.pkl")
- itchat.auto_login(
- enableCmdQR=2,
- hotReload=hotReload,
- statusStorageDir=status_path,
- qrCallback=qrCallback,
- exitCallback=self.exitCallback,
- loginCallback=self.loginCallback
- )
- self.user_id = itchat.instance.storageClass.userName
- self.name = itchat.instance.storageClass.nickName
- logger.info("Wechat login success, user_id: {}, nickname: {}".format(self.user_id, self.name))
- # start message listener
- itchat.run()
+ try:
+ itchat.instance.receivingRetryCount = 600 # 修改断线超时时间
+ # login by scan QRCode
+ hotReload = conf().get("hot_reload", False)
+ status_path = os.path.join(get_appdata_dir(), "itchat.pkl")
+ itchat.auto_login(
+ enableCmdQR=2,
+ hotReload=hotReload,
+ statusStorageDir=status_path,
+ qrCallback=qrCallback,
+ exitCallback=self.exitCallback,
+ loginCallback=self.loginCallback
+ )
+ self.user_id = itchat.instance.storageClass.userName
+ self.name = itchat.instance.storageClass.nickName
+ logger.info("Wechat login success, user_id: {}, nickname: {}".format(self.user_id, self.name))
+ # start message listener
+ itchat.run()
+ except Exception as e:
+ logger.error(e)
 
  def exitCallback(self):
- _send_logout()
- time.sleep(3)
- self.auto_login_times += 1
- if self.auto_login_times < 100:
- self.startup()
+ try:
+ from common.linkai_client import chat_client
+ if chat_client.client_id and conf().get("use_linkai"):
+ _send_logout()
+ time.sleep(2)
+ self.auto_login_times += 1
+ if self.auto_login_times < 100:
+ chat_channel.handler_pool._shutdown = False
+ self.startup()
+ except Exception as e:
+ pass
 
  def loginCallback(self):
  logger.debug("Login success")
@@ -259,7 +269,6 @@ def _send_login_success():
 def _send_logout():
  try:
  from common.linkai_client import chat_client
- time.sleep(2)
  if chat_client.client_id:
  chat_client.send_logout()
  except Exception as e:
@@ -268,7 +277,6 @@ def _send_logout():
 def _send_qr_code(qrcode_list: list):
  try:
  from common.linkai_client import chat_client
- time.sleep(2)
  if chat_client.client_id:
  chat_client.send_qrcode(qrcode_list)
  except Exception as e:

common/linkai_client.py

@@ -2,7 +2,9 @@
 from bridge.reply import Reply, ReplyType
 from common.log import logger
 from linkai import LinkAIClient, PushMsg
-from config import conf
+from config import conf, pconf, plugin_config
+from plugins import PluginManager
+
 
 chat_client: LinkAIClient
 
@@ -22,6 +24,29 @@ def on_message(self, push_msg: PushMsg):
  context["isgroup"] = push_msg.is_group
  self.channel.send(Reply(ReplyType.TEXT, content=msg_content), context)
 
+ def on_config(self, config: dict):
+ if not self.client_id:
+ return
+ logger.info(f"从控制台加载配置: {config}")
+ local_config = conf()
+ for key in local_config.keys():
+ if config.get(key) is not None:
+ local_config[key] = config.get(key)
+ if config.get("reply_voice_mode"):
+ if config.get("reply_voice_mode") == "voice_reply_voice":
+ local_config["voice_reply_voice"] = True
+ elif config.get("reply_voice_mode") == "always_reply_voice":
+ local_config["always_reply_voice"] = True
+ # if config.get("admin_password") and plugin_config["Godcmd"]:
+ # plugin_config["Godcmd"]["password"] = config.get("admin_password")
+ # PluginManager().instances["Godcmd"].reload()
+ # if config.get("group_app_map") and pconf("linkai"):
+ # local_group_map = {}
+ # for mapping in config.get("group_app_map"):
+ # local_group_map[mapping.get("group_name")] = mapping.get("app_code")
+ # pconf("linkai")["group_app_map"] = local_group_map
+ # PluginManager().instances["linkai"].reload()
+
 
 def start(channel):
  global chat_client

plugins/godcmd/godcmd.py

@@ -475,3 +475,11 @@ def model_mapping(self, model) -> str:
  if model == "gpt-4-turbo":
  return const.GPT4_TURBO_PREVIEW
  return model
+
+ def reload(self):
+ gconf = plugin_config[self.name]
+ if gconf:
+ if gconf.get("password"):
+ self.password = gconf["password"]
+ if gconf.get("admin_users"):
+ self.admin_users = gconf["admin_users"]

plugins/plugin.py

@@ -46,3 +46,6 @@ def save_config(self, config: dict):
 
  def get_help_text(self, **kwargs):
  return "暂无帮助信息"
+
+ def reload(self):
+ pass