init: RAG 知识库服务初始提交
- 后端 API（Flask + Gunicorn） - RAG 引擎（混合检索 + 云端 Reranker + 引用溯源） - 文档解析（MinerU + 多格式支持） - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
This commit is contained in:
lacerate551
287
knowledge/lazy_enhance.py
Normal file
287
knowledge/lazy_enhance.py
Normal file
@@ -0,0 +1,287 @@
|
||||
"""
|
||||
懒加载增强模块（Phase 4）
|
||||
|
||||
按需调用 LLM/VLM 生成表格摘要和图片描述
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 缓存目录（扁平化）
|
||||
VLM_CACHE_DIR = Path(".data/cache/vlm")
|
||||
LLM_CACHE_DIR = Path(".data/cache/llm")
|
||||
|
||||
|
||||
def compute_file_hash(file_path: str) -> str:
|
||||
"""计算文件哈希"""
|
||||
try:
|
||||
with open(file_path, 'rb') as f:
|
||||
return hashlib.md5(f.read()).hexdigest()
|
||||
except Exception as e:
|
||||
logger.warning(f"计算文件哈希失败: {e}")
|
||||
return hashlib.md5(file_path.encode()).hexdigest()
|
||||
|
||||
|
||||
async def lazy_vlm_description(chunk_id: str, image_path: str, kb_name: str, metadata: dict = None) -> str:
|
||||
"""
|
||||
懒加载 VLM 描述
|
||||
|
||||
触发条件：图片切片被检索命中
|
||||
|
||||
Args:
|
||||
chunk_id: 切片 ID
|
||||
image_path: 图片路径（相对路径或绝对路径）
|
||||
kb_name: 知识库名称
|
||||
metadata: 图片元数据（包含 section、page、caption、上下文等）
|
||||
|
||||
Returns:
|
||||
VLM 生成的图片描述
|
||||
"""
|
||||
import os
|
||||
from knowledge.manager import get_kb_manager
|
||||
|
||||
# 构建完整图片路径
|
||||
if not os.path.isabs(image_path):
|
||||
full_image_path = os.path.join('.data/images', image_path)
|
||||
else:
|
||||
full_image_path = image_path
|
||||
|
||||
# 1. 检查缓存
|
||||
img_hash = compute_file_hash(full_image_path)
|
||||
cache_file = VLM_CACHE_DIR / f"{img_hash}.txt"
|
||||
if cache_file.exists():
|
||||
logger.info(f"VLM 缓存命中: {image_path}")
|
||||
return cache_file.read_text(encoding='utf-8')
|
||||
|
||||
# 2. 调用 VLM（传入元数据）
|
||||
logger.info(f"VLM 懒加载: {image_path}")
|
||||
kb_manager = get_kb_manager()
|
||||
description = kb_manager._generate_image_description(full_image_path, metadata=metadata)
|
||||
|
||||
# 3. 写入缓存
|
||||
VLM_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
cache_file.write_text(description, encoding='utf-8')
|
||||
|
||||
# 4. 更新向量库（metadata + embedding）
|
||||
try:
|
||||
collection = kb_manager.get_collection(kb_name)
|
||||
result = collection.get(ids=[chunk_id], include=['metadatas'])
|
||||
if result['metadatas']:
|
||||
# 更新 metadata
|
||||
new_metadata = {
|
||||
**result['metadatas'][0],
|
||||
'has_vlm_desc': True,
|
||||
'vlm_desc': description
|
||||
}
|
||||
|
||||
# 更新 embedding（使用 VLM 描述重新计算向量）
|
||||
# 这样 VLM 描述中的关键词（如"发电量"）才能参与相似度检索
|
||||
embedding_model = kb_manager.embedding_model
|
||||
if embedding_model:
|
||||
new_vector = embedding_model.encode(description).tolist()
|
||||
if isinstance(new_vector[0], list):
|
||||
new_vector = new_vector[0]
|
||||
|
||||
collection.update(
|
||||
ids=[chunk_id],
|
||||
metadatas=[new_metadata],
|
||||
embeddings=[new_vector],
|
||||
documents=[description] # 同时更新 document 字段
|
||||
)
|
||||
logger.info(f"已更新向量库 embedding: {chunk_id}")
|
||||
else:
|
||||
# 无 embedding 模型时只更新 metadata
|
||||
collection.update(
|
||||
ids=[chunk_id],
|
||||
metadatas=[new_metadata]
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"更新向量库失败: {e}")
|
||||
|
||||
return description
|
||||
|
||||
|
||||
async def lazy_table_summary(chunk_id: str, table_md: str, kb_name: str) -> str:
|
||||
"""
|
||||
懒加载表格摘要
|
||||
|
||||
触发条件：表格切片被检索命中且相关性 > 0.7
|
||||
|
||||
Args:
|
||||
chunk_id: 切片 ID
|
||||
table_md: 表格 Markdown 内容
|
||||
kb_name: 知识库名称
|
||||
|
||||
Returns:
|
||||
LLM 生成的表格摘要
|
||||
"""
|
||||
from knowledge.manager import get_kb_manager
|
||||
|
||||
# 1. 检查缓存
|
||||
table_hash = hashlib.md5(table_md.encode()).hexdigest()
|
||||
cache_file = LLM_CACHE_DIR / f"{table_hash}.txt"
|
||||
if cache_file.exists():
|
||||
logger.info(f"LLM 缓存命中: {chunk_id}")
|
||||
return cache_file.read_text(encoding='utf-8')
|
||||
|
||||
# 2. 调用 LLM
|
||||
logger.info(f"LLM 懒加载: {chunk_id}")
|
||||
kb_manager = get_kb_manager()
|
||||
summary = kb_manager._generate_table_summary(table_md, None)
|
||||
|
||||
# 3. 写入缓存
|
||||
LLM_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
cache_file.write_text(summary, encoding='utf-8')
|
||||
|
||||
# 4. 更新向量库（可选）
|
||||
try:
|
||||
collection = kb_manager.get_collection(kb_name)
|
||||
result = collection.get(ids=[chunk_id], include=['metadatas'])
|
||||
if result['metadatas']:
|
||||
# 新增摘要切片
|
||||
embedding_model = kb_manager.embedding_model
|
||||
vector = embedding_model.encode(summary).tolist()
|
||||
if isinstance(vector[0], list):
|
||||
vector = vector[0]
|
||||
|
||||
collection.add(
|
||||
ids=[f"{chunk_id}_summary"],
|
||||
embeddings=[vector],
|
||||
documents=[summary],
|
||||
metadatas=[{
|
||||
**result['metadatas'][0],
|
||||
'is_summary': True,
|
||||
'original_doc_id': chunk_id
|
||||
}]
|
||||
)
|
||||
# 更新原切片标记
|
||||
collection.update(
|
||||
ids=[chunk_id],
|
||||
metadatas=[{**result['metadatas'][0], 'has_summary': True}]
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"更新向量库失败: {e}")
|
||||
|
||||
return summary
|
||||
|
||||
|
||||
async def enhance_retrieved_chunks(contexts: list, query: str, kb_name: str):
|
||||
"""
|
||||
检索后增强：按需调用 LLM/VLM
|
||||
|
||||
Args:
|
||||
contexts: 检索上下文列表
|
||||
query: 用户查询
|
||||
kb_name: 知识库名称
|
||||
"""
|
||||
for ctx in contexts:
|
||||
meta = ctx.get('meta', {})
|
||||
chunk_type = meta.get('chunk_type', 'text')
|
||||
image_path = meta.get('image_path', '')
|
||||
|
||||
# 图片切片：懒加载 VLM 描述
|
||||
if chunk_type in ('image', 'chart') and not meta.get('has_vlm_desc'):
|
||||
if image_path:
|
||||
try:
|
||||
# 从 doc 字段中提取图号（上下文可能包含"见图2.5"等）
|
||||
doc_text = ctx.get('doc', '')
|
||||
import re
|
||||
|
||||
# 提取图号（从前文/后文中）
|
||||
figure_number = ""
|
||||
# 匹配 "见图2.5"、"图2.5"、"见图 2.5" 等
|
||||
fig_match = re.search(r'[见如]?图\s*(\d+\.?\d*)', doc_text)
|
||||
if fig_match:
|
||||
figure_number = fig_match.group(1)
|
||||
|
||||
# 如果 doc 中没有，尝试从 section 中提取
|
||||
section = meta.get('section') or meta.get('section_path', '')
|
||||
if not figure_number and section:
|
||||
fig_match = re.search(r'[见如]?图\s*(\d+\.?\d*)', section)
|
||||
if fig_match:
|
||||
figure_number = fig_match.group(1)
|
||||
|
||||
# 传入图片元数据，增强 VLM 描述
|
||||
image_metadata = {
|
||||
'section': section,
|
||||
'page': meta.get('page'),
|
||||
'caption': meta.get('caption', ''),
|
||||
'source': meta.get('source', ''),
|
||||
'figure_number': figure_number, # 添加提取的图号
|
||||
'doc_text': doc_text # 添加完整文档文本
|
||||
}
|
||||
vlm_desc = await lazy_vlm_description(
|
||||
meta.get('id', ''),
|
||||
image_path,
|
||||
kb_name,
|
||||
metadata=image_metadata
|
||||
)
|
||||
ctx['doc'] = vlm_desc
|
||||
ctx['vlm_enhanced'] = True
|
||||
except Exception as e:
|
||||
logger.warning(f"VLM 懒加载失败: {e}")
|
||||
|
||||
# 表格切片：同时处理摘要和关联图片的 VLM 描述
|
||||
elif chunk_type == 'table':
|
||||
doc_text = ctx.get('doc', '')
|
||||
|
||||
# 1. 懒加载表格摘要（高分切片）
|
||||
if not meta.get('has_summary'):
|
||||
score = meta.get('score', 0)
|
||||
if score > 0.7: # 只对高相关表格生成摘要
|
||||
try:
|
||||
summary = await lazy_table_summary(
|
||||
meta.get('id', ''),
|
||||
doc_text,
|
||||
kb_name
|
||||
)
|
||||
# 摘要作为补充信息
|
||||
ctx['summary'] = summary
|
||||
ctx['llm_enhanced'] = True
|
||||
except Exception as e:
|
||||
logger.warning(f"表格摘要懒加载失败: {e}")
|
||||
|
||||
# 2. 表格有关联图片时，懒加载 VLM 描述
|
||||
if image_path and not meta.get('has_vlm_desc'):
|
||||
try:
|
||||
import re
|
||||
|
||||
# 提取表号（如 "表2.2"、"见表2.1"）
|
||||
table_number = ""
|
||||
# 匹配 "表2.2"、"见表2.2"、"见表 2.2" 等
|
||||
table_match = re.search(r'[见如]?表\s*(\d+\.?\d*)', doc_text)
|
||||
if table_match:
|
||||
table_number = table_match.group(1)
|
||||
|
||||
# 如果 doc 中没有，尝试从 section 中提取
|
||||
section = meta.get('section') or meta.get('section_path', '')
|
||||
if not table_number and section:
|
||||
table_match = re.search(r'[见如]?表\s*(\d+\.?\d*)', section)
|
||||
if table_match:
|
||||
table_number = table_match.group(1)
|
||||
|
||||
# 构建表格图片元数据
|
||||
table_image_metadata = {
|
||||
'section': section,
|
||||
'page': meta.get('page'),
|
||||
'caption': meta.get('caption', ''),
|
||||
'source': meta.get('source', ''),
|
||||
'table_number': table_number, # 表号
|
||||
'figure_number': table_number, # 兼容字段
|
||||
'doc_text': doc_text,
|
||||
'is_table': True # 标记为表格图片
|
||||
}
|
||||
vlm_desc = await lazy_vlm_description(
|
||||
meta.get('id', ''),
|
||||
image_path,
|
||||
kb_name,
|
||||
metadata=table_image_metadata
|
||||
)
|
||||
# 表格图片描述作为补充信息
|
||||
ctx['image_description'] = vlm_desc
|
||||
ctx['vlm_enhanced'] = True
|
||||
except Exception as e:
|
||||
logger.warning(f"表格图片 VLM 懒加载失败: {e}")
|
||||
Reference in New Issue
Block a user