|
- import os
- import re
- import threading
- import time
- from asyncio import CancelledError
- from concurrent.futures import Future, ThreadPoolExecutor
-
- from bridge.context import *
- from bridge.reply import *
- from channel.channel import Channel
- from common.dequeue import Dequeue
- from common import memory
- from plugins import *
-
- try:
- from voice.audio_convert import any_to_wav
- except Exception as e:
- pass
-
-
-
- class ChatChannel(Channel):
- name = None
- user_id = None
- futures = {}
- sessions = {}
- lock = threading.Lock()
- handler_pool = ThreadPoolExecutor(max_workers=8)
-
- def __init__(self):
- _thread = threading.Thread(target=self.consume)
- _thread.setDaemon(True)
- _thread.start()
-
-
- def _compose_context(self, ctype: ContextType, content, **kwargs):
- context = Context(ctype, content)
- context.kwargs = kwargs
-
-
-
- if "origin_ctype" not in context:
- context["origin_ctype"] = ctype
-
- first_in = "receiver" not in context
-
- if first_in:
- config = conf()
- cmsg = context["msg"]
- user_data = conf().get_user_data(cmsg.from_user_id)
- context["openai_api_key"] = user_data.get("openai_api_key")
- context["gpt_model"] = user_data.get("gpt_model")
- if context.get("isgroup", False):
- group_name = cmsg.other_user_nickname
- group_id = cmsg.other_user_id
-
- group_name_white_list = config.get("group_name_white_list", [])
- group_name_keyword_white_list = config.get("group_name_keyword_white_list", [])
- if any(
- [
- group_name in group_name_white_list,
- "ALL_GROUP" in group_name_white_list,
- check_contain(group_name, group_name_keyword_white_list),
- ]
- ):
- group_chat_in_one_session = conf().get("group_chat_in_one_session", [])
- session_id = cmsg.actual_user_id
- if any(
- [
- group_name in group_chat_in_one_session,
- "ALL_GROUP" in group_chat_in_one_session,
- ]
- ):
- session_id = group_id
- else:
- return None
- context["session_id"] = session_id
- context["receiver"] = group_id
- else:
- context["session_id"] = cmsg.other_user_id
- context["receiver"] = cmsg.other_user_id
- e_context = PluginManager().emit_event(EventContext(Event.ON_RECEIVE_MESSAGE, {"channel": self, "context": context}))
- context = e_context["context"]
- if e_context.is_pass() or context is None:
- return context
- if cmsg.from_user_id == self.user_id and not config.get("trigger_by_self", True):
- logger.debug("[WX]self message skipped")
- return None
-
-
- if ctype == ContextType.TEXT:
- if first_in and "」\n- - - - - - -" in content:
- logger.debug(content)
- logger.debug("[WX]reference query skipped")
- return None
-
- nick_name_black_list = conf().get("nick_name_black_list", [])
- if context.get("isgroup", False):
-
- match_prefix = check_prefix(content, conf().get("group_chat_prefix"))
- match_contain = check_contain(content, conf().get("group_chat_keyword"))
- flag = False
- if context["msg"].to_user_id != context["msg"].actual_user_id:
- if match_prefix is not None or match_contain is not None:
- flag = True
- if match_prefix:
- content = content.replace(match_prefix, "", 1).strip()
- if context["msg"].is_at:
- nick_name = context["msg"].actual_user_nickname
- if nick_name and nick_name in nick_name_black_list:
-
- logger.warning(f"[WX] Nickname {nick_name} in In BlackList, ignore")
- return None
-
- logger.info("[WX]receive group at")
- if not conf().get("group_at_off", False):
- flag = True
- pattern = f"@{re.escape(self.name)}(\u2005|\u0020)"
- subtract_res = re.sub(pattern, r"", content)
- if isinstance(context["msg"].at_list, list):
- for at in context["msg"].at_list:
- pattern = f"@{re.escape(at)}(\u2005|\u0020)"
- subtract_res = re.sub(pattern, r"", subtract_res)
- if subtract_res == content and context["msg"].self_display_name:
-
- pattern = f"@{re.escape(context['msg'].self_display_name)}(\u2005|\u0020)"
- subtract_res = re.sub(pattern, r"", content)
- content = subtract_res
- if not flag:
- if context["origin_ctype"] == ContextType.VOICE:
- logger.info("[WX]receive group voice, but checkprefix didn't match")
- return None
- else:
- nick_name = context["msg"].from_user_nickname
- if nick_name and nick_name in nick_name_black_list:
-
- logger.warning(f"[WX] Nickname '{nick_name}' in In BlackList, ignore")
- return None
-
- match_prefix = check_prefix(content, conf().get("single_chat_prefix", [""]))
- if match_prefix is not None:
- content = content.replace(match_prefix, "", 1).strip()
- elif context["origin_ctype"] == ContextType.VOICE:
- pass
- else:
- return None
- content = content.strip()
- img_match_prefix = check_prefix(content, conf().get("image_create_prefix"))
- if img_match_prefix:
- content = content.replace(img_match_prefix, "", 1)
- context.type = ContextType.IMAGE_CREATE
- else:
- context.type = ContextType.TEXT
- context.content = content.strip()
- if "desire_rtype" not in context and conf().get("always_reply_voice") and ReplyType.VOICE not in self.NOT_SUPPORT_REPLYTYPE:
- context["desire_rtype"] = ReplyType.VOICE
- elif context.type == ContextType.VOICE:
- if "desire_rtype" not in context and conf().get("voice_reply_voice") and ReplyType.VOICE not in self.NOT_SUPPORT_REPLYTYPE:
- context["desire_rtype"] = ReplyType.VOICE
-
- return context
-
- def _handle(self, context: Context):
- if context is None or not context.content:
- return
- logger.debug("[WX] ready to handle context: {}".format(context))
-
- reply = self._generate_reply(context)
-
- logger.debug("[WX] ready to decorate reply: {}".format(reply))
-
- reply = self._decorate_reply(context, reply)
-
-
- self._send_reply(context, reply)
-
- def _generate_reply(self, context: Context, reply: Reply = Reply()) -> Reply:
- e_context = PluginManager().emit_event(
- EventContext(
- Event.ON_HANDLE_CONTEXT,
- {"channel": self, "context": context, "reply": reply},
- )
- )
- reply = e_context["reply"]
- if not e_context.is_pass():
- logger.debug("[WX] ready to handle context: type={}, content={}".format(context.type, context.content))
- if context.type == ContextType.TEXT or context.type == ContextType.IMAGE_CREATE:
- context["channel"] = e_context["channel"]
- reply = super().build_reply_content(context.content, context)
- elif context.type == ContextType.VOICE:
- cmsg = context["msg"]
- cmsg.prepare()
- file_path = context.content
- wav_path = os.path.splitext(file_path)[0] + ".wav"
- try:
- any_to_wav(file_path, wav_path)
- except Exception as e:
- logger.warning("[WX]any to wav error, use raw path. " + str(e))
- wav_path = file_path
-
- reply = super().build_voice_to_text(wav_path)
-
- try:
- os.remove(file_path)
- if wav_path != file_path:
- os.remove(wav_path)
- except Exception as e:
- pass
-
-
- if reply.type == ReplyType.TEXT:
- new_context = self._compose_context(ContextType.TEXT, reply.content, **context.kwargs)
- if new_context:
- reply = self._generate_reply(new_context)
- else:
- return
- elif context.type == ContextType.IMAGE:
- memory.USER_IMAGE_CACHE[context["session_id"]] = {
- "path": context.content,
- "msg": context.get("msg")
- }
- elif context.type == ContextType.SHARING:
- pass
- elif context.type == ContextType.FUNCTION or context.type == ContextType.FILE:
- pass
- else:
- logger.warning("[WX] unknown context type: {}".format(context.type))
- return
- return reply
-
- def _decorate_reply(self, context: Context, reply: Reply) -> Reply:
- if reply and reply.type:
- e_context = PluginManager().emit_event(
- EventContext(
- Event.ON_DECORATE_REPLY,
- {"channel": self, "context": context, "reply": reply},
- )
- )
- reply = e_context["reply"]
- desire_rtype = context.get("desire_rtype")
- if not e_context.is_pass() and reply and reply.type:
- if reply.type in self.NOT_SUPPORT_REPLYTYPE:
- logger.error("[WX]reply type not support: " + str(reply.type))
- reply.type = ReplyType.ERROR
- reply.content = "不支持发送的消息类型: " + str(reply.type)
-
- if reply.type == ReplyType.TEXT:
- reply_text = reply.content
- if desire_rtype == ReplyType.VOICE and ReplyType.VOICE not in self.NOT_SUPPORT_REPLYTYPE:
- reply = super().build_text_to_voice(reply.content)
- return self._decorate_reply(context, reply)
- if context.get("isgroup", False):
- if not context.get("no_need_at", False):
- reply_text = "@" + context["msg"].actual_user_nickname + "\n" + reply_text.strip()
- reply_text = conf().get("group_chat_reply_prefix", "") + reply_text + conf().get("group_chat_reply_suffix", "")
- else:
- reply_text = conf().get("single_chat_reply_prefix", "") + reply_text + conf().get("single_chat_reply_suffix", "")
- reply.content = reply_text
- elif reply.type == ReplyType.ERROR or reply.type == ReplyType.INFO:
- reply.content = "[" + str(reply.type) + "]\n" + reply.content
- elif reply.type == ReplyType.IMAGE_URL or reply.type == ReplyType.VOICE or reply.type == ReplyType.IMAGE or reply.type == ReplyType.FILE or reply.type == ReplyType.VIDEO or reply.type == ReplyType.VIDEO_URL:
- pass
- else:
- logger.error("[WX] unknown reply type: {}".format(reply.type))
- return
- if desire_rtype and desire_rtype != reply.type and reply.type not in [ReplyType.ERROR, ReplyType.INFO]:
- logger.warning("[WX] desire_rtype: {}, but reply type: {}".format(context.get("desire_rtype"), reply.type))
- return reply
-
- def _send_reply(self, context: Context, reply: Reply):
- if reply and reply.type:
- e_context = PluginManager().emit_event(
- EventContext(
- Event.ON_SEND_REPLY,
- {"channel": self, "context": context, "reply": reply},
- )
- )
- reply = e_context["reply"]
- if not e_context.is_pass() and reply and reply.type:
- logger.debug("[WX] ready to send reply: {}, context: {}".format(reply, context))
- self._send(reply, context)
-
- def _send(self, reply: Reply, context: Context, retry_cnt=0):
- try:
- self.send(reply, context)
- except Exception as e:
- logger.error("[WX] sendMsg error: {}".format(str(e)))
- if isinstance(e, NotImplementedError):
- return
- logger.exception(e)
- if retry_cnt < 2:
- time.sleep(3 + 3 * retry_cnt)
- self._send(reply, context, retry_cnt + 1)
-
- def _success_callback(self, session_id, **kwargs):
- logger.debug("Worker return success, session_id = {}".format(session_id))
-
- def _fail_callback(self, session_id, exception, **kwargs):
- logger.exception("Worker return exception: {}".format(exception))
-
- def _thread_pool_callback(self, session_id, **kwargs):
- def func(worker: Future):
- try:
- worker_exception = worker.exception()
- if worker_exception:
- self._fail_callback(session_id, exception=worker_exception, **kwargs)
- else:
- self._success_callback(session_id, **kwargs)
- except CancelledError as e:
- logger.info("Worker cancelled, session_id = {}".format(session_id))
- except Exception as e:
- logger.exception("Worker raise exception: {}".format(e))
- with self.lock:
- self.sessions[session_id][1].release()
-
- return func
-
- def produce(self, context: Context):
- session_id = context["session_id"]
- with self.lock:
- if session_id not in self.sessions:
- self.sessions[session_id] = [
- Dequeue(),
- threading.BoundedSemaphore(conf().get("concurrency_in_session", 4)),
- ]
- if context.type == ContextType.TEXT and context.content.startswith("#"):
- self.sessions[session_id][0].putleft(context)
- else:
- self.sessions[session_id][0].put(context)
-
-
- def consume(self):
- while True:
- with self.lock:
- session_ids = list(self.sessions.keys())
- for session_id in session_ids:
- context_queue, semaphore = self.sessions[session_id]
- if semaphore.acquire(blocking=False):
- 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.add_done_callback(self._thread_pool_callback(session_id, context=context))
- if session_id not in self.futures:
- self.futures[session_id] = []
- self.futures[session_id].append(future)
- elif semaphore._initial_value == semaphore._value + 1:
- self.futures[session_id] = [t for t in self.futures[session_id] if not t.done()]
- assert len(self.futures[session_id]) == 0, "thread pool error"
- del self.sessions[session_id]
- else:
- semaphore.release()
- time.sleep(0.1)
-
-
- def cancel_session(self, session_id):
- with self.lock:
- if session_id in self.sessions:
- for future in self.futures[session_id]:
- future.cancel()
- cnt = self.sessions[session_id][0].qsize()
- if cnt > 0:
- logger.info("Cancel {} messages in session {}".format(cnt, session_id))
- self.sessions[session_id][0] = Dequeue()
-
- def cancel_all_session(self):
- with self.lock:
- for session_id in self.sessions:
- for future in self.futures[session_id]:
- future.cancel()
- cnt = self.sessions[session_id][0].qsize()
- if cnt > 0:
- logger.info("Cancel {} messages in session {}".format(cnt, session_id))
- self.sessions[session_id][0] = Dequeue()
-
-
- def check_prefix(content, prefix_list):
- if not prefix_list:
- return None
- for prefix in prefix_list:
- if content.startswith(prefix):
- return prefix
- return None
-
-
- def check_contain(content, keyword_list):
- if not keyword_list:
- return None
- for ky in keyword_list:
- if content.find(ky) != -1:
- return True
- return None
|