cb75b9b274
多库检索与存储修复: - RRF 融合去重改用 (collection, chunk_id) 复合键，修复同名文件结果被吞 - DocStore 存储路径加 collection 前缀，修复跨库同名切片数据覆盖 - search_multiple 去重改用复合键 - chunk_id 解析改用 rsplit 兼容下划线文件名 上传与版本管理修复: - 同名文件上传改为覆盖模式，自动清理旧切片 - 修复首次上传不创建版本记录 - 修复覆盖上传版本号回退到 v1 - sync ADDED 分支改用动态版本号生成 - _generate_version_id 改为基于全部版本递增 - 废止/恢复操作同步 SQLite 版本记录 - mark_document_as_superseded 改为仅更新 SQLite 删除清理修复: - 删除文档时同步清理 SQLite 版本记录和变更日志 - 删除向量库时同步清理该库所有版本记录 - cleanup 改为清理 SQLite 记录而非 ChromaDB 测试: - test_version_management.py: 27 条版本管理单元测试 - test_edge_cases.py: 28 条边界用例测试 - test_upload_dedup.py: 5 条上传去重测试 - e2e_risk_test.py: 27 条端到端风险测试 文档: - 新增风险边界问题修复注意事项.md（面向后端的对接文档） - 新增向量库边界风险分析.md - 更新多篇现有文档
386 lines
13 KiB
Python
386 lines
13 KiB
Python
"""
|
||
知识库管理器 - 图片/表格处理 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}")
|