""" 多向量库管理器 - 支持按部门/权限隔离的向量知识库 功能： 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 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: """加载元数据（带文件锁）""" 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 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 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 not in ('text',): # 遇到非文本类型，停止查找 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', '') # 判断是否为跨页表格 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: is_cross_page = True # 规则2: 第二个表格标题或内容包含"续表" elif '续表' in next_title or '续表' in next_content: is_cross_page = True # 规则3: 标题相似（去掉"续表"后比较） elif curr_title and next_title: clean_next = next_title.replace('续表', '').strip() if curr_title in clean_next or clean_next in curr_title: is_cross_page = True 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) merged_images = [] if curr_img: merged_images.append({'id': curr_img, 'page': curr_page_end}) if next_img: 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=100) 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=200 ) 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