From fe3684736a026fc83eb281bb836ec895157b598e Mon Sep 17 00:00:00 2001 From: tcmofashi Date: Sun, 9 Mar 2025 10:27:54 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E8=B6=85=E5=A4=A7=E5=9E=8B=E9=87=8D?= =?UTF-8?q?=E6=9E=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/plugins/chat/bot.py | 6 +- src/plugins/chat/cq_code.py | 278 +++++++-------- src/plugins/chat/emoji_manager.py | 104 ++++-- src/plugins/chat/llm_generator.py | 2 +- src/plugins/chat/message.py | 469 ++++++++++++++++--------- src/plugins/chat/message_base.py | 158 +++++++++ src/plugins/chat/message_cq.py | 188 ++++++++++ src/plugins/chat/message_sender.py | 2 +- src/plugins/chat/storage.py | 2 +- src/plugins/chat/utils.py | 23 +- src/plugins/chat/utils_image.py | 543 ++++++++++++++++++----------- 11 files changed, 1209 insertions(+), 566 deletions(-) create mode 100644 src/plugins/chat/message_base.py create mode 100644 src/plugins/chat/message_cq.py diff --git a/src/plugins/chat/bot.py b/src/plugins/chat/bot.py index a02c4a059..c9eea3ed9 100644 --- a/src/plugins/chat/bot.py +++ b/src/plugins/chat/bot.py @@ -7,10 +7,10 @@ from nonebot.adapters.onebot.v11 import Bot, GroupMessageEvent from ..memory_system.memory import hippocampus from ..moods.moods import MoodManager # 导入情绪管理器 from .config import global_config -from .cq_code import CQCode # 导入CQCode模块 +from .cq_code import CQCode,cq_code_tool # 导入CQCode模块 from .emoji_manager import emoji_manager # 导入表情包管理器 from .llm_generator import ResponseGenerator -from .message import ( +from .message_cq import ( Message, Message_Sending, Message_Thinking, # 导入 Message_Thinking 类 @@ -180,7 +180,7 @@ class ChatBot: if emoji_raw != None: emoji_path,discription = emoji_raw - emoji_cq = CQCode.create_emoji_cq(emoji_path) + emoji_cq = cq_code_tool.create_emoji_cq(emoji_path) if random() < 0.5: bot_response_time = tinking_time_point - 1 diff --git a/src/plugins/chat/cq_code.py b/src/plugins/chat/cq_code.py index 4a295e3d5..e908219b5 100644 --- a/src/plugins/chat/cq_code.py +++ b/src/plugins/chat/cq_code.py @@ -3,7 +3,7 @@ import html import os import time from dataclasses import dataclass -from typing import Dict, Optional +from typing import Dict, Optional, List, Union import requests @@ -12,12 +12,14 @@ import requests import urllib3 from nonebot import get_driver from urllib3.util import create_urllib3_context +from loguru import logger from ..models.utils_model import LLM_request from .config import global_config from .mapper import emojimapper -from .utils_image import storage_emoji, storage_image +from .utils_image import image_manager from .utils_user import get_user_nickname +from .message_base import Seg driver = get_driver() config = driver.config @@ -48,16 +50,15 @@ class CQCode: type: CQ码类型（如'image', 'at', 'face'等） params: CQ码的参数字典 raw_code: 原始CQ码字符串 - translated_plain_text: 经过处理（如AI翻译）后的文本表示 + translated_segments: 经过处理后的Seg对象列表 """ type: str params: Dict[str, str] - # raw_code: str group_id: int user_id: int group_name: str = "" user_nickname: str = "" - translated_plain_text: Optional[str] = None + translated_segments: Optional[Union[Seg, List[Seg]]] = None reply_message: Dict = None # 存储回复消息 image_base64: Optional[str] = None _llm: Optional[LLM_request] = None @@ -66,31 +67,72 @@ class CQCode: """初始化LLM实例""" self._llm = LLM_request(model=global_config.vlm, temperature=0.4, max_tokens=300) - async def translate(self): - """根据CQ码类型进行相应的翻译处理""" + def translate(self): + """根据CQ码类型进行相应的翻译处理，转换为Seg对象""" if self.type == 'text': - self.translated_plain_text = self.params.get('text', '') + self.translated_segments = Seg( + type='text', + data=self.params.get('text', '') + ) elif self.type == 'image': - if self.params.get('sub_type') == '0': - self.translated_plain_text = await self.translate_image() + base64_data = self.translate_image() + if base64_data: + if self.params.get('sub_type') == '0': + self.translated_segments = Seg( + type='image', + data=base64_data + ) + else: + self.translated_segments = Seg( + type='emoji', + data=base64_data + ) else: - self.translated_plain_text = await self.translate_emoji() + self.translated_segments = Seg( + type='text', + data='[图片]' + ) elif self.type == 'at': user_nickname = get_user_nickname(self.params.get('qq', '')) - if user_nickname: - self.translated_plain_text = f"[@{user_nickname}]" - else: - self.translated_plain_text = "@某人" + self.translated_segments = Seg( + type='text', + data=f"[@{user_nickname or '某人'}]" + ) elif self.type == 'reply': - self.translated_plain_text = await self.translate_reply() + reply_segments = self.translate_reply() + if reply_segments: + self.translated_segments = Seg( + type='seglist', + data=reply_segments + ) + else: + self.translated_segments = Seg( + type='text', + data='[回复某人消息]' + ) elif self.type == 'face': face_id = self.params.get('id', '') - # self.translated_plain_text = f"[表情{face_id}]" - self.translated_plain_text = f"[{emojimapper.get(int(face_id), '表情')}]" + self.translated_segments = Seg( + type='text', + data=f"[{emojimapper.get(int(face_id), '表情')}]" + ) elif self.type == 'forward': - self.translated_plain_text = await self.translate_forward() + forward_segments = self.translate_forward() + if forward_segments: + self.translated_segments = Seg( + type='seglist', + data=forward_segments + ) + else: + self.translated_segments = Seg( + type='text', + data='[转发消息]' + ) else: - self.translated_plain_text = f"[{self.type}]" + self.translated_segments = Seg( + type='text', + data=f"[{self.type}]" + ) def get_img(self): ''' @@ -160,155 +202,101 @@ class CQCode: return None - async def translate_emoji(self) -> str: - """处理表情包类型的CQ码""" + + def translate_image(self) -> Optional[str]: + """处理图片类型的CQ码，返回base64字符串""" if 'url' not in self.params: - return '[表情包]' - base64_str = self.get_img() - if base64_str: - # 将 base64 字符串转换为字节类型 - image_bytes = base64.b64decode(base64_str) - storage_emoji(image_bytes) - return await self.get_emoji_description(base64_str) - else: - return '[表情包]' + return None + return self.get_img() - async def translate_image(self) -> str: - """处理图片类型的CQ码，区分普通图片和表情包""" - # 没有url，直接返回默认文本 - if 'url' not in self.params: - return '[图片]' - base64_str = self.get_img() - if base64_str: - image_bytes = base64.b64decode(base64_str) - storage_image(image_bytes) - return await self.get_image_description(base64_str) - else: - return '[图片]' - - async def get_emoji_description(self, image_base64: str) -> str: - """调用AI接口获取表情包描述""" - try: - prompt = "这是一个表情包，请用简短的中文描述这个表情包传达的情感和含义。最多20个字。" - # description, _ = self._llm.generate_response_for_image_sync(prompt, image_base64) - description, _ = await self._llm.generate_response_for_image(prompt, image_base64) - return f"[表情包：{description}]" - except Exception as e: - print(f"\033[1;31m[错误]\033[0m AI接口调用失败: {str(e)}") - return "[表情包]" - - async def get_image_description(self, image_base64: str) -> str: - """调用AI接口获取普通图片描述""" - try: - prompt = "请用中文描述这张图片的内容。如果有文字，请把文字都描述出来。并尝试猜测这个图片的含义。最多200个字。" - # description, _ = self._llm.generate_response_for_image_sync(prompt, image_base64) - description, _ = await self._llm.generate_response_for_image(prompt, image_base64) - return f"[图片：{description}]" - except Exception as e: - print(f"\033[1;31m[错误]\033[0m AI接口调用失败: {str(e)}") - return "[图片]" - - async def translate_forward(self) -> str: - """处理转发消息""" + def translate_forward(self) -> Optional[List[Seg]]: + """处理转发消息，返回Seg列表""" try: if 'content' not in self.params: - return '[转发消息]' + return None - # 解析content内容（需要先反转义） content = self.unescape(self.params['content']) - # print(f"\033[1;34m[调试信息]\033[0m 转发消息内容: {content}") - # 将字符串形式的列表转换为Python对象 import ast try: messages = ast.literal_eval(content) except ValueError as e: - print(f"\033[1;31m[错误]\033[0m 解析转发消息内容失败: {str(e)}") - return '[转发消息]' + logger.error(f"解析转发消息内容失败: {str(e)}") + return None - # 处理每条消息 - formatted_messages = [] + formatted_segments = [] for msg in messages: sender = msg.get('sender', {}) nickname = sender.get('card') or sender.get('nickname', '未知用户') - - # 获取消息内容并使用Message类处理 raw_message = msg.get('raw_message', '') message_array = msg.get('message', []) if message_array and isinstance(message_array, list): - # 检查是否包含嵌套的转发消息 for message_part in message_array: if message_part.get('type') == 'forward': - content = '[转发消息]' + content_seg = Seg(type='text', data='[转发消息]') break - else: - # 处理普通消息 - if raw_message: - from .message import Message - message_obj = Message( - user_id=msg.get('user_id', 0), - message_id=msg.get('message_id', 0), - raw_message=raw_message, - plain_text=raw_message, - group_id=msg.get('group_id', 0) - ) - await message_obj.initialize() - content = message_obj.processed_plain_text else: - content = '[空消息]' + if raw_message: + from .message_cq import MessageRecvCQ + message_obj = MessageRecvCQ( + user_id=msg.get('user_id', 0), + message_id=msg.get('message_id', 0), + raw_message=raw_message, + plain_text=raw_message, + group_id=msg.get('group_id', 0) + ) + content_seg = Seg(type='seglist', data=message_obj.message_segments) + else: + content_seg = Seg(type='text', data='[空消息]') else: - # 处理普通消息 if raw_message: - from .message import Message - message_obj = Message( + from .message_cq import MessageRecvCQ + message_obj = MessageRecvCQ( user_id=msg.get('user_id', 0), message_id=msg.get('message_id', 0), raw_message=raw_message, plain_text=raw_message, group_id=msg.get('group_id', 0) ) - await message_obj.initialize() - content = message_obj.processed_plain_text + content_seg = Seg(type='seglist', data=message_obj.message_segments) else: - content = '[空消息]' + content_seg = Seg(type='text', data='[空消息]') - formatted_msg = f"{nickname}: {content}" - formatted_messages.append(formatted_msg) + formatted_segments.append(Seg(type='text', data=f"{nickname}: ")) + formatted_segments.append(content_seg) + formatted_segments.append(Seg(type='text', data='\n')) - # 合并所有消息 - combined_messages = '\n'.join(formatted_messages) - print(f"\033[1;34m[调试信息]\033[0m 合并后的转发消息: {combined_messages}") - return f"[转发消息:\n{combined_messages}]" + return formatted_segments except Exception as e: - print(f"\033[1;31m[错误]\033[0m 处理转发消息失败: {str(e)}") - return '[转发消息]' + logger.error(f"处理转发消息失败: {str(e)}") + return None - async def translate_reply(self) -> str: - """处理回复类型的CQ码""" - - # 创建Message对象 - from .message import Message - if self.reply_message == None: - # print(f"\033[1;31m[错误]\033[0m 回复消息为空") - return '[回复某人消息]' + def translate_reply(self) -> Optional[List[Seg]]: + """处理回复类型的CQ码，返回Seg列表""" + from .message_cq import MessageRecvCQ + if self.reply_message is None: + return None if self.reply_message.sender.user_id: - message_obj = Message( + message_obj = MessageRecvCQ( user_id=self.reply_message.sender.user_id, message_id=self.reply_message.message_id, raw_message=str(self.reply_message.message), group_id=self.group_id ) - await message_obj.initialize() + + segments = [] if message_obj.user_id == global_config.BOT_QQ: - return f"[回复 {global_config.BOT_NICKNAME} 的消息: {message_obj.processed_plain_text}]" + segments.append(Seg(type='text', data=f"[回复 {global_config.BOT_NICKNAME} 的消息: ")) else: - return f"[回复 {self.reply_message.sender.nickname} 的消息: {message_obj.processed_plain_text}]" - + segments.append(Seg(type='text', data=f"[回复 {self.reply_message.sender.nickname} 的消息: ")) + + segments.append(Seg(type='seglist', data=message_obj.message_segments)) + segments.append(Seg(type='text', data="]")) + return segments else: - print("\033[1;31m[错误]\033[0m 回复消息的sender.user_id为空") - return '[回复某人消息]' + return None @staticmethod def unescape(text: str) -> str: @@ -318,29 +306,12 @@ class CQCode: .replace(']', ']') \ .replace('&', '&') - @staticmethod - def create_emoji_cq(file_path: str) -> str: - """ - 创建表情包CQ码 - Args: - file_path: 本地表情包文件路径 - Returns: - 表情包CQ码字符串 - """ - # 确保使用绝对路径 - abs_path = os.path.abspath(file_path) - # 转义特殊字符 - escaped_path = abs_path.replace('&', '&') \ - .replace('[', '[') \ - .replace(']', ']') \ - .replace(',', ',') - # 生成CQ码，设置sub_type=1表示这是表情包 - return f"[CQ:image,file=file:///{escaped_path},sub_type=1]" + class CQCode_tool: @staticmethod - async def cq_from_dict_to_class(cq_code: Dict, reply: Optional[Dict] = None) -> CQCode: + def cq_from_dict_to_class(cq_code: Dict, reply: Optional[Dict] = None) -> CQCode: """ 将CQ码字典转换为CQCode对象 @@ -369,7 +340,7 @@ class CQCode_tool: ) # 进行翻译处理 - await instance.translate() + instance.translate() return instance @staticmethod @@ -382,6 +353,27 @@ class CQCode_tool: 回复CQ码字符串 """ return f"[CQ:reply,id={message_id}]" + + @staticmethod + def create_emoji_cq(file_path: str) -> str: + """ + 创建表情包CQ码 + Args: + file_path: 本地表情包文件路径 + Returns: + 表情包CQ码字符串 + """ + # 确保使用绝对路径 + abs_path = os.path.abspath(file_path) + # 转义特殊字符 + escaped_path = abs_path.replace('&', '&') \ + .replace('[', '[') \ + .replace(']', ']') \ + .replace(',', ',') + # 生成CQ码，设置sub_type=1表示这是表情包 + return f"[CQ:image,file=file:///{escaped_path},sub_type=1]" + + cq_code_tool = CQCode_tool() diff --git a/src/plugins/chat/emoji_manager.py b/src/plugins/chat/emoji_manager.py index 4f2637738..ff59220dc 100644 --- a/src/plugins/chat/emoji_manager.py +++ b/src/plugins/chat/emoji_manager.py @@ -4,6 +4,8 @@ import random import time import traceback from typing import Optional +import base64 +import hashlib from loguru import logger from nonebot import get_driver @@ -13,9 +15,11 @@ from ..chat.config import global_config from ..chat.utils import get_embedding from ..chat.utils_image import image_path_to_base64 from ..models.utils_model import LLM_request +from ..chat.utils_image import ImageManager driver = get_driver() config = driver.config +image_manager = ImageManager() class EmojiManager: @@ -142,14 +146,14 @@ class EmojiManager: emoji_similarities.sort(key=lambda x: x[1], reverse=True) # 获取前3个最相似的表情包 - top_3_emojis = emoji_similarities[:3] + top_10_emojis = emoji_similarities[:10 if len(emoji_similarities) > 10 else len(emoji_similarities)] - if not top_3_emojis: + if not top_10_emojis: logger.warning("未找到匹配的表情包") return None # 从前3个中随机选择一个 - selected_emoji, similarity = random.choice(top_3_emojis) + selected_emoji, similarity = random.choice(top_10_emojis) if selected_emoji and 'path' in selected_emoji: # 更新使用次数 @@ -172,13 +176,13 @@ class EmojiManager: return None async def _get_emoji_discription(self, image_base64: str) -> str: - """获取表情包的标签""" + """获取表情包的标签，使用image_manager的描述生成功能""" try: - prompt = '这是一个表情包，使用中文简洁的描述一下表情包的内容和表情包所表达的情感' - - content, _ = await self.vlm.generate_response_for_image(prompt, image_base64) - logger.debug(f"输出描述: {content}") - return content + # 使用image_manager获取描述，去掉前后的方括号和"表情包："前缀 + description = await image_manager.get_emoji_description(image_base64) + # 去掉[表情包：xxx]的格式，只保留描述内容 + description = description.strip('[]').replace('表情包：', '') + return description except Exception as e: logger.error(f"获取标签失败: {str(e)}") @@ -220,42 +224,94 @@ class EmojiManager: for filename in files_to_process: image_path = os.path.join(emoji_dir, filename) - # 检查是否已经注册过 - existing_emoji = self.db.db['emoji'].find_one({'filename': filename}) - if existing_emoji: - continue - - # 压缩图片并获取base64编码 + # 获取图片的base64编码和哈希值 image_base64 = image_path_to_base64(image_path) if image_base64 is None: os.remove(image_path) continue - # 获取表情包的描述 - discription = await self._get_emoji_discription(image_base64) + image_bytes = base64.b64decode(image_base64) + image_hash = hashlib.md5(image_bytes).hexdigest() + + # 检查是否已经注册过 + existing_emoji = self.db.db['emoji'].find_one({'filename': filename}) + description = None + + if existing_emoji: + # 即使表情包已存在，也检查是否需要同步到images集合 + description = existing_emoji.get('discription') + # 检查是否在images集合中存在 + existing_image = await image_manager.db.db.images.find_one({'hash': image_hash}) + if not existing_image: + # 同步到images集合 + image_doc = { + 'hash': image_hash, + 'path': image_path, + 'type': 'emoji', + 'description': description, + 'timestamp': int(time.time()) + } + await image_manager.db.db.images.update_one( + {'hash': image_hash}, + {'$set': image_doc}, + upsert=True + ) + # 保存描述到image_descriptions集合 + await image_manager._save_description_to_db(image_hash, description, 'emoji') + logger.success(f"同步已存在的表情包到images集合: {filename}") + continue + + # 检查是否在images集合中已有描述 + existing_description = await image_manager._get_description_from_db(image_hash, 'emoji') + + if existing_description: + description = existing_description + else: + # 获取表情包的描述 + description = await self._get_emoji_discription(image_base64) + if global_config.EMOJI_CHECK: check = await self._check_emoji(image_base64) if '是' not in check: os.remove(image_path) - logger.info(f"描述: {discription}") + logger.info(f"描述: {description}") logger.info(f"其不满足过滤规则，被剔除 {check}") continue logger.info(f"check通过 {check}") - embedding = await get_embedding(discription) - if discription is not None: + + if description is not None: + embedding = await get_embedding(description) # 准备数据库记录 emoji_record = { 'filename': filename, 'path': image_path, - 'embedding':embedding, - 'discription': discription, + 'embedding': embedding, + 'discription': description, + 'hash': image_hash, 'timestamp': int(time.time()) } - # 保存到数据库 + # 保存到emoji数据库 self.db.db['emoji'].insert_one(emoji_record) logger.success(f"注册新表情包: {filename}") - logger.info(f"描述: {discription}") + logger.info(f"描述: {description}") + + # 保存到images数据库 + image_doc = { + 'hash': image_hash, + 'path': image_path, + 'type': 'emoji', + 'description': description, + 'timestamp': int(time.time()) + } + await image_manager.db.db.images.update_one( + {'hash': image_hash}, + {'$set': image_doc}, + upsert=True + ) + # 保存描述到image_descriptions集合 + await image_manager._save_description_to_db(image_hash, description, 'emoji') + logger.success(f"同步保存到images集合: {filename}") else: logger.warning(f"跳过表情包: {filename}") diff --git a/src/plugins/chat/llm_generator.py b/src/plugins/chat/llm_generator.py index 1ac421e6b..2803e5a14 100644 --- a/src/plugins/chat/llm_generator.py +++ b/src/plugins/chat/llm_generator.py @@ -7,7 +7,7 @@ from nonebot import get_driver from ...common.database import Database from ..models.utils_model import LLM_request from .config import global_config -from .message import Message +from .message_cq import Message from .prompt_builder import prompt_builder from .relationship_manager import relationship_manager from .utils import process_llm_response diff --git a/src/plugins/chat/message.py b/src/plugins/chat/message.py index f1fc5569d..a561e6490 100644 --- a/src/plugins/chat/message.py +++ b/src/plugins/chat/message.py @@ -1,231 +1,360 @@ import time from dataclasses import dataclass -from typing import Dict, ForwardRef, List, Optional +from typing import Dict, ForwardRef, List, Optional, Union import urllib3 +from loguru import logger from .cq_code import CQCode, cq_code_tool from .utils_cq import parse_cq_code from .utils_user import get_groupname, get_user_cardname, get_user_nickname - -Message = ForwardRef('Message') # 添加这行 +from .utils_image import image_manager +from .message_base import Seg, GroupInfo, UserInfo, BaseMessageInfo, MessageBase # 禁用SSL警告 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) #这个类是消息数据类，用于存储和管理消息数据。 #它定义了消息的属性，包括群组ID、用户ID、消息ID、原始消息内容、纯文本内容和时间戳。 #它还定义了两个辅助属性：keywords用于提取消息的关键词，is_plain_text用于判断消息是否为纯文本。 + +@dataclass +class MessageRecv(MessageBase): + """接收消息类，用于处理从MessageCQ序列化的消息""" + + def __init__(self, message_dict: Dict): + """从MessageCQ的字典初始化 + + Args: + message_dict: MessageCQ序列化后的字典 + """ + message_info = BaseMessageInfo(**message_dict.get('message_info', {})) + message_segment = Seg(**message_dict.get('message_segment', {})) + raw_message = message_dict.get('raw_message') + + super().__init__( + message_info=message_info, + message_segment=message_segment, + raw_message=raw_message + ) + + # 处理消息内容 + self.processed_plain_text = "" # 初始化为空字符串 + self.detailed_plain_text = "" # 初始化为空字符串 + async def process(self) -> None: + """处理消息内容，生成纯文本和详细文本 + + 这个方法必须在创建实例后显式调用，因为它包含异步操作。 + """ + self.processed_plain_text = await self._process_message_segments(self.message_segment) + self.detailed_plain_text = self._generate_detailed_text() + async def _process_message_segments(self, segment: Seg) -> str: + """递归处理消息段，转换为文字描述 + + Args: + segment: 要处理的消息段 + + Returns: + str: 处理后的文本 + """ + if segment.type == 'seglist': + # 处理消息段列表 + segments_text = [] + for seg in segment.data: + processed = await self._process_message_segments(seg) + if processed: + segments_text.append(processed) + return ' '.join(segments_text) + else: + # 处理单个消息段 + return await self._process_single_segment(segment) + + async def _process_single_segment(self, seg: Seg) -> str: + """处理单个消息段 + + Args: + seg: 要处理的消息段 + + Returns: + str: 处理后的文本 + """ + try: + if seg.type == 'text': + return seg.data + elif seg.type == 'image': + # 如果是base64图片数据 + if isinstance(seg.data, str) and seg.data.startswith(('data:', 'base64:')): + return await image_manager.get_image_description(seg.data) + return '[图片]' + elif seg.type == 'emoji': + if isinstance(seg.data, str) and seg.data.startswith(('data:', 'base64:')): + return await image_manager.get_emoji_description(seg.data) + return '[表情]' + else: + return f"[{seg.type}:{str(seg.data)}]" + except Exception as e: + logger.error(f"处理消息段失败: {str(e)}, 类型: {seg.type}, 数据: {seg.data}") + return f"[处理失败的{seg.type}消息]" + + def _generate_detailed_text(self) -> str: + """生成详细文本，包含时间和用户信息""" + time_str = time.strftime("%m-%d %H:%M:%S", time.localtime(self.message_info.time)) + user_info = self.message_info.user_info + name = ( + f"{user_info.user_nickname}(ta的昵称:{user_info.user_cardname},ta的id:{user_info.user_id})" + if user_info.user_cardname!='' + else f"{user_info.user_nickname}(ta的id:{user_info.user_id})" + ) + return f"[{time_str}] {name}: {self.processed_plain_text}\n" @dataclass -class Message: - """消息数据类""" - message_id: int = None - time: float = None - - group_id: int = None - group_name: str = None # 群名称 - - user_id: int = None - user_nickname: str = None # 用户昵称 - user_cardname: str = None # 用户群昵称 - - raw_message: str = None # 原始消息，包含未解析的cq码 - plain_text: str = None # 纯文本 - - reply_message: Dict = None # 存储 回复的 源消息 - - # 延迟初始化字段 - _initialized: bool = False - message_segments: List[Dict] = None # 存储解析后的消息片段 - processed_plain_text: str = None # 用于存储处理后的plain_text - detailed_plain_text: str = None # 用于存储详细可读文本 - - # 状态标志 - is_emoji: bool = False - has_emoji: bool = False - translate_cq: bool = True - - async def initialize(self): - """显式异步初始化方法（必须调用）""" - if self._initialized: - return - - # 异步获取补充信息 - self.group_name = self.group_name or get_groupname(self.group_id) - self.user_nickname = self.user_nickname or get_user_nickname(self.user_id) - self.user_cardname = self.user_cardname or get_user_cardname(self.user_id) - - # 消息解析 - if self.raw_message: - if not isinstance(self,Message_Sending): - self.message_segments = await self.parse_message_segments(self.raw_message) - self.processed_plain_text = ' '.join( - seg.translated_plain_text - for seg in self.message_segments - ) - - # 构建详细文本 - if self.time is None: - self.time = int(time.time()) - time_str = time.strftime("%m-%d %H:%M:%S", time.localtime(self.time)) - name = ( - f"{self.user_nickname}(ta的昵称:{self.user_cardname},ta的id:{self.user_id})" - if self.user_cardname - else f"{self.user_nickname or f'用户{self.user_id}'}" - ) - if isinstance(self,Message_Sending) and self.is_emoji: - self.detailed_plain_text = f"[{time_str}] {name}: {self.detailed_plain_text}\n" - else: - self.detailed_plain_text = f"[{time_str}] {name}: {self.processed_plain_text}\n" - - self._initialized = True +class MessageProcessBase(MessageBase): + """消息处理基类，用于处理中和发送中的消息""" - async def parse_message_segments(self, message: str) -> List[CQCode]: - """ - 将消息解析为片段列表，包括纯文本和CQ码 - 返回的列表中每个元素都是字典，包含： - - cq_code_list:分割出的聊天对象，包括文本和CQ码 - - trans_list:翻译后的对象列表 - """ - # print(f"\033[1;34m[调试信息]\033[0m 正在处理消息: {message}") - cq_code_dict_list = [] - trans_list = [] - - start = 0 - while True: - # 查找下一个CQ码的开始位置 - cq_start = message.find('[CQ:', start) - #如果没有cq码，直接返回文本内容 - if cq_start == -1: - # 如果没有找到更多CQ码，添加剩余文本 - if start < len(message): - text = message[start:].strip() - if text: # 只添加非空文本 - cq_code_dict_list.append(parse_cq_code(text)) - break - # 添加CQ码前的文本 - if cq_start > start: - text = message[start:cq_start].strip() - if text: # 只添加非空文本 - cq_code_dict_list.append(parse_cq_code(text)) - # 查找CQ码的结束位置 - cq_end = message.find(']', cq_start) - if cq_end == -1: - # CQ码未闭合，作为普通文本处理 - text = message[cq_start:].strip() - if text: - cq_code_dict_list.append(parse_cq_code(text)) - break - cq_code = message[cq_start:cq_end + 1] - - #将cq_code解析成字典 - cq_code_dict_list.append(parse_cq_code(cq_code)) - # 更新start位置到当前CQ码之后 - start = cq_end + 1 - - # print(f"\033[1;34m[调试信息]\033[0m 提取的消息对象：列表: {cq_code_dict_list}") - - #判定是否是表情包消息，以及是否含有表情包 - if len(cq_code_dict_list) == 1 and cq_code_dict_list[0]['type'] == 'image': - self.is_emoji = True - self.has_emoji_emoji = True - else: - for segment in cq_code_dict_list: - if segment['type'] == 'image' and segment['data'].get('sub_type') == '1': - self.has_emoji_emoji = True - break - - - #翻译作为字典的CQ码 - for _code_item in cq_code_dict_list: - message_obj = await cq_code_tool.cq_from_dict_to_class(_code_item,reply = self.reply_message) - trans_list.append(message_obj) - return trans_list + def __init__( + self, + message_id: str, + user_id: int, + group_id: Optional[int] = None, + platform: str = "qq", + message_segment: Optional[Seg] = None, + reply: Optional['MessageRecv'] = None + ): + # 构造用户信息 + user_info = UserInfo( + platform=platform, + user_id=user_id, + user_nickname=get_user_nickname(user_id), + user_cardname=get_user_cardname(user_id) if group_id else None + ) -class Message_Thinking: - """消息思考类""" - def __init__(self, message: Message,message_id: str): - # 复制原始消息的基本属性 - self.group_id = message.group_id - self.user_id = message.user_id - self.user_nickname = message.user_nickname - self.user_cardname = message.user_cardname - self.group_name = message.group_name - - self.message_id = message_id - - # 思考状态相关属性 + # 构造群组信息（如果有） + group_info = None + if group_id: + group_info = GroupInfo( + platform=platform, + group_id=group_id, + group_name=get_groupname(group_id) + ) + + # 构造基础消息信息 + message_info = BaseMessageInfo( + platform=platform, + message_id=message_id, + time=int(time.time()), + group_info=group_info, + user_info=user_info + ) + + # 调用父类初始化 + super().__init__( + message_info=message_info, + message_segment=message_segment, + raw_message=None + ) + + # 处理状态相关属性 self.thinking_start_time = int(time.time()) self.thinking_time = 0 - self.interupt=False - - def update_thinking_time(self): - self.thinking_time = round(time.time(), 2) - self.thinking_start_time - + + # 文本处理相关属性 + self.processed_plain_text = "" + self.detailed_plain_text = "" + + # 回复消息 + self.reply = reply -@dataclass -class Message_Sending(Message): - """发送中的消息类""" - thinking_start_time: float = None # 思考开始时间 - thinking_time: float = None # 思考时间 - - reply_message_id: int = None # 存储 回复的 源消息ID - - is_head: bool = False # 是否是头部消息 - - def update_thinking_time(self): - self.thinking_time = round(time.time(), 2) - self.thinking_start_time + def update_thinking_time(self) -> float: + """更新思考时间""" + self.thinking_time = round(time.time() - self.thinking_start_time, 2) return self.thinking_time + async def _process_message_segments(self, segment: Seg) -> str: + """递归处理消息段，转换为文字描述 + + Args: + segment: 要处理的消息段 + + Returns: + str: 处理后的文本 + """ + if segment.type == 'seglist': + # 处理消息段列表 + segments_text = [] + for seg in segment.data: + processed = await self._process_message_segments(seg) + if processed: + segments_text.append(processed) + return ' '.join(segments_text) + else: + # 处理单个消息段 + return await self._process_single_segment(segment) - + async def _process_single_segment(self, seg: Seg) -> str: + """处理单个消息段 + + Args: + seg: 要处理的消息段 + + Returns: + str: 处理后的文本 + """ + try: + if seg.type == 'text': + return seg.data + elif seg.type == 'image': + # 如果是base64图片数据 + if isinstance(seg.data, str) and seg.data.startswith(('data:', 'base64:')): + return await image_manager.get_image_description(seg.data) + return '[图片]' + elif seg.type == 'emoji': + if isinstance(seg.data, str) and seg.data.startswith(('data:', 'base64:')): + return await image_manager.get_emoji_description(seg.data) + return '[表情]' + elif seg.type == 'at': + return f"[@{seg.data}]" + elif seg.type == 'reply': + if self.reply and hasattr(self.reply, 'processed_plain_text'): + return f"[回复：{self.reply.processed_plain_text}]" + else: + return f"[{seg.type}:{str(seg.data)}]" + except Exception as e: + logger.error(f"处理消息段失败: {str(e)}, 类型: {seg.type}, 数据: {seg.data}") + return f"[处理失败的{seg.type}消息]" + + def _generate_detailed_text(self) -> str: + """生成详细文本，包含时间和用户信息""" + time_str = time.strftime("%m-%d %H:%M:%S", time.localtime(self.message_info.time)) + user_info = self.message_info.user_info + name = ( + f"{user_info.user_nickname}(ta的昵称:{user_info.user_cardname},ta的id:{user_info.user_id})" + if user_info.user_cardname != '' + else f"{user_info.user_nickname}(ta的id:{user_info.user_id})" + ) + return f"[{time_str}] {name}: {self.processed_plain_text}\n" + +@dataclass +class MessageThinking(MessageProcessBase): + """思考状态的消息类""" + + def __init__( + self, + message_id: str, + user_id: int, + group_id: Optional[int] = None, + platform: str = "qq", + reply: Optional['MessageRecv'] = None + ): + # 调用父类初始化 + super().__init__( + message_id=message_id, + user_id=user_id, + group_id=group_id, + platform=platform, + message_segment=None, # 思考状态不需要消息段 + reply=reply + ) + + # 思考状态特有属性 + self.interrupt = False + +@dataclass +class MessageSending(MessageProcessBase): + """发送状态的消息类""" + + def __init__( + self, + message_id: str, + user_id: int, + message_segment: Seg, + group_id: Optional[int] = None, + reply: Optional['MessageRecv'] = None, + platform: str = "qq", + is_head: bool = False + ): + # 调用父类初始化 + super().__init__( + message_id=message_id, + user_id=user_id, + group_id=group_id, + platform=platform, + message_segment=message_segment, + reply=reply + ) + + # 发送状态特有属性 + self.reply_to_message_id = reply.message_info.message_id if reply else None + self.is_head = is_head + + async def process(self) -> None: + """处理消息内容，生成纯文本和详细文本""" + if self.message_segment: + self.processed_plain_text = await self._process_message_segments(self.message_segment) + self.detailed_plain_text = self._generate_detailed_text() + + @classmethod + def from_thinking( + cls, + thinking: MessageThinking, + message_segment: Seg, + reply: Optional['MessageRecv'] = None, + is_head: bool = False + ) -> 'MessageSending': + """从思考状态消息创建发送状态消息""" + return cls( + message_id=thinking.message_info.message_id, + user_id=thinking.message_info.user_info.user_id, + message_segment=message_segment, + group_id=thinking.message_info.group_info.group_id if thinking.message_info.group_info else None, + reply=reply or thinking.reply, + platform=thinking.message_info.platform, + is_head=is_head + ) + +@dataclass class MessageSet: """消息集合类，可以存储多个发送消息""" def __init__(self, group_id: int, user_id: int, message_id: str): self.group_id = group_id self.user_id = user_id self.message_id = message_id - self.messages: List[Message_Sending] = [] # 修改类型标注 + self.messages: List[MessageSending] = [] self.time = round(time.time(), 2) - def add_message(self, message: Message_Sending) -> None: - """添加消息到集合，只接受Message_Sending类型""" - if not isinstance(message, Message_Sending): - raise TypeError("MessageSet只能添加Message_Sending类型的消息") + def add_message(self, message: MessageSending) -> None: + """添加消息到集合""" + if not isinstance(message, MessageSending): + raise TypeError("MessageSet只能添加MessageSending类型的消息") self.messages.append(message) - # 按时间排序 - self.messages.sort(key=lambda x: x.time) + self.messages.sort(key=lambda x: x.message_info.time) - def get_message_by_index(self, index: int) -> Optional[Message_Sending]: + def get_message_by_index(self, index: int) -> Optional[MessageSending]: """通过索引获取消息""" if 0 <= index < len(self.messages): return self.messages[index] return None - def get_message_by_time(self, target_time: float) -> Optional[Message_Sending]: + def get_message_by_time(self, target_time: float) -> Optional[MessageSending]: """获取最接近指定时间的消息""" if not self.messages: return None - # 使用二分查找找到最接近的消息 left, right = 0, len(self.messages) - 1 while left < right: mid = (left + right) // 2 - if self.messages[mid].time < target_time: + if self.messages[mid].message_info.time < target_time: left = mid + 1 else: right = mid return self.messages[left] - def clear_messages(self) -> None: """清空所有消息""" self.messages.clear() - def remove_message(self, message: Message_Sending) -> bool: + def remove_message(self, message: MessageSending) -> bool: """移除指定消息""" if message in self.messages: self.messages.remove(message) diff --git a/src/plugins/chat/message_base.py b/src/plugins/chat/message_base.py new file mode 100644 index 000000000..77694ad8c --- /dev/null +++ b/src/plugins/chat/message_base.py @@ -0,0 +1,158 @@ +from dataclasses import dataclass, asdict +from typing import List, Optional, Union, Any, Dict + +@dataclass +class Seg(dict): + """消息片段类，用于表示消息的不同部分 + + Attributes: + type: 片段类型，可以是 'text'、'image'、'seglist' 等 + data: 片段的具体内容 + - 对于 text 类型，data 是字符串 + - 对于 image 类型，data 是 base64 字符串 + - 对于 seglist 类型，data 是 Seg 列表 + translated_data: 经过翻译处理的数据（可选） + """ + type: str + data: Union[str, List['Seg']] + translated_data: Optional[str] = None + + def __init__(self, type: str, data: Union[str, List['Seg']], translated_data: Optional[str] = None): + """初始化实例，确保字典和属性同步""" + # 先初始化字典 + super().__init__(type=type, data=data) + if translated_data is not None: + self['translated_data'] = translated_data + + # 再初始化属性 + object.__setattr__(self, 'type', type) + object.__setattr__(self, 'data', data) + object.__setattr__(self, 'translated_data', translated_data) + + # 验证数据类型 + self._validate_data() + + def _validate_data(self) -> None: + """验证数据类型的正确性""" + if self.type == 'seglist' and not isinstance(self.data, list): + raise ValueError("seglist类型的data必须是列表") + elif self.type == 'text' and not isinstance(self.data, str): + raise ValueError("text类型的data必须是字符串") + elif self.type == 'image' and not isinstance(self.data, str): + raise ValueError("image类型的data必须是字符串") + + def __setattr__(self, name: str, value: Any) -> None: + """重写属性设置，同时更新字典值""" + # 更新属性 + object.__setattr__(self, name, value) + # 同步更新字典 + if name in ['type', 'data', 'translated_data']: + self[name] = value + + def __setitem__(self, key: str, value: Any) -> None: + """重写字典值设置，同时更新属性""" + # 更新字典 + super().__setitem__(key, value) + # 同步更新属性 + if key in ['type', 'data', 'translated_data']: + object.__setattr__(self, key, value) + + def to_dict(self) -> Dict: + """转换为字典格式""" + result = {'type': self.type} + if self.type == 'seglist': + result['data'] = [seg.to_dict() for seg in self.data] + else: + result['data'] = self.data + if self.translated_data is not None: + result['translated_data'] = self.translated_data + return result + +@dataclass +class GroupInfo: + """群组信息类""" + platform: Optional[str] = None + group_id: Optional[int] = None + group_name: Optional[str] = None # 群名称 + + def to_dict(self) -> Dict: + """转换为字典格式""" + return {k: v for k, v in asdict(self).items() if v is not None} + +@dataclass +class UserInfo: + """用户信息类""" + platform: Optional[str] = None + user_id: Optional[int] = None + user_nickname: Optional[str] = None # 用户昵称 + user_cardname: Optional[str] = None # 用户群昵称 + + def to_dict(self) -> Dict: + """转换为字典格式""" + return {k: v for k, v in asdict(self).items() if v is not None} + +@dataclass +class BaseMessageInfo: + """消息信息类""" + platform: Optional[str] = None + message_id: Optional[int,str] = None + time: Optional[int] = None + group_info: Optional[GroupInfo] = None + user_info: Optional[UserInfo] = None + + def to_dict(self) -> Dict: + """转换为字典格式""" + result = {} + for field, value in asdict(self).items(): + if value is not None: + if isinstance(value, (GroupInfo, UserInfo)): + result[field] = value.to_dict() + else: + result[field] = value + return result + +@dataclass +class MessageBase: + """消息类""" + message_info: BaseMessageInfo + message_segment: Seg + raw_message: Optional[str] = None # 原始消息，包含未解析的cq码 + + def to_dict(self) -> Dict: + """转换为字典格式 + + Returns: + Dict: 包含所有非None字段的字典，其中： + - message_info: 转换为字典格式 + - message_segment: 转换为字典格式 + - raw_message: 如果存在则包含 + """ + result = { + 'message_info': self.message_info.to_dict(), + 'message_segment': self.message_segment.to_dict() + } + if self.raw_message is not None: + result['raw_message'] = self.raw_message + return result + + @classmethod + def from_dict(cls, data: Dict) -> 'MessageBase': + """从字典创建MessageBase实例 + + Args: + data: 包含必要字段的字典 + + Returns: + MessageBase: 新的实例 + """ + message_info = BaseMessageInfo(**data.get('message_info', {})) + message_segment = Seg(**data.get('message_segment', {})) + raw_message = data.get('raw_message') + return cls( + message_info=message_info, + message_segment=message_segment, + raw_message=raw_message + ) + + + diff --git a/src/plugins/chat/message_cq.py b/src/plugins/chat/message_cq.py new file mode 100644 index 000000000..80df4e340 --- /dev/null +++ b/src/plugins/chat/message_cq.py @@ -0,0 +1,188 @@ +import time +from dataclasses import dataclass +from typing import Dict, ForwardRef, List, Optional, Union + +import urllib3 + +from .cq_code import CQCode, cq_code_tool +from .utils_cq import parse_cq_code +from .utils_user import get_groupname, get_user_cardname, get_user_nickname +from .message_base import Seg, GroupInfo, UserInfo, BaseMessageInfo, MessageBase +# 禁用SSL警告 +urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) + +#这个类是消息数据类，用于存储和管理消息数据。 +#它定义了消息的属性，包括群组ID、用户ID、消息ID、原始消息内容、纯文本内容和时间戳。 +#它还定义了两个辅助属性：keywords用于提取消息的关键词，is_plain_text用于判断消息是否为纯文本。 + +@dataclass +class MessageCQ(MessageBase): + """QQ消息基类，继承自MessageBase + + 最小必要参数: + - message_id: 消息ID + - user_id: 发送者/接收者ID + - platform: 平台标识（默认为"qq"） + """ + def __init__( + self, + message_id: int, + user_id: int, + group_id: Optional[int] = None, + platform: str = "qq" + ): + # 构造用户信息 + user_info = UserInfo( + platform=platform, + user_id=user_id, + user_nickname=get_user_nickname(user_id), + user_cardname=get_user_cardname(user_id) if group_id else None + ) + + # 构造群组信息（如果有） + group_info = None + if group_id: + group_info = GroupInfo( + platform=platform, + group_id=group_id, + group_name=get_groupname(group_id) + ) + + # 构造基础消息信息 + message_info = BaseMessageInfo( + platform=platform, + message_id=message_id, + time=int(time.time()), + group_info=group_info, + user_info=user_info + ) + + # 调用父类初始化，message_segment 由子类设置 + super().__init__( + message_info=message_info, + message_segment=None, + raw_message=None + ) + +@dataclass +class MessageRecvCQ(MessageCQ): + """QQ接收消息类，用于解析raw_message到Seg对象""" + + def __init__( + self, + message_id: int, + user_id: int, + raw_message: str, + group_id: Optional[int] = None, + reply_message: Optional[Dict] = None, + platform: str = "qq" + ): + # 调用父类初始化 + super().__init__(message_id, user_id, group_id, platform) + + # 解析消息段 + self.message_segment = self._parse_message(raw_message, reply_message) + self.raw_message = raw_message + + def _parse_message(self, message: str, reply_message: Optional[Dict] = None) -> Seg: + """解析消息内容为Seg对象""" + cq_code_dict_list = [] + segments = [] + + start = 0 + while True: + cq_start = message.find('[CQ:', start) + if cq_start == -1: + if start < len(message): + text = message[start:].strip() + if text: + cq_code_dict_list.append(parse_cq_code(text)) + break + + if cq_start > start: + text = message[start:cq_start].strip() + if text: + cq_code_dict_list.append(parse_cq_code(text)) + + cq_end = message.find(']', cq_start) + if cq_end == -1: + text = message[cq_start:].strip() + if text: + cq_code_dict_list.append(parse_cq_code(text)) + break + + cq_code = message[cq_start:cq_end + 1] + cq_code_dict_list.append(parse_cq_code(cq_code)) + start = cq_end + 1 + + # 转换CQ码为Seg对象 + for code_item in cq_code_dict_list: + message_obj = cq_code_tool.cq_from_dict_to_class(code_item, reply=reply_message) + if message_obj.translated_segments: + segments.append(message_obj.translated_segments) + + # 如果只有一个segment，直接返回 + if len(segments) == 1: + return segments[0] + + # 否则返回seglist类型的Seg + return Seg(type='seglist', data=segments) + + def to_dict(self) -> Dict: + """转换为字典格式，包含所有必要信息""" + base_dict = super().to_dict() + return base_dict + +@dataclass +class MessageSendCQ(MessageCQ): + """QQ发送消息类，用于将Seg对象转换为raw_message""" + + def __init__( + self, + message_id: int, + user_id: int, + message_segment: Seg, + group_id: Optional[int] = None, + reply_to_message_id: Optional[int] = None, + platform: str = "qq" + ): + # 调用父类初始化 + super().__init__(message_id, user_id, group_id, platform) + + self.message_segment = message_segment + self.raw_message = self._generate_raw_message(reply_to_message_id) + + def _generate_raw_message(self, reply_to_message_id: Optional[int] = None) -> str: + """将Seg对象转换为raw_message""" + segments = [] + + # 添加回复消息 + if reply_to_message_id: + segments.append(cq_code_tool.create_reply_cq(reply_to_message_id)) + + # 处理消息段 + if self.message_segment.type == 'seglist': + for seg in self.message_segment.data: + segments.append(self._seg_to_cq_code(seg)) + else: + segments.append(self._seg_to_cq_code(self.message_segment)) + + return ''.join(segments) + + def _seg_to_cq_code(self, seg: Seg) -> str: + """将单个Seg对象转换为CQ码字符串""" + if seg.type == 'text': + return str(seg.data) + elif seg.type == 'image': + # 如果是base64图片数据 + if seg.data.startswith(('data:', 'base64:')): + return f"[CQ:image,file=base64://{seg.data}]" + # 如果是表情包（本地文件） + return cq_code_tool.create_emoji_cq(seg.data) + elif seg.type == 'at': + return f"[CQ:at,qq={seg.data}]" + elif seg.type == 'reply': + return cq_code_tool.create_reply_cq(int(seg.data)) + else: + return f"[{seg.data}]" + diff --git a/src/plugins/chat/message_sender.py b/src/plugins/chat/message_sender.py index 050c59d74..8ed30d69c 100644 --- a/src/plugins/chat/message_sender.py +++ b/src/plugins/chat/message_sender.py @@ -5,7 +5,7 @@ from typing import Dict, List, Optional, Union from nonebot.adapters.onebot.v11 import Bot from .cq_code import cq_code_tool -from .message import Message, Message_Sending, Message_Thinking, MessageSet +from .message_cq import Message, Message_Sending, Message_Thinking, MessageSet from .storage import MessageStorage from .utils import calculate_typing_time from .config import global_config diff --git a/src/plugins/chat/storage.py b/src/plugins/chat/storage.py index 6a87480b7..4b1bf0ca4 100644 --- a/src/plugins/chat/storage.py +++ b/src/plugins/chat/storage.py @@ -1,7 +1,7 @@ from typing import Optional from ...common.database import Database -from .message import Message +from .message_cq import Message class MessageStorage: diff --git a/src/plugins/chat/utils.py b/src/plugins/chat/utils.py index b2583e86f..07c3e2a6f 100644 --- a/src/plugins/chat/utils.py +++ b/src/plugins/chat/utils.py @@ -11,31 +11,12 @@ from nonebot import get_driver from ..models.utils_model import LLM_request from ..utils.typo_generator import ChineseTypoGenerator from .config import global_config -from .message import Message +from .message_cq import Message driver = get_driver() config = driver.config -def combine_messages(messages: List[Message]) -> str: - """将消息列表组合成格式化的字符串 - - Args: - messages: Message对象列表 - - Returns: - str: 格式化后的消息字符串 - """ - result = "" - for message in messages: - time_str = time.strftime("%m-%d %H:%M:%S", time.localtime(message.time)) - name = message.user_nickname or f"用户{message.user_id}" - content = message.processed_plain_text or message.plain_text - - result += f"[{time_str}] {name}: {content}\n" - - return result - def db_message_to_str(message_dict: Dict) -> str: print(f"message_dict: {message_dict}") @@ -159,7 +140,7 @@ async def get_recent_group_messages(db, group_id: int, limit: int = 12) -> list: return [] # 转换为 Message对象列表 - from .message import Message + from .message_cq import Message message_objects = [] for msg_data in recent_messages: try: diff --git a/src/plugins/chat/utils_image.py b/src/plugins/chat/utils_image.py index 8a8b3ce5a..aba09714c 100644 --- a/src/plugins/chat/utils_image.py +++ b/src/plugins/chat/utils_image.py @@ -2,7 +2,11 @@ import base64 import io import os import time -import zlib # 用于 CRC32 +import zlib +import aiohttp +import hashlib +from typing import Optional, Tuple, Union +from urllib.parse import urlparse from loguru import logger from nonebot import get_driver @@ -10,213 +14,348 @@ from PIL import Image from ...common.database import Database from ..chat.config import global_config - +from ..models.utils_model import LLM_request driver = get_driver() config = driver.config - - -def storage_compress_image(base64_data: str, max_size: int = 200) -> str: - """ - 压缩base64格式的图片到指定大小（单位：KB）并在数据库中记录图片信息 - Args: - base64_data: base64编码的图片数据 - max_size: 最大文件大小（KB） - Returns: - str: 压缩后的base64图片数据 - """ - try: - # 将base64转换为字节数据 - image_data = base64.b64decode(base64_data) - - # 使用 CRC32 计算哈希值 - hash_value = format(zlib.crc32(image_data) & 0xFFFFFFFF, 'x') - - # 确保图片目录存在 - images_dir = "data/images" - os.makedirs(images_dir, exist_ok=True) - - # 连接数据库 - db = Database( - host=config.mongodb_host, - port=int(config.mongodb_port), - db_name=config.database_name, - username=config.mongodb_username, - password=config.mongodb_password, - auth_source=config.mongodb_auth_source - ) - - # 检查是否已存在相同哈希值的图片 - collection = db.db['images'] - existing_image = collection.find_one({'hash': hash_value}) - - if existing_image: - print(f"\033[1;33m[提示]\033[0m 发现重复图片，使用已存在的文件: {existing_image['path']}") - return base64_data - - # 将字节数据转换为图片对象 - img = Image.open(io.BytesIO(image_data)) - - # 如果是动图，直接返回原图 - if getattr(img, 'is_animated', False): - return base64_data - - # 计算当前大小（KB） - current_size = len(image_data) / 1024 - - # 如果已经小于目标大小，直接使用原图 - if current_size <= max_size: - compressed_data = image_data - else: - # 压缩逻辑 - # 先缩放到50% - new_width = int(img.width * 0.5) - new_height = int(img.height * 0.5) - img = img.resize((new_width, new_height), Image.Resampling.LANCZOS) - - # 如果缩放后的最大边长仍然大于400，继续缩放 - max_dimension = 400 - max_current = max(new_width, new_height) - if max_current > max_dimension: - ratio = max_dimension / max_current - new_width = int(new_width * ratio) - new_height = int(new_height * ratio) - img = img.resize((new_width, new_height), Image.Resampling.LANCZOS) - - # 转换为RGB模式（去除透明通道） - if img.mode in ('RGBA', 'P'): - img = img.convert('RGB') - - # 使用固定质量参数压缩 - output = io.BytesIO() - img.save(output, format='JPEG', quality=85, optimize=True) - compressed_data = output.getvalue() - - # 生成文件名（使用时间戳和哈希值确保唯一性） - timestamp = int(time.time()) - filename = f"{timestamp}_{hash_value}.jpg" - image_path = os.path.join(images_dir, filename) - - # 保存文件 - with open(image_path, "wb") as f: - f.write(compressed_data) - - print(f"\033[1;32m[成功]\033[0m 保存图片到: {image_path}") - - try: - # 准备数据库记录 - image_record = { - 'filename': filename, - 'path': image_path, - 'size': len(compressed_data) / 1024, - 'timestamp': timestamp, - 'width': img.width, - 'height': img.height, - 'description': '', - 'tags': [], - 'type': 'image', - 'hash': hash_value - } - - # 保存记录 - collection.insert_one(image_record) - print("\033[1;32m[成功]\033[0m 保存图片记录到数据库") - - except Exception as db_error: - print(f"\033[1;31m[错误]\033[0m 数据库操作失败: {str(db_error)}") - - # 将压缩后的数据转换为base64 - compressed_base64 = base64.b64encode(compressed_data).decode('utf-8') - return compressed_base64 - - except Exception as e: - print(f"\033[1;31m[错误]\033[0m 压缩图片失败: {str(e)}") - import traceback - print(traceback.format_exc()) - return base64_data - -def storage_emoji(image_data: bytes) -> bytes: - """ - 存储表情包到本地文件夹 - Args: - image_data: 图片字节数据 - group_id: 群组ID（仅用于日志） - user_id: 用户ID（仅用于日志） - Returns: - bytes: 原始图片数据 - """ - if not global_config.EMOJI_SAVE: - return image_data - try: - # 使用 CRC32 计算哈希值 - hash_value = format(zlib.crc32(image_data) & 0xFFFFFFFF, 'x') - - # 确保表情包目录存在 - emoji_dir = "data/emoji" - os.makedirs(emoji_dir, exist_ok=True) - - # 检查是否已存在相同哈希值的文件 - for filename in os.listdir(emoji_dir): - if hash_value in filename: - # print(f"\033[1;33m[提示]\033[0m 发现重复表情包: {filename}") - return image_data - - # 生成文件名 - timestamp = int(time.time()) - filename = f"{timestamp}_{hash_value}.jpg" - emoji_path = os.path.join(emoji_dir, filename) - - # 直接保存原始文件 - with open(emoji_path, "wb") as f: - f.write(image_data) - - print(f"\033[1;32m[成功]\033[0m 保存表情包到: {emoji_path}") - return image_data - - except Exception as e: - print(f"\033[1;31m[错误]\033[0m 保存表情包失败: {str(e)}") - return image_data +class ImageManager: + _instance = None + IMAGE_DIR = "data" # 图像存储根目录 - -def storage_image(image_data: bytes) -> bytes: - """ - 存储图片到本地文件夹 - Args: - image_data: 图片字节数据 - group_id: 群组ID（仅用于日志） - user_id: 用户ID（仅用于日志） - Returns: - bytes: 原始图片数据 - """ - try: - # 使用 CRC32 计算哈希值 - hash_value = format(zlib.crc32(image_data) & 0xFFFFFFFF, 'x') - - # 确保表情包目录存在 - image_dir = "data/image" - os.makedirs(image_dir, exist_ok=True) - - # 检查是否已存在相同哈希值的文件 - for filename in os.listdir(image_dir): - if hash_value in filename: - # print(f"\033[1;33m[提示]\033[0m 发现重复表情包: {filename}") - return image_data - - # 生成文件名 - timestamp = int(time.time()) - filename = f"{timestamp}_{hash_value}.jpg" - image_path = os.path.join(image_dir, filename) - - # 直接保存原始文件 - with open(image_path, "wb") as f: - f.write(image_data) + def __new__(cls): + if cls._instance is None: + cls._instance = super().__new__(cls) + cls._instance.db = None + cls._instance._initialized = False + return cls._instance + + def __init__(self): + if not self._initialized: + self.db = Database.get_instance() + self._ensure_image_collection() + self._ensure_description_collection() + self._ensure_image_dir() + self._initialized = True + self._llm = LLM_request(model=global_config.vlm, temperature=0.4, max_tokens=300) - print(f"\033[1;32m[成功]\033[0m 保存图片到: {image_path}") - return image_data + def _ensure_image_dir(self): + """确保图像存储目录存在""" + os.makedirs(self.IMAGE_DIR, exist_ok=True) - except Exception as e: - print(f"\033[1;31m[错误]\033[0m 保存图片失败: {str(e)}") - return image_data + def _ensure_image_collection(self): + """确保images集合存在并创建索引""" + if 'images' not in self.db.db.list_collection_names(): + self.db.db.create_collection('images') + # 创建索引 + self.db.db.images.create_index([('hash', 1)], unique=True) + self.db.db.images.create_index([('url', 1)]) + self.db.db.images.create_index([('path', 1)]) + + def _ensure_description_collection(self): + """确保image_descriptions集合存在并创建索引""" + if 'image_descriptions' not in self.db.db.list_collection_names(): + self.db.db.create_collection('image_descriptions') + # 创建索引 + self.db.db.image_descriptions.create_index([('hash', 1)], unique=True) + self.db.db.image_descriptions.create_index([('type', 1)]) + + async def _get_description_from_db(self, image_hash: str, description_type: str) -> Optional[str]: + """从数据库获取图片描述 + + Args: + image_hash: 图片哈希值 + description_type: 描述类型 ('emoji' 或 'image') + + Returns: + Optional[str]: 描述文本，如果不存在则返回None + """ + result = await self.db.db.image_descriptions.find_one({ + 'hash': image_hash, + 'type': description_type + }) + return result['description'] if result else None + + async def _save_description_to_db(self, image_hash: str, description: str, description_type: str) -> None: + """保存图片描述到数据库 + + Args: + image_hash: 图片哈希值 + description: 描述文本 + description_type: 描述类型 ('emoji' 或 'image') + """ + await self.db.db.image_descriptions.update_one( + {'hash': image_hash, 'type': description_type}, + { + '$set': { + 'description': description, + 'timestamp': int(time.time()) + } + }, + upsert=True + ) + + async def save_image(self, + image_data: Union[str, bytes], + url: str = None, + description: str = None, + is_base64: bool = False) -> Optional[str]: + """保存图像 + Args: + image_data: 图像数据(base64字符串或字节) + url: 图像URL + description: 图像描述 + is_base64: image_data是否为base64格式 + Returns: + str: 保存后的文件路径,失败返回None + """ + try: + # 转换为字节格式 + if is_base64: + if isinstance(image_data, str): + image_bytes = base64.b64decode(image_data) + else: + return None + else: + if isinstance(image_data, bytes): + image_bytes = image_data + else: + return None + + # 计算哈希值 + image_hash = hashlib.md5(image_bytes).hexdigest() + + # 查重 + existing = self.db.db.images.find_one({'hash': image_hash}) + if existing: + return existing['path'] + + # 生成文件名和路径 + timestamp = int(time.time()) + filename = f"{timestamp}_{image_hash[:8]}.jpg" + file_path = os.path.join(self.IMAGE_DIR, filename) + + # 保存文件 + with open(file_path, "wb") as f: + f.write(image_bytes) + + # 保存到数据库 + image_doc = { + 'hash': image_hash, + 'path': file_path, + 'url': url, + 'description': description, + 'timestamp': timestamp + } + self.db.db.images.insert_one(image_doc) + + return file_path + + except Exception as e: + logger.error(f"保存图像失败: {str(e)}") + return None + + async def get_image_by_url(self, url: str) -> Optional[str]: + """根据URL获取图像路径(带查重) + Args: + url: 图像URL + Returns: + str: 本地文件路径,不存在返回None + """ + try: + # 先查找是否已存在 + existing = self.db.db.images.find_one({'url': url}) + if existing: + return existing['path'] + + # 下载图像 + async with aiohttp.ClientSession() as session: + async with session.get(url) as resp: + if resp.status == 200: + image_bytes = await resp.read() + return await self.save_image(image_bytes, url=url) + return None + + except Exception as e: + logger.error(f"获取图像失败: {str(e)}") + return None + + async def get_base64_by_url(self, url: str) -> Optional[str]: + """根据URL获取base64(带查重) + Args: + url: 图像URL + Returns: + str: base64字符串,失败返回None + """ + try: + image_path = await self.get_image_by_url(url) + if not image_path: + return None + + with open(image_path, 'rb') as f: + image_bytes = f.read() + return base64.b64encode(image_bytes).decode('utf-8') + + except Exception as e: + logger.error(f"获取base64失败: {str(e)}") + return None + + async def save_base64_image(self, base64_str: str, description: str = None) -> Optional[str]: + """保存base64图像(带查重) + Args: + base64_str: base64字符串 + description: 图像描述 + Returns: + str: 保存路径,失败返回None + """ + return await self.save_image(base64_str, description=description, is_base64=True) + + def check_url_exists(self, url: str) -> bool: + """检查URL是否已存在 + Args: + url: 图像URL + Returns: + bool: 是否存在 + """ + return self.db.db.images.find_one({'url': url}) is not None + + def check_hash_exists(self, image_data: Union[str, bytes], is_base64: bool = False) -> bool: + """检查图像是否已存在 + Args: + image_data: 图像数据(base64或字节) + is_base64: 是否为base64格式 + Returns: + bool: 是否存在 + """ + try: + if is_base64: + if isinstance(image_data, str): + image_bytes = base64.b64decode(image_data) + else: + return False + else: + if isinstance(image_data, bytes): + image_bytes = image_data + else: + return False + + image_hash = hashlib.md5(image_bytes).hexdigest() + return self.db.db.images.find_one({'hash': image_hash}) is not None + + except Exception as e: + logger.error(f"检查哈希失败: {str(e)}") + return False + + async def get_emoji_description(self, image_base64: str) -> str: + """获取表情包描述，带查重和保存功能""" + try: + # 计算图片哈希 + image_bytes = base64.b64decode(image_base64) + image_hash = hashlib.md5(image_bytes).hexdigest() + + # 查询缓存的描述 + cached_description = await self._get_description_from_db(image_hash, 'emoji') + if cached_description: + return f"[表情包：{cached_description}]" + + # 调用AI获取描述 + prompt = "这是一个表情包，使用中文简洁的描述一下表情包的内容和表情包所表达的情感" + description, _ = await self._llm.generate_response_for_image(prompt, image_base64) + + # 根据配置决定是否保存图片 + if global_config.EMOJI_SAVE: + # 生成文件名和路径 + timestamp = int(time.time()) + filename = f"emoji_{timestamp}_{image_hash[:8]}.jpg" + file_path = os.path.join(self.IMAGE_DIR, filename) + + try: + # 保存文件 + with open(file_path, "wb") as f: + f.write(image_bytes) + + # 保存到数据库 + image_doc = { + 'hash': image_hash, + 'path': file_path, + 'type': 'emoji', + 'description': description, + 'timestamp': timestamp + } + await self.db.db.images.update_one( + {'hash': image_hash}, + {'$set': image_doc}, + upsert=True + ) + logger.success(f"保存表情包: {file_path}") + except Exception as e: + logger.error(f"保存表情包文件失败: {str(e)}") + + # 保存描述到数据库 + await self._save_description_to_db(image_hash, description, 'emoji') + + return f"[表情包：{description}]" + except Exception as e: + logger.error(f"获取表情包描述失败: {str(e)}") + return "[表情包]" + + async def get_image_description(self, image_base64: str) -> str: + """获取普通图片描述，带查重和保存功能""" + try: + # 计算图片哈希 + image_bytes = base64.b64decode(image_base64) + image_hash = hashlib.md5(image_bytes).hexdigest() + + # 查询缓存的描述 + cached_description = await self._get_description_from_db(image_hash, 'image') + if cached_description: + return f"[图片：{cached_description}]" + + # 调用AI获取描述 + prompt = "请用中文描述这张图片的内容。如果有文字，请把文字都描述出来。并尝试猜测这个图片的含义。最多200个字。" + description, _ = await self._llm.generate_response_for_image(prompt, image_base64) + + # 根据配置决定是否保存图片 + if global_config.EMOJI_SAVE: + # 生成文件名和路径 + timestamp = int(time.time()) + filename = f"image_{timestamp}_{image_hash[:8]}.jpg" + file_path = os.path.join(self.IMAGE_DIR, filename) + + try: + # 保存文件 + with open(file_path, "wb") as f: + f.write(image_bytes) + + # 保存到数据库 + image_doc = { + 'hash': image_hash, + 'path': file_path, + 'type': 'image', + 'description': description, + 'timestamp': timestamp + } + await self.db.db.images.update_one( + {'hash': image_hash}, + {'$set': image_doc}, + upsert=True + ) + logger.success(f"保存图片: {file_path}") + except Exception as e: + logger.error(f"保存图片文件失败: {str(e)}") + + # 保存描述到数据库 + await self._save_description_to_db(image_hash, description, 'image') + + return f"[图片：{description}]" + except Exception as e: + logger.error(f"获取图片描述失败: {str(e)}") + return "[图片]" + + + +# 创建全局单例 +image_manager = ImageManager() + def compress_base64_image_by_scale(base64_data: str, target_size: int = 0.8 * 1024 * 1024) -> str: """压缩base64格式的图片到指定大小