""" 知识库管理器 - 图片/表格处理 Mixin 提供图片和表格的智能处理功能，包括： - 图片过滤：判断图片是否值得处理（过滤 logo、icon 等垃圾图片） - 图片描述生成：生成用于向量检索的轻量级描述 - 表格摘要生成：生成表格的语义摘要 - 原始数据存储：将表格/图片引用存储到 DocStore 处理策略： - 图片：通过 VLM（视觉语言模型）生成描述，用于语义检索 - 表格：通过 LLM 生成摘要，提取关键字段和示例数据 主要方法： - should_process_image: 判断图片是否值得处理 - generate_image_short_summary: 生成图片短摘要 - generate_lightweight_image_description: 生成轻量级图片描述 """ import os import re import json import logging from pathlib import Path from typing import Optional, List logger = logging.getLogger(__name__) class ProcessingMixin: """ 图片/表格处理 Mixin 提供图片和表格的智能处理能力，支持： - 垃圾图片过滤（logo、icon、二维码等） - 图片描述生成（用于向量检索） - 表格摘要生成（提取语义信息） """ def should_process_image(self, image_path: str, context_text: str, caption: str = "") -> bool: """ 判断图片是否值得处理 通过多维度判断过滤无意义的图片： 1. 文件名过滤：排除 logo、icon、qr、watermark 等 2. 尺寸过滤：排除小于 100x100 的图片 3. 上下文判断：有 caption 或上下文文本时保留 Args: image_path: 图片文件路径 context_text: 图片周围的上下文文本 caption: 图片标题/说明 Returns: 值得处理返回 True，应该跳过返回 False Example: >>> should_process_image("/path/to/logo.png", "", "") False # 文件名包含 logo >>> should_process_image("/path/to/chart.png", "如图所示...", "图2.1 架构图") True # 有 caption """ filename = os.path.basename(image_path).lower() junk_keywords = ["logo", "icon", "qr", "watermark", "banner", "button", "avatar"] if any(kw in filename for kw in junk_keywords): logger.debug(f"图片过滤：文件名包含垃圾关键词 - {filename}") return False try: from PIL import Image with Image.open(image_path) as img: width, height = img.size if width < 100 or height < 100: logger.debug(f"图片过滤：尺寸过小 ({width}x{height}) - {filename}") return False except Exception as e: logger.debug(f"图片尺寸检查失败: {e}") pass if len(caption) >= 3: return True if len(context_text) >= 10: return True logger.debug(f"图片保留：{filename}") return True def _get_vlm_cache(self, image_path: str) -> Optional[str]: """ 检查是否有 VLM 缓存描述 通过图片 MD5 哈希查找缓存文件，避免重复调用 VLM。 Args: image_path: 图片文件路径 Returns: 缓存的描述文本；无缓存返回 None """ import hashlib try: if not os.path.exists(image_path): return None img_hash = hashlib.md5(open(image_path, 'rb').read()).hexdigest() cache_file = Path(f'.data/cache/vlm/{img_hash}.txt') if cache_file.exists(): return cache_file.read_text(encoding='utf-8') except Exception as e: logger.debug(f"检查 VLM 缓存失败: {e}") return None def generate_image_short_summary(self, chunk, page_info: dict, full_description: str = "") -> str: """ 生成图片短摘要（用于向量检索匹配） 提取关键信息构建简洁摘要，包括： - 图号/表号 - 关键数值（年份、金额等） - 章节主题 - 图表类型 Args: chunk: 文档切片对象 page_info: 页面信息字典 full_description: 完整描述（可选） Returns: 短摘要字符串，不超过 45 字符 Example: >>> generate_image_short_summary(chunk, {"section": "1.2 概述"}, "2020-2025年数据...") "图2.1：2020-2025年，柱状图" """ figure_number = "" table_number = "" section = page_info.get('section_path', '') or page_info.get('section', '') title = chunk.title if hasattr(chunk, 'title') and chunk.title else "" caption = page_info.get('caption', '') for source in [section, title, caption, full_description]: if not source: continue fig_match = re.search(r'图\s*(\d+\.?\d*)', str(source)) if fig_match and not figure_number: figure_number = fig_match.group(1) table_match = re.search(r'表\s*(\d+\.?\d*)', str(source)) if table_match and not table_number: table_number = table_match.group(1) chunk_type = page_info.get('chunk_type', 'image') keywords = [] year_match = re.search(r'(\d{4})\s*[-至到]\s*(\d{4})', full_description) if year_match: keywords.append(f"{year_match.group(1)}-{year_match.group(2)}年") num_match = re.search(r'(\d+\.?\d*)\s*(亿|万|千瓦时|吨|米)', full_description) if num_match: keywords.append(f"{num_match.group(1)}{num_match.group(2)}") if section: section_parts = section.split('>') if section_parts: last_part = section_parts[-1].strip() theme_match = re.search(r'(\d+\.?\d*)\s*(.+)', last_part) if theme_match: keywords.append(theme_match.group(2).strip()[:10]) if chunk_type == 'chart': type_label = "图表" if '柱状图' in full_description: type_label = "柱状图" elif '折线图' in full_description or '曲线图' in full_description: type_label = "折线图" elif '饼图' in full_description: type_label = "饼图" elif chunk_type == 'table': type_label = "表格" else: type_label = "图片" if figure_number: summary = f"图{figure_number}：" elif table_number: summary = f"表{table_number}：" else: summary = f"{type_label}：" if keywords: summary += "，".join(keywords[:3]) if full_description and len(keywords) < 2: first_sentence = full_description.split('。')[0][:30] if first_sentence: summary += first_sentence if len(summary) > 45: summary = summary[:42] + "..." return summary def _extract_figure_references(self, text: str) -> List[str]: """ 提取文本中的图表引用 从文本中提取所有"见图X.X"、"见表X.X"等引用。 Args: text: 待提取的文本 Returns: 图表编号列表（去重） """ references = [] fig_matches = re.findall(r'(?:[见如及和与])?图\s*(\d+\.?\d*)', text) references.extend(fig_matches) table_matches = re.findall(r'(?:[见如及和与])?表\s*(\d+\.?\d*)', text) references.extend(table_matches) return list(set(references)) def generate_lightweight_image_description(self, image_path: str, chunk, page_info: dict) -> str: """ 生成轻量级图片描述（用于语义检索） 构建包含上下文信息的描述，不调用 VLM， 仅利用已有的元数据信息。 Args: image_path: 图片路径（或图片标识） chunk: 文档切片对象 page_info: 页面信息字典 Returns: 多行描述字符串，包含： - 图号/表号 - 标题/caption - 章节位置 - 页码 - 前后文上下文 Example: >>> generate_lightweight_image_description("/path/to/img", chunk, page_info) 图2.1，系统架构图，位于「第一章 > 1.2 概述」，第5页 前文：系统由三个模块组成... 后文：如图所示，各模块之间... """ parts = [] chunk_type = page_info.get('chunk_type', 'image') is_chart = chunk_type == 'chart' type_label = "图表" if is_chart else "图片" figure_number = "" table_number = "" sources_to_check = [] section = page_info.get('section_path', '') or page_info.get('section', '') if section: sources_to_check.append(section) title = chunk.title if hasattr(chunk, 'title') and chunk.title else "" if title: sources_to_check.append(title) caption = page_info.get('caption', '') if caption: sources_to_check.append(caption) context_before = "" context_after = "" if hasattr(chunk, 'context_before') and chunk.context_before: context_before = chunk.context_before[:500] sources_to_check.append(context_before) if hasattr(chunk, 'context_after') and chunk.context_after: context_after = chunk.context_after[:500] sources_to_check.append(context_after) for source_text in sources_to_check: if not figure_number: fig_match = re.search(r'(?:[见如及和与])?图\s*(\d+\.?\d*)', source_text) if fig_match: figure_number = fig_match.group(1) if not table_number: table_match = re.search(r'(?:[见如及和与])?表\s*(\d+\.?\d*)', source_text) if table_match: table_number = table_match.group(1) page = page_info.get('page', 0) if figure_number: parts.append(f"图{figure_number}") if table_number: parts.append(f"表{table_number}") if caption and caption not in ("图片", "图表"): parts.append(caption) elif title and title not in ("图片", "图表"): parts.append(title) if section: parts.append(f"位于「{section}」") parts.append(f"第{page}页") description_parts = [f"{type_label}：{'，'.join(parts)}"] if context_before: description_parts.append(f"前文：{context_before}") if context_after: description_parts.append(f"后文：{context_after}") return "\n".join(description_parts) def _store_original_table(self, doc_id: str, table_md: str, metadata: dict) -> None: """ 存储原始表格到 DocStore 将表格的 Markdown 内容持久化存储，用于后续检索展示。 Args: doc_id: 文档 ID（切片 ID） table_md: 表格的 Markdown 内容 metadata: 元数据字典 """ try: docstore_dir = Path(".data/docstore") docstore_dir.mkdir(parents=True, exist_ok=True) record = { "content_type": "table", "markdown": table_md, "meta": metadata } # 使用 collection/doc_id 组合路径，防止跨库同名文件覆盖 coll = metadata.get('collection', '') safe_id = f"{coll}_{doc_id}" if coll else doc_id doc_path = docstore_dir / f"{safe_id}.json" with open(doc_path, 'w', encoding='utf-8') as f: json.dump(record, f, ensure_ascii=False, indent=2) except Exception as e: logger.warning(f"存储原始表格失败: {e}") def _store_image_reference(self, doc_id: str, image_path: str, metadata: dict) -> None: """ 存储图片引用到 DocStore 记录图片文件路径，用于后续访问和展示。 Args: doc_id: 文档 ID（切片 ID） image_path: 图片文件路径 metadata: 元数据字典 """ try: docstore_dir = Path(".data/docstore") docstore_dir.mkdir(parents=True, exist_ok=True) record = { "content_type": "image", "storage_type": "file", "file_path": image_path, "meta": metadata } # 使用 collection/doc_id 组合路径，防止跨库同名文件覆盖 coll = metadata.get('collection', '') safe_id = f"{coll}_{doc_id}" if coll else doc_id doc_path = docstore_dir / f"{safe_id}.json" with open(doc_path, 'w', encoding='utf-8') as f: json.dump(record, f, ensure_ascii=False, indent=2) except Exception as e: logger.warning(f"存储图片引用失败: {e}")