457 lines
20 KiB
Python
457 lines
20 KiB
Python
from enum import Enum
|
||
from maim_message import MessageBase, Seg
|
||
from typing import List, Tuple, Optional, Dict, TYPE_CHECKING
|
||
|
||
import base64
|
||
import hashlib
|
||
import msgpack
|
||
import random
|
||
import re
|
||
|
||
from src.common.data_models.message_component_data_model import (
|
||
MessageSequence,
|
||
StandardMessageComponents,
|
||
TextComponent,
|
||
ImageComponent,
|
||
EmojiComponent,
|
||
VoiceComponent,
|
||
AtComponent,
|
||
ReplyComponent,
|
||
DictComponent,
|
||
UnknownUser,
|
||
ForwardNodeComponent,
|
||
)
|
||
from src.config.config import global_config
|
||
|
||
from .math_utils import number_to_short_id, TimestampMode, translate_timestamp_to_human_readable
|
||
|
||
if TYPE_CHECKING:
|
||
from src.chat.message_receive.message import SessionMessage
|
||
|
||
|
||
class MessageUtils:
|
||
@staticmethod
|
||
def from_db_record_msg_to_MaiSeq(raw_content: bytes) -> MessageSequence:
|
||
unpacked_data = msgpack.unpackb(raw_content)
|
||
return MessageSequence.from_dict(unpacked_data)
|
||
|
||
@staticmethod
|
||
def from_MaiSeq_to_db_record_msg(msg: MessageSequence) -> bytes:
|
||
dict_representation = msg.to_dict()
|
||
return msgpack.packb(dict_representation) # type: ignore
|
||
|
||
@staticmethod
|
||
def from_maim_message_segments_to_MaiSeq(message: "MessageBase") -> MessageSequence:
|
||
"""从maim_message.MessageBase.message_segment转换为MessageSequence"""
|
||
raw_msg_seq = message.message_segment
|
||
components: List[StandardMessageComponents] = []
|
||
if not raw_msg_seq:
|
||
return MessageSequence(components)
|
||
if raw_msg_seq.type == "seglist":
|
||
assert isinstance(raw_msg_seq.data, list), "seglist类型的message_segment数据应该是一个列表"
|
||
components.extend(MessageUtils._parse_maim_message_segment_to_component(item) for item in raw_msg_seq.data)
|
||
elif raw_msg_seq.type in {"text", "image", "emoji", "voice", "at", "reply"}:
|
||
components.append(MessageUtils._parse_maim_message_segment_to_component(raw_msg_seq))
|
||
else:
|
||
raise NotImplementedError(f"暂时不支持的消息片段类型: {raw_msg_seq.type}")
|
||
return MessageSequence(components)
|
||
|
||
@staticmethod
|
||
async def from_MaiSeq_to_maim_message_segments(msg_seq: MessageSequence) -> List[Seg]:
|
||
"""从MessageSequence转换为maim_message.MessageBase.message_segment格式的列表"""
|
||
segments = []
|
||
for component in msg_seq.components:
|
||
if isinstance(component, DictComponent):
|
||
seg = Seg(type="dict", data=component.data) # type: ignore
|
||
else:
|
||
seg = await component.to_seg()
|
||
segments.append(seg)
|
||
return segments
|
||
|
||
@staticmethod
|
||
def _parse_maim_message_segment_to_component(seg: Seg) -> "StandardMessageComponents":
|
||
if seg.type == "text":
|
||
assert isinstance(seg.data, str), "text类型的seg数据应该是字符串"
|
||
return TextComponent(text=seg.data)
|
||
elif seg.type == "image":
|
||
assert isinstance(seg.data, str), "image类型的seg数据应该是base64字符串"
|
||
image_bytes = base64.b64decode(seg.data)
|
||
binary_hash = hashlib.md5(image_bytes).hexdigest()
|
||
return ImageComponent(binary_hash=binary_hash, binary_data=image_bytes)
|
||
elif seg.type == "emoji":
|
||
assert isinstance(seg.data, str), "emoji类型的seg数据应该是base64字符串"
|
||
emoji_bytes = base64.b64decode(seg.data)
|
||
binary_hash = hashlib.md5(emoji_bytes).hexdigest()
|
||
return EmojiComponent(binary_hash=binary_hash, binary_data=emoji_bytes)
|
||
elif seg.type == "voice":
|
||
assert isinstance(seg.data, str), "voice类型的seg数据应该是base64字符串"
|
||
voice_bytes = base64.b64decode(seg.data)
|
||
binary_hash = hashlib.md5(voice_bytes).hexdigest()
|
||
return VoiceComponent(binary_hash=binary_hash, binary_data=voice_bytes)
|
||
elif seg.type == "at":
|
||
assert isinstance(seg.data, str), "at类型的seg数据应该是字符串"
|
||
return AtComponent(target_user_id=seg.data)
|
||
elif seg.type == "reply":
|
||
assert isinstance(seg.data, str), "reply类型的seg数据应该是字符串"
|
||
return ReplyComponent(target_message_id=seg.data)
|
||
else:
|
||
raise NotImplementedError(f"暂时不支持的消息片段类型: {seg.type}")
|
||
|
||
@staticmethod
|
||
def check_ban_words(text: str) -> Tuple[bool, Optional[str]]:
|
||
"""检查消息是否包含过滤词
|
||
|
||
Args:
|
||
text: 待检查的文本
|
||
|
||
Returns:
|
||
bool: 是否包含过滤词
|
||
"""
|
||
if not text:
|
||
return False, None
|
||
return next(
|
||
((True, word) for word in global_config.message_receive.ban_words if word in text),
|
||
(False, None),
|
||
)
|
||
|
||
@staticmethod
|
||
def check_ban_regex(text: str) -> Tuple[bool, Optional[str]]:
|
||
"""检查消息是否匹配过滤正则表达式
|
||
|
||
Args:
|
||
text: 待检查的文本
|
||
chat: 聊天对象
|
||
userinfo: 用户信息
|
||
|
||
Returns:
|
||
bool: 是否匹配过滤正则
|
||
"""
|
||
# 检查text是否为None或空字符串
|
||
if not text:
|
||
return False, None
|
||
return next(
|
||
((True, pattern) for pattern in global_config.message_receive.ban_msgs_regex if re.search(pattern, text)),
|
||
(False, None),
|
||
)
|
||
|
||
@staticmethod
|
||
def store_message_to_db(message: "SessionMessage"):
|
||
"""存储消息到数据库"""
|
||
from src.common.database.database import get_db_session
|
||
|
||
with get_db_session() as session:
|
||
db_message = message.to_db_instance()
|
||
session.add(db_message)
|
||
|
||
@staticmethod
|
||
async def build_readable_message(
|
||
messages: List["SessionMessage"],
|
||
*,
|
||
anonymize: bool = False,
|
||
show_lineno: bool = False,
|
||
extract_pictures: bool = False,
|
||
replace_bot_name: bool = False,
|
||
target_bot_name: Optional[str] = None,
|
||
timestamp_mode: Optional[TimestampMode] = None,
|
||
show_message_id_prefix: bool = False,
|
||
truncate_message: bool = False,
|
||
) -> Tuple[str, Dict[str, Tuple[str, str]]]:
|
||
"""
|
||
将消息构建为LLM可读的文本格式
|
||
|
||
Args:
|
||
messages (List[SessionMessage]): 消息列表
|
||
anonymize (bool): 是否匿名化用户信息
|
||
show_lineno (bool): 是否在每条消息前显示行号
|
||
extract_pictures (bool): 是否提取图片信息并在文本中显示占位符
|
||
replace_bot_name (bool): 是否将消息中的机器人名称替换为统一的占位符
|
||
target_bot_name (Optional[str]): 如果replace_bot_name为True，指定要替换的机器人名称，比如可以把机器人名称替换为“你”
|
||
timestamp_mode (Optional[TimestampMode]): 时间戳显示模式，默认为None表示不显示时间戳
|
||
show_message_id_prefix (bool): 是否在每条消息前显示消息ID前缀
|
||
truncate_message (bool): 是否截断过长的消息文本，避免生成过长的输入给LLM
|
||
Returns:
|
||
return (Tuple[str, Dict[str, Tuple[str, str]]]): 构建后的消息文本，以及映射表（匿名ID, 原始名称）
|
||
"""
|
||
msg_list: List["SessionMessage"] = messages
|
||
user_id_mapping: Dict[str, Tuple[str, str]] = {} # user_id -> (匿名ID, 原始名称)
|
||
copied: bool = False # 标记是否已经复制过消息列表，避免不必要的复制开销
|
||
img_map: Optional[Dict[str, Tuple[int, str]]] = None
|
||
emoji_map: Optional[Dict[str, Tuple[int, str]]] = None
|
||
if replace_bot_name and not target_bot_name:
|
||
raise ValueError("当replace_bot_name为True时，必须指定target_bot_name参数")
|
||
if anonymize or replace_bot_name:
|
||
user_id_mapping = {} # 利用弱引用直接传入并得到修改结果
|
||
anonymous_messages: List["SessionMessage"] = []
|
||
salt_str = str(random.randint(100000, 999999)) # 每次调用生成一个随机盐，确保匿名ID不可预测
|
||
anonymous_messages.extend(
|
||
MessageUtils._process_usr_info(
|
||
msg,
|
||
user_id_mapping,
|
||
salt_str,
|
||
anonymize,
|
||
replace_bot_name,
|
||
target_bot_name,
|
||
)
|
||
for msg in messages
|
||
)
|
||
msg_list = anonymous_messages
|
||
copied = True
|
||
|
||
processed_plain_texts: List[str] = []
|
||
if extract_pictures:
|
||
img_map = {} # binary_hash -> (图片ID, 描述信息)
|
||
emoji_map = {} # binary_hash -> (表情ID, 描述信息)
|
||
msg_list = [
|
||
MessageUtils._extract_pictures_from_message(msg, img_map, emoji_map, copied) for msg in msg_list
|
||
]
|
||
processed_plain_texts.extend(f"[图片{img_id}: {desc}]" for img_id, desc in img_map.values())
|
||
processed_plain_texts.append("") # 图片和表情之间添加一个换行，避免连在一起
|
||
processed_plain_texts.extend(f"[表情{emoji_id}: {desc}]" for emoji_id, desc in emoji_map.values())
|
||
processed_plain_texts.append("") # 表情和消息文本之间添加两个换行，避免连在一起
|
||
|
||
lineno_counter = 1
|
||
for msg in msg_list:
|
||
await msg.process()
|
||
plain_text: str = msg.processed_plain_text # type: ignore
|
||
usr_info = msg.message_info.user_info
|
||
usr_name = usr_info.user_cardname or usr_info.user_nickname or "未知用户"
|
||
header = f"[{lineno_counter}] {usr_name}说：" if show_lineno else f"{usr_name}说："
|
||
lineno_counter += 1
|
||
processed_plain_texts.append("".join([header, plain_text]))
|
||
|
||
return "\n".join(processed_plain_texts), user_id_mapping
|
||
|
||
@staticmethod
|
||
def _process_usr_info(
|
||
message: "SessionMessage",
|
||
anonymize_mapping: Dict[str, Tuple[str, str]],
|
||
salt: str,
|
||
anonymize: bool,
|
||
replace_bot_name: bool,
|
||
target_bot_name: Optional[str] = None,
|
||
):
|
||
"""处理消息中的用户信息，进行匿名化显示"""
|
||
new_message = message.deepcopy()
|
||
new_component_list = [
|
||
MessageUtils._process_msg_component(
|
||
component,
|
||
anonymize_mapping,
|
||
salt,
|
||
anonymize,
|
||
replace_bot_name,
|
||
target_bot_name,
|
||
)
|
||
for component in new_message.raw_message.components
|
||
]
|
||
new_message.raw_message.components = new_component_list
|
||
msg_usr_info = message.message_info.user_info
|
||
if anonymize:
|
||
if msg_usr_info.user_id not in anonymize_mapping:
|
||
num = len(anonymize_mapping) + 1
|
||
anonymous_id = number_to_short_id(num, salt, length=6)
|
||
original_name = msg_usr_info.user_cardname or msg_usr_info.user_nickname or msg_usr_info.user_id
|
||
anonymize_mapping[msg_usr_info.user_id] = (anonymous_id, original_name)
|
||
anonymous_name = anonymize_mapping[msg_usr_info.user_id][0]
|
||
new_message.message_info.user_info.user_nickname = anonymous_name
|
||
new_message.message_info.user_info.user_cardname = anonymous_name
|
||
if replace_bot_name and target_bot_name and is_bot_self(msg_usr_info.user_id):
|
||
new_message.message_info.user_info.user_nickname = target_bot_name
|
||
new_message.message_info.user_info.user_cardname = target_bot_name
|
||
return new_message
|
||
|
||
@staticmethod
|
||
def _process_msg_component(
|
||
component: StandardMessageComponents,
|
||
anonymize_mapping: Dict[str, Tuple[str, str]],
|
||
salt: str,
|
||
anonymize: bool,
|
||
replace_bot_name: bool,
|
||
target_bot_name: Optional[str] = None,
|
||
) -> StandardMessageComponents:
|
||
"""将消息组件中的用户信息匿名化"""
|
||
if isinstance(component, AtComponent):
|
||
return MessageUtils.__handle_at_component(
|
||
component,
|
||
anonymize_mapping,
|
||
salt,
|
||
anonymize,
|
||
replace_bot_name,
|
||
target_bot_name,
|
||
)
|
||
elif isinstance(component, ReplyComponent):
|
||
return MessageUtils.__handle_reply_component(
|
||
component,
|
||
anonymize_mapping,
|
||
salt,
|
||
anonymize,
|
||
replace_bot_name,
|
||
target_bot_name,
|
||
)
|
||
elif isinstance(component, ForwardNodeComponent):
|
||
return MessageUtils.__handle_forward_node_component(
|
||
component,
|
||
anonymize_mapping,
|
||
salt,
|
||
anonymize,
|
||
replace_bot_name,
|
||
target_bot_name,
|
||
)
|
||
return component
|
||
|
||
@staticmethod
|
||
def __handle_at_component(
|
||
component: AtComponent,
|
||
anonymize_mapping: Dict[str, Tuple[str, str]],
|
||
salt: str,
|
||
anonymize: bool,
|
||
replace_bot_name: bool,
|
||
target_bot_name: Optional[str] = None,
|
||
):
|
||
user_id = component.target_user_id # user_id一定存在
|
||
if anonymize:
|
||
if user_id not in anonymize_mapping:
|
||
# 新人物? 编号 + 1，生成一个新的匿名ID
|
||
num = len(anonymize_mapping) + 1
|
||
anonymous_id = number_to_short_id(num, salt, length=6)
|
||
original_name = component.target_user_cardname or component.target_user_nickname or user_id
|
||
anonymize_mapping[user_id] = (anonymous_id, original_name)
|
||
# 替换昵称和备注为匿名ID
|
||
anonymous_name = anonymize_mapping[user_id][0]
|
||
component.target_user_nickname = anonymous_name
|
||
component.target_user_cardname = anonymous_name
|
||
if replace_bot_name and target_bot_name and is_bot_self(user_id):
|
||
component.target_user_nickname = target_bot_name
|
||
component.target_user_cardname = target_bot_name
|
||
return component
|
||
|
||
@staticmethod
|
||
def __handle_forward_node_component(
|
||
component: ForwardNodeComponent,
|
||
anonymize_mapping: Dict[str, Tuple[str, str]],
|
||
salt: str,
|
||
anonymize: bool,
|
||
replace_bot_name: bool,
|
||
target_bot_name: Optional[str] = None,
|
||
):
|
||
for comp in component.forward_components:
|
||
user_id = comp.user_id
|
||
if not user_id: # 如果转发节点的用户ID不存在，直接设置为未知用户
|
||
comp.user_id = "unknown_user"
|
||
comp.user_cardname = "未知用户"
|
||
comp.user_nickname = "未知用户"
|
||
continue
|
||
if isinstance(user_id, UnknownUser): # 如果用户ID是UnknownUser类型，直接设置为未知用户
|
||
comp.user_id = "unknown_user"
|
||
comp.user_cardname = "未知用户"
|
||
comp.user_nickname = "未知用户"
|
||
continue
|
||
if anonymize:
|
||
if user_id not in anonymize_mapping:
|
||
num = len(anonymize_mapping) + 1
|
||
anonymous_id = number_to_short_id(num, salt, length=6)
|
||
original_name = comp.user_cardname or comp.user_nickname or user_id
|
||
anonymize_mapping[user_id] = (anonymous_id, original_name)
|
||
anonymous_name = anonymize_mapping[user_id][0]
|
||
comp.user_nickname = anonymous_name
|
||
comp.user_cardname = anonymous_name
|
||
if replace_bot_name and target_bot_name and is_bot_self(user_id):
|
||
comp.user_nickname = target_bot_name
|
||
comp.user_cardname = target_bot_name
|
||
comp.content = [ # 递归处理转发消息中的组件
|
||
MessageUtils._process_msg_component(
|
||
c,
|
||
anonymize_mapping,
|
||
salt,
|
||
anonymize,
|
||
replace_bot_name,
|
||
target_bot_name,
|
||
)
|
||
for c in comp.content
|
||
]
|
||
return component
|
||
|
||
@staticmethod
|
||
def __handle_reply_component(
|
||
component: ReplyComponent,
|
||
anonymize_mapping: Dict[str, Tuple[str, str]],
|
||
salt: str,
|
||
anonymize: bool,
|
||
replace_bot_name: bool,
|
||
target_bot_name: Optional[str] = None,
|
||
):
|
||
if user_id := component.target_message_sender_id:
|
||
if anonymize:
|
||
if user_id not in anonymize_mapping:
|
||
num = len(anonymize_mapping) + 1
|
||
anonymous_id = number_to_short_id(num, salt, length=6)
|
||
original_name = (
|
||
component.target_message_sender_cardname or component.target_message_sender_nickname or user_id
|
||
)
|
||
anonymize_mapping[user_id] = (anonymous_id, original_name)
|
||
anonymous_name = anonymize_mapping[user_id][0]
|
||
component.target_message_sender_nickname = anonymous_name
|
||
component.target_message_sender_cardname = anonymous_name
|
||
if replace_bot_name and target_bot_name and is_bot_self(user_id):
|
||
component.target_message_sender_nickname = target_bot_name
|
||
component.target_message_sender_cardname = target_bot_name
|
||
else:
|
||
component.target_message_sender_nickname = "未知用户" # 如果没有Reply消息的发送者ID，直接设置为未知用户
|
||
component.target_message_sender_cardname = "未知用户"
|
||
return component
|
||
|
||
@staticmethod
|
||
def _extract_pictures_from_message(
|
||
message: "SessionMessage",
|
||
img_map: Dict[str, Tuple[int, str]],
|
||
emoji_map: Dict[str, Tuple[int, str]],
|
||
copied: bool,
|
||
):
|
||
"""从消息中提取图片组件，返回列表包含(图片ID, 描述信息)"""
|
||
if not copied:
|
||
message = message.deepcopy() # 避免修改原消息
|
||
new_component_list: List[StandardMessageComponents] = []
|
||
new_component_list.extend(
|
||
MessageUtils._extract_pictures_from_component(component, img_map, emoji_map)
|
||
for component in message.raw_message.components
|
||
)
|
||
message.raw_message.components = new_component_list
|
||
return message
|
||
|
||
@staticmethod
|
||
def _extract_pictures_from_component(
|
||
component: StandardMessageComponents,
|
||
img_map: Dict[str, Tuple[int, str]],
|
||
emoji_map: Dict[str, Tuple[int, str]],
|
||
) -> StandardMessageComponents:
|
||
"""从消息组件中提取图片信息"""
|
||
if isinstance(component, ImageComponent):
|
||
if component.binary_hash in img_map:
|
||
img_id, _ = img_map[component.binary_hash]
|
||
else:
|
||
img_id = len(img_map) + 1
|
||
img_map[component.binary_hash] = (img_id, component.content)
|
||
component.content = f"图片{img_id}"
|
||
elif isinstance(component, EmojiComponent):
|
||
if component.binary_hash in emoji_map:
|
||
emoji_id, _ = emoji_map[component.binary_hash]
|
||
else:
|
||
emoji_id = len(emoji_map) + 1
|
||
emoji_map[component.binary_hash] = (emoji_id, component.content)
|
||
component.content = f"表情{emoji_id}"
|
||
elif isinstance(component, ForwardNodeComponent):
|
||
for comp in component.forward_components:
|
||
comp.content = [
|
||
MessageUtils._extract_pictures_from_component(c, img_map, emoji_map) for c in comp.content
|
||
]
|
||
return component
|
||
|
||
|
||
# TODO: 这个函数的实现非常临时，后续需要替换为更完善的实现，比如直接从配置文件中读取机器人自己的ID，或者通过API获取机器人自己的信息等
|
||
def is_bot_self(user_id: str) -> bool:
|
||
"""
|
||
判断用户ID是否是机器人自己
|
||
|
||
临时方法，后续会替换为更完善的实现
|
||
"""
|
||
return user_id == "bot_self"
|