init: RAG 知识库服务初始提交
- 后端 API（Flask + Gunicorn） - RAG 引擎（混合检索 + 云端 Reranker + 引用溯源） - 文档解析（MinerU + 多格式支持） - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
This commit is contained in:
lacerate551
379
knowledge/processing.py
Normal file
379
knowledge/processing.py
Normal file
@@ -0,0 +1,379 @@
|
||||
"""
|
||||
知识库管理器 - 图片/表格处理 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
|
||||
}
|
||||
|
||||
doc_path = docstore_dir / f"{doc_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
|
||||
}
|
||||
|
||||
doc_path = docstore_dir / f"{doc_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}")
|
||||
Reference in New Issue
Block a user