From 4f3d12129c35787d1c54d1175d4f5e5c9e2dd390 Mon Sep 17 00:00:00 2001 From: lanvent Date: Fri, 31 Mar 2023 01:29:04 +0800 Subject: [PATCH] refactor: wrap wechat msg to make logic general --- channel/chat_message.py | 64 ++++++++++++++++++++++ channel/wechat/wechat_channel.py | 93 ++++++++++++++------------------ channel/wechat/wechat_message.py | 57 ++++++++++++++++++++ 3 files changed, 160 insertions(+), 54 deletions(-) create mode 100644 channel/chat_message.py create mode 100644 channel/wechat/wechat_message.py diff --git a/channel/chat_message.py b/channel/chat_message.py new file mode 100644 index 0000000..28ee054 --- /dev/null +++ b/channel/chat_message.py @@ -0,0 +1,64 @@ + +""" +本类表示聊天消息，用于对itchat和wechaty的消息进行统一的封装 +ChatMessage +msg_id: 消息id +create_time: 消息创建时间 + +ctype: 消息类型 : ContextType +content: 消息内容, 如果是声音/图片，这里是文件路径 + +from_user_id: 发送者id +from_user_nickname: 发送者昵称 +to_user_id: 接收者id +to_user_nickname: 接收者昵称 + +other_user_id: 对方的id，如果你是发送者，那这个就是接收者id，如果你是接收者，那这个就是发送者id +other_user_nickname: 同上 + +is_group: 是否是群消息 +is_at: 是否被at + +- (群消息时，一般会存在实际发送者，是群内某个成员的id和昵称，下列项仅在群消息时存在) +actual_user_id: 实际发送者id +actual_user_nickname：实际发送者昵称 + + + + +_prepare_fn: 准备函数，用于准备消息的内容，比如下载图片等, +_prepared: 是否已经调用过准备函数 +_rawmsg: 原始消息对象 + +""" +class ChatMessage(object): + msg_id = None + create_time = None + + ctype = None + content = None + + from_user_id = None + from_user_nickname = None + to_user_id = None + to_user_nickname = None + other_user_id = None + other_user_nickname = None + + is_group = False + is_at = False + actual_user_id = None + actual_user_nickname = None + + _prepare_fn = None + _prepared = False + _rawmsg = None + + + def __init__(self,_rawmsg): + self._rawmsg = _rawmsg + + def prepare(self): + if self._prepare_fn and not self._prepared: + self._prepared = True + self._prepare_fn() \ No newline at end of file diff --git a/channel/wechat/wechat_channel.py b/channel/wechat/wechat_channel.py index 2e7b075..cdbb333 100644 --- a/channel/wechat/wechat_channel.py +++ b/channel/wechat/wechat_channel.py @@ -9,6 +9,7 @@ import re import requests import io import time +from channel.wechat.wechat_message import * from common.singleton import singleton from lib import itchat import json @@ -38,36 +39,36 @@ def thread_pool_callback(worker): @itchat.msg_register(TEXT) def handler_single_msg(msg): - WechatChannel().handle_text(msg) + WechatChannel().handle_text(WeChatMessage(msg)) return None @itchat.msg_register(TEXT, isGroupChat=True) def handler_group_msg(msg): - WechatChannel().handle_group(msg) + WechatChannel().handle_group(WeChatMessage(msg,True)) return None @itchat.msg_register(VOICE) def handler_single_voice(msg): - WechatChannel().handle_voice(msg) + WechatChannel().handle_voice(WeChatMessage(msg)) return None @itchat.msg_register(VOICE, isGroupChat=True) def handler_group_voice(msg): - WechatChannel().handle_group_voice(msg) + WechatChannel().handle_group_voice(WeChatMessage(msg,True)) return None def _check(func): - def wrapper(self, msg): - msgId = msg['MsgId'] + def wrapper(self, cmsg: ChatMessage): + msgId = cmsg.msg_id if msgId in self.receivedMsgs: logger.info("Wechat message {} already received, ignore".format(msgId)) return - self.receivedMsgs[msgId] = msg - create_time = msg['CreateTime'] # 消息时间 + self.receivedMsgs[msgId] = cmsg + create_time = cmsg.create_time # 消息时间戳 if conf().get('hot_reload') == True and int(create_time) < int(time.time()) - 60: # 跳过1分钟前的历史消息 logger.debug("[WX]history message {} skipped".format(msgId)) return - return func(self, msg) + return func(self, cmsg) return wrapper @singleton @@ -112,57 +113,40 @@ class WechatChannel(Channel): @time_checker @_check - def handle_voice(self, msg): + def handle_voice(self, cmsg : ChatMessage): if conf().get('speech_recognition') != True: return - logger.debug("[WX]receive voice msg: " + msg['FileName']) - to_user_id = msg['ToUserName'] - from_user_id = msg['FromUserName'] - try: - other_user_id = msg['User']['UserName'] # 对手方id - except Exception as e: - logger.warn("[WX]get other_user_id failed: " + str(e)) - if from_user_id == self.userName: - other_user_id = to_user_id - else: - other_user_id = from_user_id + logger.debug("[WX]receive voice msg: {}".format(cmsg.content)) + from_user_id = cmsg.from_user_id + other_user_id = cmsg.other_user_id if from_user_id == other_user_id: - context = self._compose_context(ContextType.VOICE, msg['FileName'], isgroup=False, msg=msg, receiver=other_user_id, session_id=other_user_id) + context = self._compose_context(ContextType.VOICE, cmsg.content, isgroup=False, msg=cmsg, receiver=other_user_id, session_id=other_user_id) if context: thread_pool.submit(self.handle, context).add_done_callback(thread_pool_callback) @time_checker @_check - def handle_text(self, msg): - logger.debug("[WX]receive text msg: " + json.dumps(msg, ensure_ascii=False)) - content = msg['Text'] - from_user_id = msg['FromUserName'] - to_user_id = msg['ToUserName'] # 接收人id - try: - other_user_id = msg['User']['UserName'] # 对手方id - except Exception as e: - logger.warn("[WX]get other_user_id failed: " + str(e)) - if from_user_id == self.userName: - other_user_id = to_user_id - else: - other_user_id = from_user_id - if "」\n- - - - - - - - - - - - - - -" in content: + def handle_text(self, cmsg : ChatMessage): + logger.debug("[WX]receive text msg: " + json.dumps(cmsg._rawmsg, ensure_ascii=False)) + content = cmsg.content + other_user_id = cmsg.other_user_id + if "」\n- - - - - - - - - - - - - - -" in cmsg.content: logger.debug("[WX]reference query skipped") return - context = self._compose_context(ContextType.TEXT, content, isgroup=False, msg=msg, receiver=other_user_id, session_id=other_user_id) + context = self._compose_context(ContextType.TEXT, content, isgroup=False, msg=cmsg, receiver=other_user_id, session_id=other_user_id) if context: thread_pool.submit(self.handle, context).add_done_callback(thread_pool_callback) @time_checker @_check - def handle_group(self, msg): - logger.debug("[WX]receive group msg: " + json.dumps(msg, ensure_ascii=False)) - group_name = msg['User'].get('NickName', None) - group_id = msg['User'].get('UserName', None) + def handle_group(self, cmsg : ChatMessage): + logger.debug("[WX]receive group msg: " + json.dumps(cmsg._rawmsg, ensure_ascii=False)) + group_name = cmsg.other_user_nickname + group_id = cmsg.other_user_id if not group_name: return "" - content = msg.content + content = cmsg.content if "」\n- - - - - - - - - - - - - - -" in content: logger.debug("[WX]reference query skipped") return "" @@ -175,21 +159,22 @@ class WechatChannel(Channel): 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 = msg['ActualUserName'] + 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 - context = self._compose_context(ContextType.TEXT, content, isgroup=True, msg=msg, receiver=group_id, session_id=session_id) + + context = self._compose_context(ContextType.TEXT, content, isgroup=True, msg=cmsg, receiver=group_id, session_id=session_id) if context: thread_pool.submit(self.handle, context).add_done_callback(thread_pool_callback) @time_checker @_check - def handle_group_voice(self, msg): + def handle_group_voice(self, cmsg : ChatMessage): if conf().get('group_speech_recognition', False) != True: return - logger.debug("[WX]receive voice for group msg: " + msg['FileName']) - group_name = msg['User'].get('NickName', None) - group_id = msg['User'].get('UserName', None) + logger.debug("[WX]receive voice for group msg: {}".format(cmsg.content)) + group_name = cmsg.other_user_nickname + group_id = cmsg.other_user_id # 验证群名 if not group_name: return "" @@ -199,10 +184,10 @@ class WechatChannel(Channel): 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 =msg['ActualUserName'] + 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 - context = self._compose_context(ContextType.VOICE, msg['FileName'], isgroup=True, msg=msg, receiver=group_id, session_id=session_id) + context = self._compose_context(ContextType.VOICE, cmsg.content, isgroup=True, msg=cmsg, receiver=group_id, session_id=session_id) if context: thread_pool.submit(self.handle, context).add_done_callback(thread_pool_callback) @@ -222,7 +207,7 @@ class WechatChannel(Channel): # 判断如果匹配到自定义前缀，则返回过滤掉前缀+空格后的内容，用于实现类似自定义+前缀触发生成AI图片的功能 if match_prefix: content = content.replace(match_prefix, '', 1).strip() - elif context['msg']['IsAt'] and not conf().get("group_at_off", False): + elif context['msg'].is_at and not conf().get("group_at_off", False): logger.info("[WX]receive group at, continue") elif context["origin_ctype"] == ContextType.VOICE: logger.info("[WX]receive group voice, checkprefix didn't match") @@ -306,8 +291,8 @@ class WechatChannel(Channel): reply = super().build_reply_content(context.content, context) elif context.type == ContextType.VOICE: # 语音消息 msg = context['msg'] - mp3_path = TmpDir().path() + context.content - msg.download(mp3_path) + msg.prepare() + mp3_path = context.content # mp3转wav wav_path = os.path.splitext(mp3_path)[0] + '.wav' try: @@ -349,7 +334,7 @@ class WechatChannel(Channel): reply = super().build_text_to_voice(reply.content) return self._decorate_reply(context, reply) if context['isgroup']: - reply_text = '@' + context['msg']['ActualNickName'] + ' ' + reply_text.strip() + reply_text = '@' + context['msg'].actual_user_nickname + ' ' + reply_text.strip() reply_text = conf().get("group_chat_reply_prefix", "")+reply_text else: reply_text = conf().get("single_chat_reply_prefix", "")+reply_text diff --git a/channel/wechat/wechat_message.py b/channel/wechat/wechat_message.py new file mode 100644 index 0000000..48a5867 --- /dev/null +++ b/channel/wechat/wechat_message.py @@ -0,0 +1,57 @@ + + +from bridge.context import ContextType +from channel.chat_message import ChatMessage +from common.tmp_dir import TmpDir +from common.log import logger +from lib.itchat.content import * +from lib import itchat + +class WeChatMessage(ChatMessage): + + def __init__(self, itchat_msg, is_group=False): + super().__init__( itchat_msg) + self.msg_id = itchat_msg['MsgId'] + self.create_time = itchat_msg['CreateTime'] + self.is_group = is_group + + if itchat_msg['Type'] == TEXT: + self.ctype = ContextType.TEXT + self.content = itchat_msg['Text'] + elif itchat_msg['Type'] == VOICE: + self.ctype = ContextType.VOICE + self.content = TmpDir().path() + itchat_msg['FileName'] # content直接存临时目录路径 + self._prepare_fn = lambda: itchat_msg.download(self.content) + else: + raise NotImplementedError("Unsupported message type: {}".format(itchat_msg['Type'])) + + self.from_user_id = itchat_msg['FromUserName'] + self.to_user_id = itchat_msg['ToUserName'] + + user_id = itchat.instance.storageClass.userName + nickname = itchat.instance.storageClass.nickName + + # 虽然from_user_id和to_user_id用的少，但是为了保持一致性，还是要填充一下 + # 以下很繁琐，一句话总结：能填的都填了。 + if self.from_user_id == user_id: + self.from_user_nickname = nickname + if self.to_user_id == user_id: + self.to_user_nickname = nickname + try: # 陌生人时候, 'User'字段可能不存在 + self.other_user_id = itchat_msg['User']['UserName'] + self.other_user_nickname = itchat_msg['User']['NickName'] + if self.other_user_id == self.from_user_id: + self.from_user_nickname = self.other_user_nickname + if self.other_user_id == self.to_user_id: + self.to_user_nickname = self.other_user_nickname + except KeyError as e: + logger.warn("[WX]get other_user_id failed: " + str(e)) + if self.from_user_id == user_id: + self.other_user_id = self.to_user_id + else: + self.other_user_id = self.from_user_id + + if self.is_group: + self.is_at = itchat_msg['IsAt'] + self.actual_user_id = itchat_msg['ActualUserName'] + self.actual_user_nickname = itchat_msg['ActualNickName']