""" 知识库管理器 - 切片管理 Mixin 提供文档切片（Chunk）的 CRUD 操作，支持： - 新增切片：手动添加文本切片到向量库 - 修改切片：更新切片内容和元数据 - 删除切片：从向量库移除切片 - 查询切片：分页获取切片列表 切片是向量检索的基本单位，每个切片包含： - 文档内容（用于向量化） - 元数据（来源、页码、状态等） - 向量嵌入（由 embedding 模型生成） 主要方法： - add_chunk: 新增切片 - update_chunk: 修改切片 - delete_chunk: 删除切片 - list_chunks: 查询切片列表 """ import hashlib import logging from datetime import datetime from typing import List, Dict, Optional, Tuple logger = logging.getLogger(__name__) class ChunkMixin: """ 切片管理 Mixin 提供切片级别的 CRUD 操作，同时维护向量索引和 BM25 索引。 依赖属性（需由主类提供）： - self.get_collection: 获取向量库集合的方法 - self.get_bm25_index: 获取 BM25 索引的方法 - self.save_bm25_index: 保存 BM25 索引的方法 """ def add_chunk( self, kb_name: str, content: str, metadata: dict = None ) -> str: """ 新增单个切片 将文本内容切片添加到指定向量库，自动生成向量嵌入和 BM25 索引。 Args: kb_name: 目标向量库名称 content: 切片文本内容 metadata: 可选的元数据字典，可包含： - source: 来源文件名 - page: 页码 - chunk_type: 切片类型 ('text' | 'table' | 'image') - status: 状态 ('active' | 'deprecated') - 其他自定义字段 Returns: 新创建的切片 ID，格式为: manual_{时间戳}_{内容哈希} Raises: ValueError: 向量库不存在 Exception: 向量生成失败 Example: >>> chunk_id = kb_manager.add_chunk( ... "public_kb", ... "这是要添加的文本内容", ... {"source": "manual", "page": 1} ... ) """ collection = self.get_collection(kb_name) if not collection: raise ValueError(f"向量库 '{kb_name}' 不存在") content_hash = hashlib.md5(content.encode('utf-8')).hexdigest()[:8] chunk_id = f"manual_{datetime.now().strftime('%Y%m%d%H%M%S')}_{content_hash}" chunk_metadata = { "chunk_id": chunk_id, "chunk_type": "text", "collection": kb_name, "source": "manual", "status": "active", "version": "v1", "change_time": datetime.now().isoformat(), } if metadata: chunk_metadata.update(metadata) try: from core.engine import get_engine engine = get_engine() if not engine._initialized: engine.initialize() embedding = engine.embedding_model.encode(content).tolist() except Exception as e: logger.error(f"生成向量失败: {e}") raise collection.add( ids=[chunk_id], documents=[content], metadatas=[chunk_metadata], embeddings=[embedding] ) bm25 = self.get_bm25_index(kb_name) if bm25: bm25.add_documents([chunk_id], [content], [chunk_metadata]) self.save_bm25_index(kb_name) logger.info(f"新增切片: {chunk_id} -> {kb_name}") return chunk_id def update_chunk( self, kb_name: str, chunk_id: str, content: str = None, metadata: dict = None ) -> bool: """ 修改切片 更新切片的内容和/或元数据。如果更新内容，会重新生成向量嵌入。 Args: kb_name: 向量库名称 chunk_id: 切片 ID content: 新的文本内容（None 表示不更新） metadata: 要更新的元数据字段（None 表示不更新） Returns: 更新成功返回 True，切片不存在或更新失败返回 False Example: >>> kb_manager.update_chunk( ... "public_kb", ... "manual_20260517_abc123", ... content="更新后的内容", ... metadata={"status": "updated"} ... ) True """ collection = self.get_collection(kb_name) if not collection: return False existing = collection.get(ids=[chunk_id]) if not existing['ids']: logger.warning(f"切片不存在: {chunk_id}") return False update_kwargs = {"ids": [chunk_id]} if content is not None: update_kwargs["documents"] = [content] try: from core.engine import get_engine engine = get_engine() if not engine._initialized: engine.initialize() embedding = engine.embedding_model.encode(content).tolist() update_kwargs["embeddings"] = [embedding] except Exception as e: logger.error(f"生成向量失败: {e}") return False if metadata is not None: existing_metadata = existing['metadatas'][0] if existing['metadatas'] else {} existing_metadata.update(metadata) update_kwargs["metadatas"] = [existing_metadata] collection.update(**update_kwargs) if content is not None: bm25 = self.get_bm25_index(kb_name) if bm25: bm25.add_documents([chunk_id], [content], [metadata or {}]) self.save_bm25_index(kb_name) logger.info(f"更新切片: {chunk_id}") return True def get_chunk(self, kb_name: str, chunk_id: str) -> Optional[Dict]: """ 获取单个切片信息 Args: kb_name: 向量库名称 chunk_id: 切片 ID Returns: 切片信息字典，包含 id, document, metadata 等字段 切片不存在时返回 None """ collection = self.get_collection(kb_name) if not collection: return None result = collection.get(ids=[chunk_id], include=["documents", "metadatas"]) if not result['ids']: return None return { 'id': result['ids'][0], 'document': result['documents'][0] if result['documents'] else '', 'metadata': result['metadatas'][0] if result['metadatas'] else {} } def delete_chunk(self, kb_name: str, chunk_id: str) -> Tuple[bool, Optional[str]]: """ 删除切片 从向量库中移除指定的切片。 Args: kb_name: 向量库名称 chunk_id: 切片 ID Returns: 元组 (success, source_file)： - success: 删除是否成功 - source_file: 被删除切片的来源文件名（用于清理哈希记录） Note: 删除后 BM25 索引不会立即更新，需要手动调用 rebuild_bm25_index。 """ collection = self.get_collection(kb_name) if not collection: return False, None # 获取切片信息（删除前） existing = collection.get(ids=[chunk_id], include=["metadatas"]) if not existing['ids']: logger.warning(f"切片不存在: {chunk_id}") return False, None source_file = None if existing['metadatas'] and existing['metadatas'][0]: source_file = existing['metadatas'][0].get('source') collection.delete(ids=[chunk_id]) logger.info(f"删除切片: {chunk_id}") return True, source_file def delete_chunks_by_source(self, kb_name: str, source: str) -> int: """ 批量删除指定文件的所有切片 Args: kb_name: 向量库名称 source: 文件名（source 字段值） Returns: 删除的切片数量 Note: 删除后会重建 BM25 索引。 """ collection = self.get_collection(kb_name) if not collection: return 0 # 获取所有切片，找到匹配 source 的 all_chunks = collection.get(include=["metadatas"]) # 找到要删除的切片 ID chunk_ids_to_delete = [] for i, chunk_id in enumerate(all_chunks['ids']): metadata = all_chunks['metadatas'][i] if all_chunks['metadatas'] else {} if metadata.get('source') == source: chunk_ids_to_delete.append(chunk_id) if not chunk_ids_to_delete: logger.warning(f"未找到文件 {source} 的切片") return 0 # 批量删除 collection.delete(ids=chunk_ids_to_delete) logger.info(f"批量删除切片: {source}, 共 {len(chunk_ids_to_delete)} 个") # 重建 BM25 索引 try: 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) except Exception as e: logger.warning(f"重建 BM25 索引失败: {e}") return len(chunk_ids_to_delete) def list_chunks( self, kb_name: str, document_id: str = None, limit: int = 100, offset: int = 0 ) -> List[Dict]: """ 获取向量库中的切片列表 支持分页和按文档过滤。 Args: kb_name: 向量库名称 document_id: 文档 ID（source 字段），用于过滤特定文档的切片 limit: 返回数量限制（默认 100） offset: 偏移量（默认 0） Returns: 切片字典列表，每个字典包含： - id: 切片 ID - document: 切片内容 - metadata: 元数据字典 - status: 状态 - version: 版本号 Example: >>> chunks = kb_manager.list_chunks("public_kb", limit=10) >>> for chunk in chunks: ... print(chunk["id"], chunk["status"]) """ collection = self.get_collection(kb_name) if not collection: return [] where_filter = {} if document_id: where_filter["source"] = document_id try: result = collection.get( where=where_filter if where_filter else None, limit=limit, offset=offset ) except Exception as e: logger.warning(f"获取切片失败: {e}") return [] return [ { "id": id, "document": doc, "metadata": meta, "status": meta.get("status", "active") if meta else "active", "version": meta.get("version", "v1") if meta else "v1" } for id, doc, meta in zip( result['ids'], result.get('documents', []), result.get('metadatas', []) ) ]