diff --git a/src/chat/message_receive/storage.py b/src/chat/message_receive/storage.py index 862354db7..23afe6c87 100644 --- a/src/chat/message_receive/storage.py +++ b/src/chat/message_receive/storage.py @@ -4,7 +4,7 @@ from typing import Union # from ...common.database.database import db # db is now Peewee's SqliteDatabase instance from .message import MessageSending, MessageRecv from .chat_stream import ChatStream -from ...common.database.database_model import Messages, RecalledMessages # Import Peewee models +from ...common.database.database_model import Messages, RecalledMessages, Images # Import Peewee models from src.common.logger import get_logger logger = get_logger("message_storage") @@ -25,6 +25,7 @@ class MessageStorage: # print(processed_plain_text) if processed_plain_text: + processed_plain_text = MessageStorage.replace_image_descriptions(processed_plain_text) filtered_processed_plain_text = re.sub(pattern, "", processed_plain_text, flags=re.DOTALL) else: filtered_processed_plain_text = "" @@ -136,3 +137,28 @@ class MessageStorage: except Exception as e: logger.error(f"更新消息ID失败: {e}") + + @staticmethod + def replace_image_descriptions(text: str) -> str: + """将[图片:描述]替换为[picid:image_id]""" + # 先检查文本中是否有图片标记 + pattern = r'\[图片:([^\]]+)\]' + matches = re.findall(pattern, text) + + if not matches: + logger.debug("文本中没有图片标记,直接返回原文本") + return text + def replace_match(match): + description = match.group(1).strip() + try: + image_record = (Images.select() + .where(Images.description == description) + .order_by(Images.timestamp.desc()) + .first()) + if image_record: + return f"[picid:{image_record.image_id}]" + else: + return match.group(0) # 保持原样 + except Exception as e: + return match.group(0) + return re.sub(r'\[图片:([^\]]+)\]', replace_match, text) diff --git a/src/chat/utils/utils_image.py b/src/chat/utils/utils_image.py index 25b753bab..eed65ad88 100644 --- a/src/chat/utils/utils_image.py +++ b/src/chat/utils/utils_image.py @@ -178,12 +178,24 @@ class ImageManager: """获取普通图片描述,带查重和保存功能""" try: # 计算图片哈希 - # 确保base64字符串只包含ASCII字符 if isinstance(image_base64, str): image_base64 = image_base64.encode("ascii", errors="ignore").decode("ascii") image_bytes = base64.b64decode(image_base64) image_hash = hashlib.md5(image_bytes).hexdigest() - image_format = Image.open(io.BytesIO(image_bytes)).format.lower() + + # 检查图片是否已存在 + existing_image = Images.get_or_none(Images.emoji_hash == image_hash) + if existing_image: + # 更新计数 + if hasattr(existing_image, 'count') and existing_image.count is not None: + existing_image.count += 1 + else: + existing_image.count = 1 + existing_image.save() + + # 如果已有描述,直接返回 + if existing_image.description: + return f"[图片:{existing_image.description}]" # 查询缓存的描述 cached_description = self._get_description_from_db(image_hash, "image") @@ -192,6 +204,7 @@ class ImageManager: return f"[图片:{cached_description}]" # 调用AI获取描述 + image_format = Image.open(io.BytesIO(image_bytes)).format.lower() prompt = "请用中文描述这张图片的内容。如果有文字,请把文字都描述出来,请留意其主题,直观感受,输出为一段平文本,最多50字" description, _ = await self._llm.generate_response_for_image(prompt, image_base64, image_format) @@ -199,17 +212,7 @@ class ImageManager: logger.warning("AI未能生成图片描述") return "[图片(描述生成失败)]" - # 再次检查缓存 - cached_description = self._get_description_from_db(image_hash, "image") - if cached_description: - logger.warning(f"虽然生成了描述,但是找到缓存图片描述 {cached_description}") - return f"[图片:{cached_description}]" - - logger.debug(f"描述是{description}") - - # 根据配置决定是否保存图片 - - # 生成文件名和路径 + # 保存图片和描述 current_timestamp = time.time() filename = f"{int(current_timestamp)}_{image_hash[:8]}.{image_format}" image_dir = os.path.join(self.IMAGE_DIR, "image") @@ -221,26 +224,31 @@ class ImageManager: with open(file_path, "wb") as f: f.write(image_bytes) - # 保存到数据库 (Images表) - try: - img_obj = Images.get((Images.emoji_hash == image_hash) & (Images.type == "image")) - img_obj.path = file_path - img_obj.description = description - img_obj.timestamp = current_timestamp - img_obj.save() - except Images.DoesNotExist: + # 保存到数据库,补充缺失字段 + if existing_image: + existing_image.path = file_path + existing_image.description = description + existing_image.timestamp = current_timestamp + if not hasattr(existing_image, 'image_id') or not existing_image.image_id: + existing_image.image_id = str(uuid.uuid4()) + if not hasattr(existing_image, 'vlm_processed') or existing_image.vlm_processed is None: + existing_image.vlm_processed = True + existing_image.save() + else: Images.create( + image_id=str(uuid.uuid4()), emoji_hash=image_hash, path=file_path, type="image", description=description, timestamp=current_timestamp, + vlm_processed=True, + count=1, ) - logger.debug(f"保存图片元数据: {file_path}") except Exception as e: logger.error(f"保存图片文件或元数据失败: {str(e)}") - # 保存描述到数据库 (ImageDescriptions表) + # 保存描述到ImageDescriptions表 self._save_description_to_db(image_hash, description, "image") return f"[图片:{description}]"