849 lines
32 KiB
Python
849 lines
32 KiB
Python
"""
|
||
多向量库管理器 - 支持按部门/权限隔离的向量知识库
|
||
|
||
功能：
|
||
1. 多向量库管理 - 创建、删除、列举向量库
|
||
2. 多 BM25 索引管理 - 每个向量库独立的 BM25 索引
|
||
3. 权限过滤 - 根据用户角色和部门返回可访问的向量库
|
||
4. 并行检索 - 支持同时检索多个向量库
|
||
|
||
向量库命名规范：
|
||
- public_kb: 公开知识库，所有人可访问
|
||
- dept_{部门名}: 部门知识库，如 dept_finance, dept_hr, dept_tech
|
||
|
||
使用方式：
|
||
from knowledge.manager import KnowledgeBaseManager
|
||
|
||
kb_manager = KnowledgeBaseManager()
|
||
|
||
# 获取向量库
|
||
collection = kb_manager.get_collection("dept_finance")
|
||
|
||
# 列出所有向量库
|
||
collections = kb_manager.list_collections()
|
||
|
||
# 获取用户可访问的向量库
|
||
accessible = kb_manager.get_accessible_collections("manager", "finance")
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
import threading
|
||
try:
|
||
import fcntl
|
||
_HAS_FCNTL = True
|
||
except ImportError:
|
||
_HAS_FCNTL = False # Windows 环境无 fcntl
|
||
from typing import List, Dict, Optional, Tuple
|
||
from pathlib import Path
|
||
import logging
|
||
|
||
import chromadb
|
||
|
||
# 从 base.py 导入基础类和常量
|
||
from .base import (
|
||
BM25Index,
|
||
CollectionInfo,
|
||
SearchResult,
|
||
_get_doc_type,
|
||
_extract_figure_number,
|
||
_extract_section,
|
||
_build_semantic_content_for_text,
|
||
_build_semantic_content_for_table,
|
||
CHROMA_DB_BASE_PATH,
|
||
BM25_INDEX_BASE_PATH,
|
||
KB_METADATA_FILE,
|
||
PUBLIC_KB_NAME,
|
||
DEFAULT_DEPARTMENTS,
|
||
DEPARTMENT_NAME_MAP,
|
||
normalize_department_name,
|
||
)
|
||
|
||
# 导入 Mixin 类
|
||
from .collection import CollectionMixin
|
||
from .document import DocumentMixin
|
||
from .chunk import ChunkMixin
|
||
from .index import IndexMixin
|
||
from .search import SearchMixin
|
||
from .processing import ProcessingMixin
|
||
from .permission import PermissionMixin
|
||
|
||
# 导入 LLM 工具函数
|
||
from core.llm_utils import call_llm
|
||
|
||
# 设置日志
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
# ==================== 多向量库管理器 ====================
|
||
|
||
class KnowledgeBaseManager(
|
||
CollectionMixin,
|
||
DocumentMixin,
|
||
ChunkMixin,
|
||
IndexMixin,
|
||
SearchMixin,
|
||
ProcessingMixin,
|
||
PermissionMixin
|
||
):
|
||
"""
|
||
多向量库管理器
|
||
|
||
管理多个独立的 ChromaDB 集合，每个集合对应一个知识库。
|
||
支持按部门隔离，每个部门有独立的向量库和 BM25 索引。
|
||
|
||
通过 Mixin 模式组合功能：
|
||
- CollectionMixin: 向量库管理
|
||
- DocumentMixin: 文档管理
|
||
- ChunkMixin: 切片管理
|
||
- IndexMixin: BM25 索引管理
|
||
- SearchMixin: 检索功能
|
||
- ProcessingMixin: 图片/表格处理
|
||
- PermissionMixin: 权限控制
|
||
"""
|
||
|
||
def __init__(self, base_path: str = None, bm25_base_path: str = None):
|
||
"""
|
||
初始化
|
||
|
||
Args:
|
||
base_path: 向量库存储路径
|
||
bm25_base_path: BM25 索引存储路径
|
||
"""
|
||
self.base_path = base_path or CHROMA_DB_BASE_PATH
|
||
self.bm25_base_path = bm25_base_path or BM25_INDEX_BASE_PATH
|
||
|
||
# 缓存
|
||
self._collections: Dict[str, chromadb.Collection] = {}
|
||
self._bm25_indexes: Dict[str, BM25Index] = {}
|
||
self._clients: Dict[str, chromadb.PersistentClient] = {}
|
||
self._lock = threading.Lock()
|
||
|
||
# 确保目录存在
|
||
os.makedirs(self.base_path, exist_ok=True)
|
||
os.makedirs(self.bm25_base_path, exist_ok=True)
|
||
|
||
# 加载元数据
|
||
self._metadata = self._load_metadata()
|
||
|
||
# 初始化公开知识库
|
||
self._ensure_public_kb()
|
||
|
||
# 扫描并列出所有已存在的向量库
|
||
existing_kbs = []
|
||
if os.path.exists(self.base_path):
|
||
for item in os.listdir(self.base_path):
|
||
if os.path.isdir(os.path.join(self.base_path, item)) and not item.startswith('.'):
|
||
if os.path.exists(os.path.join(self.base_path, item, 'chroma.sqlite3')):
|
||
existing_kbs.append(item)
|
||
logger.info(f"知识库管理器初始化完成，路径: {self.base_path}，发现 {len(existing_kbs)} 个向量库: {existing_kbs}")
|
||
|
||
def _load_metadata(self) -> dict:
|
||
"""加载元数据（带文件锁，确保多 worker 进程间一致）"""
|
||
metadata_path = os.path.join(self.base_path, KB_METADATA_FILE)
|
||
if os.path.exists(metadata_path):
|
||
try:
|
||
with open(metadata_path, 'r', encoding='utf-8') as f:
|
||
if _HAS_FCNTL:
|
||
fcntl.flock(f.fileno(), fcntl.LOCK_SH)
|
||
try:
|
||
return json.load(f)
|
||
finally:
|
||
if _HAS_FCNTL:
|
||
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
|
||
except Exception as e:
|
||
logger.error(f"加载元数据失败: {e}")
|
||
return {"collections": {}}
|
||
|
||
def _save_metadata(self):
|
||
"""保存元数据（带文件锁，防止并发写入数据覆盖）"""
|
||
metadata_path = os.path.join(self.base_path, KB_METADATA_FILE)
|
||
try:
|
||
with open(metadata_path, 'w', encoding='utf-8') as f:
|
||
if _HAS_FCNTL:
|
||
fcntl.flock(f.fileno(), fcntl.LOCK_EX)
|
||
try:
|
||
json.dump(self._metadata, f, ensure_ascii=False, indent=2)
|
||
f.flush()
|
||
os.fsync(f.fileno())
|
||
finally:
|
||
if _HAS_FCNTL:
|
||
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
|
||
except Exception as e:
|
||
logger.error(f"保存元数据失败: {e}")
|
||
|
||
def _ensure_public_kb(self):
|
||
"""确保公开知识库存在"""
|
||
if PUBLIC_KB_NAME not in self._metadata.get("collections", {}):
|
||
self.create_collection(
|
||
PUBLIC_KB_NAME,
|
||
display_name="公开知识库",
|
||
department="",
|
||
description="所有人可访问的公开文档"
|
||
)
|
||
|
||
# ==================== 文件添加 ====================
|
||
|
||
def add_file_to_kb(
|
||
self,
|
||
kb_name: str,
|
||
filepath: str,
|
||
embedding_model=None,
|
||
extra_metadata: dict = None,
|
||
enable_table_summary: bool = True,
|
||
enable_image_description: bool = False,
|
||
file_content: bytes = None
|
||
) -> int:
|
||
"""
|
||
添加文件到指定向量库（v6 支持企业文件系统）
|
||
|
||
使用统一的 parse_document() 入口，支持：
|
||
- PDF/DOCX/PPTX/图片 → MinerU 解析
|
||
- XLSX/XLS → Pandas 解析
|
||
- TXT → 文本解析
|
||
|
||
Args:
|
||
kb_name: 向量库名称
|
||
filepath: 文件路径（相对路径或绝对路径）
|
||
embedding_model: 向量模型（可选，默认使用 engine 的）
|
||
extra_metadata: 额外的元数据（如 status, version 等）
|
||
enable_table_summary: 是否启用表格摘要管道（LLM 生成摘要）
|
||
enable_image_description: 是否启用图片描述管道（VLM 生成描述）
|
||
file_content: 文件二进制内容（可选，用于企业文件系统集成）
|
||
|
||
Returns:
|
||
添加的片段数量
|
||
"""
|
||
collection = self.get_collection(kb_name)
|
||
if not collection:
|
||
raise ValueError(f"向量库 '{kb_name}' 不存在")
|
||
|
||
# 解析文档
|
||
from parsers import parse_document
|
||
result = parse_document(filepath, file_content=file_content)
|
||
|
||
if not result:
|
||
logger.warning(f"文档解析结果为空: {filepath}")
|
||
return 0
|
||
|
||
chunks = result.get('chunks', [])
|
||
if not chunks:
|
||
logger.warning(f"文档解析后无有效切片: {filepath}")
|
||
return 0
|
||
|
||
# 合并跨页表格
|
||
chunks = self._merge_cross_page_tables(chunks)
|
||
|
||
filename = Path(filepath).name
|
||
|
||
# 入库前清理同名旧切片，防止重复上传导致新旧切片共存
|
||
existing = collection.get(where={"source": filename})
|
||
if existing and existing['ids']:
|
||
old_count = len(existing['ids'])
|
||
collection.delete(ids=existing['ids'])
|
||
logger.info(f"替换模式: 清理旧切片 {filename} -> {kb_name}, 共 {old_count} 个")
|
||
|
||
# 清理关联的 DocStore 文件
|
||
try:
|
||
docstore_dir = Path('.data/docstore')
|
||
if docstore_dir.exists():
|
||
for ds_file in docstore_dir.glob(f'{kb_name}_{filename}_*.json'):
|
||
ds_file.unlink()
|
||
logger.debug(f"清理 DocStore: {ds_file.name}")
|
||
except Exception as e:
|
||
logger.warning(f"清理 DocStore 失败: {e}")
|
||
|
||
# 重建 BM25 索引（移除旧条目）
|
||
self._bm25_indexes.pop(kb_name, None)
|
||
bm25 = self.get_bm25_index(kb_name)
|
||
if bm25:
|
||
remaining = collection.get(include=["documents", "metadatas"])
|
||
if remaining['ids']:
|
||
bm25.add_documents(
|
||
remaining['ids'],
|
||
remaining['documents'] or [],
|
||
remaining['metadatas'] or []
|
||
)
|
||
self.save_bm25_index(kb_name)
|
||
|
||
# 准备向量模型
|
||
if embedding_model is None:
|
||
from core.engine import get_engine
|
||
engine = get_engine()
|
||
if not engine._initialized:
|
||
engine.initialize()
|
||
embedding_model = engine.embedding_model
|
||
|
||
# 处理切片
|
||
ids = []
|
||
documents = []
|
||
metadatas = []
|
||
embeddings = []
|
||
|
||
doc_type = _get_doc_type(filename)
|
||
|
||
for i, chunk in enumerate(chunks):
|
||
chunk_id = f"{filename}_{i}"
|
||
|
||
# 获取切片类型（优先从 chunk 直接获取，兼容 MinerUChunk 和其他格式）
|
||
chunk_type = getattr(chunk, 'chunk_type', None)
|
||
if not chunk_type:
|
||
page_info = getattr(chunk, 'page_info', {}) or {}
|
||
chunk_type = page_info.get('chunk_type', 'text')
|
||
|
||
# 获取页码信息（兼容 MinerUChunk 和 page_info 格式）
|
||
page_start = getattr(chunk, 'page_start', None)
|
||
page_end = getattr(chunk, 'page_end', None)
|
||
if page_start is None:
|
||
page_info = getattr(chunk, 'page_info', {}) or {}
|
||
page_start = page_info.get('page', 0)
|
||
page_end = page_info.get('page_end', page_start)
|
||
page_start = page_start or 0
|
||
page_end = page_end or page_start
|
||
|
||
# 获取章节信息
|
||
section_path = getattr(chunk, 'section_path', '') or ''
|
||
if not section_path:
|
||
page_info = getattr(chunk, 'page_info', {}) or {}
|
||
section_path = page_info.get('section_path', '') or page_info.get('section', '')
|
||
|
||
if chunk_type == 'table':
|
||
# 表格内容优先使用 table_html（完整表格），fallback 到 content（标题）
|
||
table_md = getattr(chunk, 'table_html', None) or chunk.content
|
||
semantic_content = _build_semantic_content_for_table(
|
||
table_md, {'page': page_start, 'page_end': page_end, 'section_path': section_path}, chunk, doc_type
|
||
)
|
||
elif chunk_type in ('image', 'chart'):
|
||
semantic_content = self.generate_lightweight_image_description(
|
||
chunk.content, chunk, {'page': page_start, 'page_end': page_end, 'section_path': section_path}
|
||
)
|
||
else:
|
||
semantic_content = _build_semantic_content_for_text(
|
||
chunk, {'page': page_start, 'page_end': page_end, 'section_path': section_path}, doc_type
|
||
)
|
||
|
||
# 构建元数据
|
||
metadata = {
|
||
"chunk_id": chunk_id,
|
||
"chunk_index": i,
|
||
"chunk_type": chunk_type,
|
||
"source": filename,
|
||
"collection": kb_name,
|
||
"doc_type": doc_type, # 文档类型(pdf/word/excel/ppt/other)，驱动前端差异化溯源展示
|
||
"page": page_start,
|
||
"page_end": page_end,
|
||
"section": section_path,
|
||
"status": "active",
|
||
"version": "v1",
|
||
}
|
||
|
||
if extra_metadata:
|
||
metadata.update(extra_metadata)
|
||
|
||
# 序列化图片信息（修复图片召回断链）
|
||
if hasattr(chunk, 'images') and chunk.images:
|
||
metadata['images_json'] = json.dumps(chunk.images, ensure_ascii=False)
|
||
|
||
if hasattr(chunk, 'image_path') and chunk.image_path:
|
||
metadata['image_path'] = chunk.image_path
|
||
|
||
# 生成向量
|
||
try:
|
||
embedding = embedding_model.encode(semantic_content).tolist()
|
||
except Exception as e:
|
||
logger.error(f"生成向量失败: {chunk_id}, 错误: {e}")
|
||
continue
|
||
|
||
ids.append(chunk_id)
|
||
documents.append(semantic_content)
|
||
metadatas.append(metadata)
|
||
embeddings.append(embedding)
|
||
|
||
# 处理表格摘要
|
||
if chunk_type == 'table' and enable_table_summary:
|
||
table_md = chunk.content
|
||
summary = self._generate_table_summary(table_md, chunk)
|
||
if summary:
|
||
# 存储原始表格到 DocStore
|
||
self._store_original_table(chunk_id, table_md, metadata)
|
||
|
||
# 处理图片描述
|
||
if chunk_type in ('image', 'chart') and enable_image_description:
|
||
image_path = chunk.content
|
||
if self.should_process_image(image_path, "", page_info.get('caption', '')):
|
||
desc = self._generate_image_description(image_path, metadata)
|
||
if desc:
|
||
self._store_image_reference(chunk_id, image_path, metadata)
|
||
|
||
if not ids:
|
||
return 0
|
||
|
||
# 批量添加到向量库
|
||
collection.add(
|
||
ids=ids,
|
||
documents=documents,
|
||
metadatas=metadatas,
|
||
embeddings=embeddings
|
||
)
|
||
|
||
# 更新 BM25 索引
|
||
bm25 = self.get_bm25_index(kb_name)
|
||
if bm25:
|
||
bm25.add_documents(ids, documents, metadatas)
|
||
self.save_bm25_index(kb_name)
|
||
|
||
logger.info(f"添加文件: {filename} -> {kb_name}, 片段数: {len(ids)}")
|
||
return len(ids)
|
||
|
||
def _merge_cross_page_tables(self, chunks: list) -> list:
|
||
"""
|
||
合并跨页表格
|
||
|
||
检测规则：
|
||
1. 表格切片后面跟着"续表"文本或另一个表格
|
||
2. 页码连续或无法判断时，通过标题匹配
|
||
3. 第二个表格标题包含"续表"或标题相似
|
||
|
||
合并操作：
|
||
1. 合并 table_html
|
||
2. 将两个 image_path 存入 images 字段
|
||
3. 更新 page_end
|
||
"""
|
||
if len(chunks) < 2:
|
||
return chunks
|
||
|
||
# 检测页码是否可靠：若所有 chunk 的 page_start 相同（如 Word 文档 page_idx 全为 0），
|
||
# 则页码信息不可用，需要启用降级合并规则
|
||
page_values = set(getattr(c, 'page_start', 0) for c in chunks)
|
||
pages_unavailable = len(page_values) <= 1
|
||
|
||
merged_chunks = []
|
||
i = 0
|
||
merge_count = 0
|
||
|
||
while i < len(chunks):
|
||
current = chunks[i]
|
||
|
||
# 检查是否为表格类型
|
||
if getattr(current, 'chunk_type', '') == 'table':
|
||
# 查找下一个表格（跳过中间的"续表"文本）
|
||
next_table_idx = None
|
||
next_chunk = None
|
||
intermediate_texts = [] # 收集中间文本用于降级判断
|
||
|
||
for j in range(i + 1, min(i + 4, len(chunks))): # 最多向前看3个切片
|
||
candidate = chunks[j]
|
||
candidate_type = getattr(candidate, 'chunk_type', '')
|
||
candidate_title = getattr(candidate, 'title', '') or ''
|
||
candidate_content = getattr(candidate, 'content', '')
|
||
|
||
if candidate_type == 'table':
|
||
next_table_idx = j
|
||
next_chunk = candidate
|
||
break
|
||
elif candidate_type == 'text' and ('续表' in candidate_title or '续表' in candidate_content):
|
||
# 遇到"续表"文本，继续查找下一个表格
|
||
continue
|
||
elif candidate_type == 'text':
|
||
# 非"续表"文本，收集后停止查找
|
||
intermediate_texts.append(candidate)
|
||
break
|
||
else:
|
||
# 遇到非文本类型，停止查找
|
||
break
|
||
|
||
if next_chunk is not None:
|
||
# 获取页码信息
|
||
curr_page_end = getattr(current, 'page_end', getattr(current, 'page_start', 0))
|
||
next_page_start = getattr(next_chunk, 'page_start', 0)
|
||
|
||
# 获取标题
|
||
curr_title = getattr(current, 'title', '') or ''
|
||
next_title = getattr(next_chunk, 'title', '') or ''
|
||
if isinstance(curr_title, list):
|
||
curr_title = curr_title[0] if curr_title else ''
|
||
if isinstance(next_title, list):
|
||
next_title = next_title[0] if next_title else ''
|
||
|
||
# 获取内容（用于检测"续表"）
|
||
next_content = getattr(next_chunk, 'content', '')
|
||
|
||
# 通用/无意义标题集合，这些标题不能用于"标题相似"判定
|
||
_GENERIC_TITLES = {'表格', 'table', '表格', ''}
|
||
|
||
# 判断是否为跨页表格
|
||
is_cross_page = False
|
||
|
||
# 规则1: 页码连续（如果页码有效）
|
||
page_valid = curr_page_end > 0 and next_page_start > 0
|
||
if page_valid and curr_page_end + 1 == next_page_start:
|
||
# 页码连续时，还需标题匹配或为通用标题才合并
|
||
# 避免把不同页面上不相关的表格错误合并
|
||
if curr_title == next_title or curr_title in _GENERIC_TITLES and next_title in _GENERIC_TITLES:
|
||
is_cross_page = True
|
||
elif curr_title and next_title:
|
||
clean_next_r1 = next_title.replace('续表', '').strip()
|
||
if curr_title in clean_next_r1 or clean_next_r1 in curr_title:
|
||
is_cross_page = True
|
||
|
||
# 规则2: 第二个表格标题或内容包含"续表"
|
||
elif '续表' in next_title or '续表' in next_content:
|
||
is_cross_page = True
|
||
|
||
# 规则3: 标题相似（去掉"续表"后比较）
|
||
# 排除通用标题（如"表格"），防止把所有标题为"表格"的相邻表格都误合并
|
||
elif (curr_title and next_title
|
||
and curr_title not in _GENERIC_TITLES
|
||
and next_title not in _GENERIC_TITLES):
|
||
clean_next = next_title.replace('续表', '').strip()
|
||
if clean_next and (curr_title in clean_next or clean_next in curr_title):
|
||
is_cross_page = True
|
||
|
||
# 规则4（降级）: 页码不可用（如 Word 文档 page_idx 全为 0）
|
||
# 仅当页码信息缺失时才启用此规则，避免 PDF 正常页码时被误合并
|
||
if (not is_cross_page
|
||
and pages_unavailable
|
||
and curr_title in _GENERIC_TITLES
|
||
and next_title in _GENERIC_TITLES):
|
||
# 检查中间文本是否暗示跨页延续（空、短文本、续表标记等）
|
||
has_separating_content = False
|
||
for text_chunk in intermediate_texts:
|
||
tc = (getattr(text_chunk, 'content', '') or '').strip()
|
||
tt = (getattr(text_chunk, 'title', '') or '').strip()
|
||
if not tc:
|
||
continue # 空文本不算分隔
|
||
if '续表' in tc or '续表' in tt:
|
||
continue # 续表标记，说明是跨页
|
||
# 有实质性中间内容（如分类标题"A3类：xxx"），不合并
|
||
has_separating_content = True
|
||
break
|
||
if not has_separating_content:
|
||
is_cross_page = True
|
||
logger.debug(f"降级合并(页码不可用): '{curr_title}' + '{next_title}'")
|
||
|
||
if is_cross_page:
|
||
# 执行合并
|
||
merge_count += 1
|
||
logger.info(f"合并跨页表格: {curr_title} (页{curr_page_end}) + {next_title} (页{next_page_start})")
|
||
|
||
# 合并 table_html
|
||
curr_html = getattr(current, 'table_html', '') or ''
|
||
next_html = getattr(next_chunk, 'table_html', '') or ''
|
||
if curr_html and next_html:
|
||
# 合并两个表格的 HTML
|
||
current.table_html = curr_html + '\n' + next_html
|
||
|
||
# 合并 image_path 和嵌入图片到 images
|
||
curr_img = getattr(current, 'image_path', None)
|
||
next_img = getattr(next_chunk, 'image_path', None)
|
||
curr_images = getattr(current, 'images', None) or []
|
||
next_images = getattr(next_chunk, 'images', None) or []
|
||
|
||
# 合并两个表格的所有图片（image_path + 嵌入图片）
|
||
merged_images = list(curr_images) # 保留当前表格的嵌入图片
|
||
# 添加 image_path 图片（如果不在列表中）
|
||
existing_ids = {img.get('id', '') for img in merged_images if isinstance(img, dict)}
|
||
if curr_img and curr_img not in existing_ids:
|
||
merged_images.append({'id': curr_img, 'page': curr_page_end})
|
||
existing_ids.add(curr_img)
|
||
for img in next_images: # 添加下一个表格的嵌入图片
|
||
img_id = img.get('id', '') if isinstance(img, dict) else ''
|
||
if img_id and img_id not in existing_ids:
|
||
merged_images.append(img)
|
||
existing_ids.add(img_id)
|
||
if next_img and next_img not in existing_ids:
|
||
merged_images.append({'id': next_img, 'page': next_page_start})
|
||
|
||
if merged_images:
|
||
current.images = merged_images
|
||
# 保留第一个图片作为主 image_path
|
||
current.image_path = curr_img
|
||
|
||
# 更新页码范围
|
||
current.page_end = getattr(next_chunk, 'page_end', next_page_start)
|
||
|
||
# 添加合并后的切片，跳过中间所有切片
|
||
merged_chunks.append(current)
|
||
i = next_table_idx + 1
|
||
continue
|
||
|
||
# 不需要合并，直接添加
|
||
merged_chunks.append(current)
|
||
i += 1
|
||
|
||
if merge_count > 0:
|
||
logger.info(f"跨页表格合并完成: {merge_count} 组表格被合并")
|
||
|
||
return merged_chunks
|
||
|
||
def _generate_table_summary(self, table_md: str, chunk) -> str:
|
||
"""生成表格摘要（LLM）"""
|
||
# 提取表头
|
||
lines = table_md.split('\n')
|
||
headers = []
|
||
for line in lines:
|
||
if line.startswith('|') and '---' not in line:
|
||
headers = [h.strip() for h in line.split('|') if h.strip()]
|
||
break
|
||
|
||
if not headers:
|
||
return ""
|
||
|
||
# 构建提示词
|
||
prompt = f"""请用一句话总结以下表格的主要内容，不超过50字。
|
||
|
||
表头：{', '.join(headers)}
|
||
|
||
表格内容（前5行）：
|
||
{chr(10).join(lines[:6])}
|
||
|
||
摘要："""
|
||
|
||
try:
|
||
from config import get_llm_client, DASHSCOPE_MODEL
|
||
client = get_llm_client()
|
||
summary = call_llm(client, prompt, DASHSCOPE_MODEL, max_tokens=512)
|
||
return summary.strip() if summary else ""
|
||
except Exception as e:
|
||
logger.warning(f"生成表格摘要失败: {e}")
|
||
return ""
|
||
|
||
@staticmethod
|
||
def _extract_table_title(table_md: str) -> str:
|
||
"""从表格 Markdown 中提取标题"""
|
||
lines = table_md.split('\n')
|
||
for line in lines[:3]:
|
||
if line.startswith('#'):
|
||
return line.lstrip('#').strip()
|
||
return ""
|
||
|
||
def _generate_image_description(self, image_path: str, metadata: dict = None) -> str:
|
||
"""生成图片描述（VLM）"""
|
||
# 检查缓存
|
||
cached = self._get_vlm_cache(image_path)
|
||
if cached:
|
||
return cached
|
||
|
||
# 调用 VLM
|
||
try:
|
||
import base64
|
||
from pathlib import Path
|
||
|
||
# 读取图片
|
||
img_path = Path(image_path)
|
||
if not img_path.exists():
|
||
return ""
|
||
|
||
img_data = base64.b64encode(img_path.read_bytes()).decode()
|
||
|
||
# 构建提示词
|
||
prompt = """请描述这张图片的内容，包括：
|
||
1. 图片类型（如流程图、架构图、数据图表等）
|
||
2. 主要内容和关键信息
|
||
3. 如果是图表，描述数据趋势或关键数值
|
||
|
||
描述应简洁，不超过100字。"""
|
||
|
||
# 调用 VLM（需要支持视觉的模型）
|
||
from config import DASHSCOPE_API_KEY, DASHSCOPE_BASE_URL, VLM_MODEL
|
||
from openai import OpenAI
|
||
|
||
client = OpenAI(api_key=DASHSCOPE_API_KEY, base_url=DASHSCOPE_BASE_URL)
|
||
|
||
response = client.chat.completions.create(
|
||
model=VLM_MODEL,
|
||
messages=[
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "text", "text": prompt},
|
||
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_data}"}}
|
||
]
|
||
}
|
||
],
|
||
max_tokens=512
|
||
)
|
||
|
||
description = response.choices[0].message.content
|
||
|
||
# 缓存结果
|
||
import hashlib
|
||
img_hash = hashlib.md5(img_path.read_bytes()).hexdigest()
|
||
cache_dir = Path('.data/cache/vlm')
|
||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||
(cache_dir / f'{img_hash}.txt').write_text(description, encoding='utf-8')
|
||
|
||
return description
|
||
|
||
except Exception as e:
|
||
logger.warning(f"生成图片描述失败: {e}")
|
||
return ""
|
||
|
||
def update_image_descriptions(self, kb_name: str) -> dict:
|
||
"""更新向量库中所有图片的描述"""
|
||
collection = self.get_collection(kb_name)
|
||
if not collection:
|
||
return {"success": False, "error": "向量库不存在"}
|
||
|
||
result = collection.get(where={"chunk_type": {"$in": ["image", "chart"]}})
|
||
|
||
if not result['ids']:
|
||
return {"success": True, "updated": 0, "message": "没有图片切片"}
|
||
|
||
updated = 0
|
||
failed = 0
|
||
|
||
for chunk_id, doc, meta in zip(
|
||
result['ids'],
|
||
result['documents'],
|
||
result['metadatas']
|
||
):
|
||
image_path = meta.get('image_path', '')
|
||
if not image_path:
|
||
continue
|
||
|
||
description = self._generate_image_description(image_path, meta)
|
||
if description:
|
||
try:
|
||
collection.update(
|
||
ids=[chunk_id],
|
||
documents=[description],
|
||
metadatas=[{**meta, 'image_description': description}]
|
||
)
|
||
updated += 1
|
||
except Exception as e:
|
||
logger.warning(f"更新图片描述失败: {chunk_id}, {e}")
|
||
failed += 1
|
||
|
||
# 重建 BM25 索引
|
||
self.rebuild_bm25_index(kb_name)
|
||
|
||
return {
|
||
"success": True,
|
||
"updated": updated,
|
||
"failed": failed,
|
||
"total": len(result['ids'])
|
||
}
|
||
|
||
# ==================== 文档生命周期 ====================
|
||
|
||
def mark_document_as_superseded(
|
||
self,
|
||
kb_name: str,
|
||
filename: str,
|
||
new_version: str = "",
|
||
reason: str = "版本更新"
|
||
) -> Dict:
|
||
"""
|
||
标记文档旧版本为已替代（仅更新 SQLite 版本记录）
|
||
|
||
ChromaDB 中的旧切片由 Phase 3 去重逻辑自动清理，
|
||
此方法只负责在 document_versions 表中将 active 版本改为 superseded。
|
||
|
||
Args:
|
||
kb_name: 向量库名称
|
||
filename: 文件名
|
||
new_version: 新版本号（如 "v2"），用于日志记录
|
||
reason: 替代原因
|
||
|
||
Returns:
|
||
{"success": True, "superseded_version": "v1"}
|
||
"""
|
||
from datetime import datetime
|
||
try:
|
||
from knowledge.document_versions import get_version_query
|
||
from data.db import get_connection
|
||
|
||
vq = get_version_query()
|
||
active = vq.get_active_version(kb_name, filename)
|
||
if not active:
|
||
return {"success": True, "superseded_version": None,
|
||
"message": "无 active 版本需要标记"}
|
||
|
||
superseded_date = datetime.now().isoformat()
|
||
with get_connection("knowledge") as conn:
|
||
conn.execute("""
|
||
UPDATE document_versions
|
||
SET status = 'superseded',
|
||
deprecated_date = ?,
|
||
deprecated_reason = ?
|
||
WHERE collection = ? AND document_id = ? AND version = ?
|
||
""", (superseded_date, reason,
|
||
kb_name, filename, active.version))
|
||
conn.commit()
|
||
|
||
# 记录变更日志
|
||
vq.log_version_change(
|
||
kb_name, filename,
|
||
change_type="supersede",
|
||
old_version=active.version,
|
||
new_version=new_version,
|
||
old_status="active",
|
||
new_status="superseded",
|
||
reason=reason
|
||
)
|
||
|
||
logger.info(
|
||
f"标记版本替代: {kb_name}/{filename} "
|
||
f"{active.version} -> {new_version or '(待创建)'}"
|
||
)
|
||
|
||
return {
|
||
"success": True,
|
||
"superseded_version": active.version,
|
||
"new_version": new_version,
|
||
"collection": kb_name
|
||
}
|
||
except Exception as e:
|
||
logger.warning(f"标记替代版本失败: {e}")
|
||
return {"success": False, "error": str(e)}
|
||
|
||
# ==================== 辅助检索方法 ====================
|
||
|
||
def search_with_status_filter(
|
||
self,
|
||
kb_name: str,
|
||
query_vector: List[float],
|
||
top_k: int = 5,
|
||
status_filter: str = None
|
||
) -> Optional[SearchResult]:
|
||
"""带状态过滤的检索"""
|
||
collection = self.get_collection(kb_name)
|
||
if not collection:
|
||
return None
|
||
|
||
where_filter = None
|
||
if status_filter:
|
||
where_filter = {"status": status_filter}
|
||
|
||
result = collection.query(
|
||
query_embeddings=[query_vector],
|
||
n_results=top_k,
|
||
where=where_filter
|
||
)
|
||
|
||
return SearchResult(
|
||
ids=result['ids'][0] if result['ids'] else [],
|
||
documents=result['documents'][0] if result['documents'] else [],
|
||
metadatas=result['metadatas'][0] if result['metadatas'] else [],
|
||
distances=result['distances'][0] if result['distances'] else [],
|
||
collection_name=kb_name
|
||
)
|
||
|
||
def _rebuild_bm25_index(self, kb_name: str):
|
||
"""重建 BM25 索引（私有方法，兼容旧代码）"""
|
||
return self.rebuild_bm25_index(kb_name)
|
||
|
||
|
||
# ==================== 全局实例 ====================
|
||
|
||
_kb_manager: Optional[KnowledgeBaseManager] = None
|
||
|
||
|
||
def get_kb_manager() -> KnowledgeBaseManager:
|
||
"""获取全局知识库管理器实例"""
|
||
global _kb_manager
|
||
if _kb_manager is None:
|
||
_kb_manager = KnowledgeBaseManager()
|
||
return _kb_manager
|