import asyncio import json import os import shutil import tempfile from src.knowledge.base import KBNotFoundError, KnowledgeBase from src.knowledge.factory import KnowledgeBaseFactory from src.utils import logger from src.utils.datetime_utils import coerce_any_to_utc_datetime, utc_isoformat class KnowledgeBaseManager: """ 知识库管理器 统一管理多种类型的知识库实例，提供统一的外部接口 """ def __init__(self, work_dir: str): """ 初始化知识库管理器 Args: work_dir: 工作目录 """ self.work_dir = work_dir os.makedirs(work_dir, exist_ok=True) # 知识库实例缓存 {kb_type: kb_instance} self.kb_instances: dict[str, KnowledgeBase] = {} # 全局数据库元信息 {db_id: metadata_with_kb_type} self.global_databases_meta: dict[str, dict] = {} # 元数据锁 self._metadata_lock = asyncio.Lock() # 加载全局元数据 self._load_global_metadata() self._normalize_global_metadata() # 初始化已存在的知识库实例 self._initialize_existing_kbs() logger.info("KnowledgeBaseManager initialized") # 在后台运行数据一致性检测（不阻塞初始化） # try: # # 尝试获取当前事件循环，如果没有则创建新的 # try: # loop = asyncio.get_event_loop() # if loop.is_running(): # # 如果已经在事件循环中，创建任务 # asyncio.create_task(self.detect_data_inconsistencies()) # else: # # 如果事件循环未运行，直接运行 # loop.run_until_complete(self.detect_data_inconsistencies()) # except RuntimeError: # # 没有事件循环，创建一个来运行检测 # asyncio.run(self.detect_data_inconsistencies()) # except Exception as e: # logger.warning(f"初始化时运行数据一致性检测失败: {e}") def _load_global_metadata(self): """加载全局元数据""" meta_file = os.path.join(self.work_dir, "global_metadata.json") if os.path.exists(meta_file): try: with open(meta_file, encoding="utf-8") as f: data = json.load(f) self.global_databases_meta = data.get("databases", {}) logger.info(f"Loaded global metadata for {len(self.global_databases_meta)} databases") except Exception as e: logger.error(f"Failed to load global metadata: {e}") # 尝试从备份恢复 backup_file = f"{meta_file}.backup" if os.path.exists(backup_file): try: with open(backup_file, encoding="utf-8") as f: data = json.load(f) self.global_databases_meta = data.get("databases", {}) logger.info("Loaded global metadata from backup") # 恢复备份文件 shutil.copy2(backup_file, meta_file) return except Exception as backup_e: logger.error(f"Failed to load backup: {backup_e}") # 如果加载失败，初始化为空状态 logger.warning("Initializing empty global metadata") self.global_databases_meta = {} def _save_global_metadata(self): """保存全局元数据""" self._normalize_global_metadata() meta_file = os.path.join(self.work_dir, "global_metadata.json") backup_file = f"{meta_file}.backup" try: # 创建简单备份 if os.path.exists(meta_file): shutil.copy2(meta_file, backup_file) # 准备数据 data = {"databases": self.global_databases_meta, "updated_at": utc_isoformat(), "version": "2.0"} # 原子性写入（使用临时文件） with tempfile.NamedTemporaryFile( mode="w", dir=os.path.dirname(meta_file), prefix=".tmp_", suffix=".json", delete=False ) as tmp_file: json.dump(data, tmp_file, ensure_ascii=False, indent=2) temp_path = tmp_file.name os.replace(temp_path, meta_file) logger.debug("Saved global metadata") except Exception as e: logger.error(f"Failed to save global metadata: {e}") # 尝试恢复备份 if os.path.exists(backup_file): try: shutil.copy2(backup_file, meta_file) logger.info("Restored global metadata from backup") except Exception as restore_e: logger.error(f"Failed to restore backup: {restore_e}") raise e def _normalize_global_metadata(self) -> None: """Normalize stored timestamps within the global metadata cache.""" for meta in self.global_databases_meta.values(): if "created_at" in meta: try: dt_value = coerce_any_to_utc_datetime(meta.get("created_at")) if dt_value: meta["created_at"] = utc_isoformat(dt_value) continue except Exception as exc: # noqa: BLE001 logger.warning(f"Failed to normalize database metadata timestamp {meta.get('created_at')!r}: {exc}") def _initialize_existing_kbs(self): """初始化已存在的知识库实例""" kb_types_in_use = set() for db_meta in self.global_databases_meta.values(): kb_type = db_meta.get("kb_type", "lightrag") # 默认为lightrag kb_types_in_use.add(kb_type) # 为每种使用中的知识库类型创建实例 for kb_type in kb_types_in_use: try: self._get_or_create_kb_instance(kb_type) except Exception as e: logger.error(f"Failed to initialize {kb_type} knowledge base: {e}") def _get_or_create_kb_instance(self, kb_type: str) -> KnowledgeBase: """ 获取或创建知识库实例 Args: kb_type: 知识库类型 Returns: 知识库实例 """ if kb_type in self.kb_instances: return self.kb_instances[kb_type] # 创建新的知识库实例 kb_work_dir = os.path.join(self.work_dir, f"{kb_type}_data") kb_instance = KnowledgeBaseFactory.create(kb_type, kb_work_dir) self.kb_instances[kb_type] = kb_instance logger.info(f"Created {kb_type} knowledge base instance") return kb_instance def _get_kb_for_database(self, db_id: str) -> KnowledgeBase: """ 根据数据库ID获取对应的知识库实例 Args: db_id: 数据库ID Returns: 知识库实例 Raises: KBNotFoundError: 数据库不存在或知识库类型不支持 """ if db_id not in self.global_databases_meta: raise KBNotFoundError(f"Database {db_id} not found") kb_type = self.global_databases_meta[db_id].get("kb_type", "lightrag") if not KnowledgeBaseFactory.is_type_supported(kb_type): raise KBNotFoundError(f"Unsupported knowledge base type: {kb_type}") return self._get_or_create_kb_instance(kb_type) # ============================================================================= # 统一的外部接口 - 与原始 LightRagBasedKB 兼容 # ============================================================================= def get_kb(self, db_id: str) -> KnowledgeBase: """Public accessor to fetch the underlying knowledge base instance by database id. This provides a simple compatibility layer for callers that expect a `get_kb` method on the manager. """ return self._get_kb_for_database(db_id) def get_databases(self) -> dict: """获取所有数据库信息""" all_databases = [] # 收集所有知识库的数据库信息 for kb_type, kb_instance in self.kb_instances.items(): kb_databases = kb_instance.get_databases()["databases"] all_databases.extend(kb_databases) return {"databases": all_databases} async def create_database( self, database_name: str, description: str, kb_type: str, embed_info: dict | None = None, **kwargs ) -> dict: """ 创建数据库 Args: database_name: 数据库名称 description: 数据库描述 kb_type: 知识库类型，默认为lightrag embed_info: 嵌入模型信息 **kwargs: 其他配置参数，包括chunk_size和chunk_overlap Returns: 数据库信息字典 """ if not KnowledgeBaseFactory.is_type_supported(kb_type): available_types = list(KnowledgeBaseFactory.get_available_types().keys()) raise ValueError(f"Unsupported knowledge base type: {kb_type}. Available types: {available_types}") kb_instance = self._get_or_create_kb_instance(kb_type) db_info = kb_instance.create_database(database_name, description, embed_info, **kwargs) db_id = db_info["db_id"] async with self._metadata_lock: self.global_databases_meta[db_id] = { "name": database_name, "description": description, "kb_type": kb_type, "created_at": utc_isoformat(), "additional_params": kwargs.copy(), } self._save_global_metadata() logger.info(f"Created {kb_type} database: {database_name} ({db_id}) with {kwargs}") return db_info async def delete_database(self, db_id: str) -> dict: """删除数据库""" try: kb_instance = self._get_kb_for_database(db_id) result = kb_instance.delete_database(db_id) async with self._metadata_lock: if db_id in self.global_databases_meta: del self.global_databases_meta[db_id] self._save_global_metadata() return result except KBNotFoundError as e: logger.warning(f"Database {db_id} not found during deletion: {e}") return {"message": "删除成功"} async def add_content(self, db_id: str, items: list[str], params: dict | None = None) -> list[dict]: """添加内容（文件/URL）""" kb_instance = self._get_kb_for_database(db_id) return await kb_instance.add_content(db_id, items, params or {}) async def aquery(self, query_text: str, db_id: str, **kwargs) -> str: """异步查询知识库""" kb_instance = self._get_kb_for_database(db_id) return await kb_instance.aquery(query_text, db_id, **kwargs) async def export_data(self, db_id: str, format: str = "zip", **kwargs) -> str: """导出知识库数据""" kb_instance = self._get_kb_for_database(db_id) return await kb_instance.export_data(db_id, format=format, **kwargs) def query(self, query_text: str, db_id: str, **kwargs) -> str: """同步查询知识库（兼容性方法）""" kb_instance = self._get_kb_for_database(db_id) return kb_instance.query(query_text, db_id, **kwargs) def get_database_info(self, db_id: str) -> dict | None: """获取数据库详细信息""" try: kb_instance = self._get_kb_for_database(db_id) db_info = kb_instance.get_database_info(db_id) # 添加全局元数据中的additional_params信息 if db_info and db_id in self.global_databases_meta: global_meta = self.global_databases_meta[db_id] additional_params = global_meta.get("additional_params", {}) if additional_params: db_info["additional_params"] = additional_params return db_info except KBNotFoundError: return None async def delete_file(self, db_id: str, file_id: str) -> None: """删除文件""" kb_instance = self._get_kb_for_database(db_id) await kb_instance.delete_file(db_id, file_id) async def get_file_basic_info(self, db_id: str, file_id: str) -> dict: """获取文件基本信息（仅元数据）""" kb_instance = self._get_kb_for_database(db_id) return await kb_instance.get_file_basic_info(db_id, file_id) async def get_file_content(self, db_id: str, file_id: str) -> dict: """获取文件内容信息（chunks和lines）""" kb_instance = self._get_kb_for_database(db_id) return await kb_instance.get_file_content(db_id, file_id) async def get_file_info(self, db_id: str, file_id: str) -> dict: """获取文件完整信息（基本信息+内容信息）- 保持向后兼容""" kb_instance = self._get_kb_for_database(db_id) return await kb_instance.get_file_info(db_id, file_id) def get_db_upload_path(self, db_id: str | None = None) -> str: """获取数据库上传路径""" if db_id: try: kb_instance = self._get_kb_for_database(db_id) return kb_instance.get_db_upload_path(db_id) except KBNotFoundError: # 如果数据库不存在，创建通用上传路径 pass # 通用上传路径 general_uploads = os.path.join(self.work_dir, "uploads") os.makedirs(general_uploads, exist_ok=True) return general_uploads def file_existed_in_db(self, db_id: str | None, content_hash: str | None) -> bool: """检查指定数据库中是否存在相同内容哈希的文件""" if not db_id or not content_hash: return False try: kb_instance = self._get_kb_for_database(db_id) except KBNotFoundError: return False for file_info in kb_instance.files_meta.values(): if file_info.get("database_id") != db_id: continue if file_info.get("status") == "failed": continue if file_info.get("content_hash") == content_hash: return True return False async def update_database(self, db_id: str, name: str, description: str, llm_info: dict = None) -> dict: """更新数据库""" kb_instance = self._get_kb_for_database(db_id) result = kb_instance.update_database(db_id, name, description, llm_info) async with self._metadata_lock: if db_id in self.global_databases_meta: self.global_databases_meta[db_id]["name"] = name self.global_databases_meta[db_id]["description"] = description self._save_global_metadata() return result def get_retrievers(self) -> dict[str, dict]: """获取所有检索器""" all_retrievers = {} # 收集所有知识库的检索器 for kb_instance in self.kb_instances.values(): retrievers = kb_instance.get_retrievers() all_retrievers.update(retrievers) return all_retrievers # ============================================================================= # 管理器特有的方法 # ============================================================================= def get_supported_kb_types(self) -> dict[str, dict]: """获取支持的知识库类型""" return KnowledgeBaseFactory.get_available_types() def get_kb_instance_info(self) -> dict[str, dict]: """获取知识库实例信息""" info = {} for kb_type, kb_instance in self.kb_instances.items(): info[kb_type] = { "work_dir": kb_instance.work_dir, "database_count": len(kb_instance.databases_meta), "file_count": len(kb_instance.files_meta), } return info def get_statistics(self) -> dict: """获取统计信息""" stats = {"total_databases": len(self.global_databases_meta), "kb_types": {}, "total_files": 0} # 按知识库类型统计 for db_meta in self.global_databases_meta.values(): kb_type = db_meta.get("kb_type", "lightrag") if kb_type not in stats["kb_types"]: stats["kb_types"][kb_type] = 0 stats["kb_types"][kb_type] += 1 # 统计文件总数 for kb_instance in self.kb_instances.values(): stats["total_files"] += len(kb_instance.files_meta) return stats # ============================================================================= # 兼容性方法 - 为了支持现有的 graph_router.py # ============================================================================= async def _get_lightrag_instance(self, db_id: str): """ 获取 LightRAG 实例（兼容性方法） Args: db_id: 数据库ID Returns: LightRAG 实例，如果数据库不是 lightrag 类型则返回 None Raises: ValueError: 如果数据库不存在或不是 lightrag 类型 """ try: # 检查数据库是否存在 if db_id not in self.global_databases_meta: logger.error(f"Database {db_id} not found in global metadata") return None # 检查是否是 LightRAG 类型 kb_type = self.global_databases_meta[db_id].get("kb_type", "lightrag") if kb_type != "lightrag": logger.error(f"Database {db_id} is not a LightRAG type (actual type: {kb_type})") raise ValueError(f"Database {db_id} is not a LightRAG knowledge base") # 获取 LightRAG 知识库实例 kb_instance = self._get_kb_for_database(db_id) # 如果不是 LightRagKB 实例，返回错误 if not hasattr(kb_instance, "_get_lightrag_instance"): logger.error(f"Knowledge base instance for {db_id} is not LightRagKB") return None # 调用 LightRagKB 的方法获取 LightRAG 实例 return await kb_instance._get_lightrag_instance(db_id) except Exception as e: logger.error(f"Failed to get LightRAG instance for {db_id}: {e}") return None def is_lightrag_database(self, db_id: str) -> bool: """ 检查数据库是否是 LightRAG 类型 Args: db_id: 数据库ID Returns: 是否是 LightRAG 类型的数据库 """ if db_id not in self.global_databases_meta: return False kb_type = self.global_databases_meta[db_id].get("kb_type", "lightrag") return kb_type == "lightrag" def get_lightrag_databases(self) -> list[dict]: """ 获取所有 LightRAG 类型的数据库 Returns: LightRAG 数据库列表 """ lightrag_databases = [] all_databases = self.get_databases()["databases"] for db in all_databases: if db.get("kb_type", "lightrag") == "lightrag": lightrag_databases.append(db) return lightrag_databases # ============================================================================= # 数据一致性检测方法 # ============================================================================= async def detect_data_inconsistencies(self) -> dict: """ 检测向量数据库中存在但在 metadata 中缺失的数据 Returns: 包含不一致信息的字典，按知识库类型分组 """ inconsistencies = { "chroma": {"missing_collections": [], "missing_files": []}, "milvus": {"missing_collections": [], "missing_files": []}, "total_missing_collections": 0, "total_missing_files": 0, } logger.info("开始检测向量数据库与元数据的一致性...") # 检测 ChromaDB 数据不一致 if "chroma" in self.kb_instances: try: chroma_inconsistencies = await self._detect_chroma_inconsistencies() inconsistencies["chroma"] = chroma_inconsistencies inconsistencies["total_missing_collections"] += len(chroma_inconsistencies["missing_collections"]) inconsistencies["total_missing_files"] += len(chroma_inconsistencies["missing_files"]) except Exception as e: logger.error(f"检测 ChromaDB 数据不一致时出错: {e}") # 检测 Milvus 数据不一致 if "milvus" in self.kb_instances: try: milvus_inconsistencies = await self._detect_milvus_inconsistencies() inconsistencies["milvus"] = milvus_inconsistencies inconsistencies["total_missing_collections"] += len(milvus_inconsistencies["missing_collections"]) inconsistencies["total_missing_files"] += len(milvus_inconsistencies["missing_files"]) except Exception as e: logger.error(f"检测 Milvus 数据不一致时出错: {e}") # 输出检测结果到日志 self._log_inconsistencies(inconsistencies) return inconsistencies async def _detect_chroma_inconsistencies(self) -> dict: """检测 ChromaDB 中的数据不一致""" inconsistencies = {"missing_collections": [], "missing_files": []} chroma_kb = self.kb_instances["chroma"] # 获取 ChromaDB 中所有实际的集合 try: actual_collections = chroma_kb.chroma_client.list_collections() actual_collection_names = {col.name for col in actual_collections} # 获取 metadata 中记录的数据库ID metadata_collection_names = set() for db_id, db_meta in chroma_kb.databases_meta.items(): metadata_collection_names.add(db_id) # 找出存在于 ChromaDB 但不在 metadata 中的集合 missing_collections = actual_collection_names - metadata_collection_names for collection_name in missing_collections: # 跳过一些系统集合 if not collection_name.startswith("kb_"): continue collection_info = {"collection_name": collection_name, "detected_at": utc_isoformat()} # 尝试获取集合的基本信息 try: collection = chroma_kb.chroma_client.get_collection(name=collection_name) collection_info["count"] = collection.count() collection_info["metadata"] = collection.metadata except Exception as e: logger.warning(f"无法获取集合 {collection_name} 的详细信息: {e}") collection_info["count"] = "unknown" inconsistencies["missing_collections"].append(collection_info) logger.warning( f"发现 ChromaDB 中存在但 metadata 中缺失的集合: {collection_name} " f"(文档数: {collection_info['count']})" ) # 检查文件级别的不一致（针对已知的数据库） for db_id in metadata_collection_names: try: collection = chroma_kb.chroma_client.get_collection(name=db_id) actual_count = collection.count() # 获取 metadata 中记录的文件数量 metadata_files_count = sum( 1 for file_info in chroma_kb.files_meta.values() if file_info.get("database_id") == db_id ) # 如果向量数据库中有数据但 metadata 中没有文件记录，可能存在文件缺失 if actual_count > 0 and metadata_files_count == 0: inconsistencies["missing_files"].append( { "database_id": db_id, "vector_count": actual_count, "metadata_files_count": metadata_files_count, "detected_at": utc_isoformat(), } ) logger.warning( f"发现数据库 {db_id} 在 ChromaDB 中有 {actual_count} 条向量数据，但 metadata 中没有文件记录" ) except Exception as e: logger.debug(f"检查数据库 {db_id} 的文件一致性时出错: {e}") except Exception as e: logger.error(f"检测 ChromaDB 数据不一致时出错: {e}") return inconsistencies async def _detect_milvus_inconsistencies(self) -> dict: """检测 Milvus 中的数据不一致""" inconsistencies = {"missing_collections": [], "missing_files": []} milvus_kb = self.kb_instances["milvus"] try: from pymilvus import utility # 获取 Milvus 中所有实际的集合 actual_collection_names = set(utility.list_collections(using=milvus_kb.connection_alias)) # 获取 metadata 中记录的数据库ID metadata_collection_names = set(milvus_kb.databases_meta.keys()) # 找出存在于 Milvus 但不在 metadata 中的集合 missing_collections = actual_collection_names - metadata_collection_names for collection_name in missing_collections: # 跳过一些系统集合 if not collection_name.startswith("kb_"): continue collection_info = {"collection_name": collection_name, "detected_at": utc_isoformat()} # 尝试获取集合的基本信息 try: from pymilvus import Collection collection = Collection(name=collection_name, using=milvus_kb.connection_alias) collection_info["count"] = collection.num_entities collection_info["description"] = collection.description except Exception as e: logger.warning(f"无法获取集合 {collection_name} 的详细信息: {e}") collection_info["count"] = "unknown" inconsistencies["missing_collections"].append(collection_info) logger.warning( f"发现 Milvus 中存在但 metadata 中缺失的集合: {collection_name} " f"(实体数: {collection_info['count']})" ) # 检查文件级别的不一致（针对已知的数据库） for db_id in metadata_collection_names: try: if utility.has_collection(db_id, using=milvus_kb.connection_alias): from pymilvus import Collection collection = Collection(name=db_id, using=milvus_kb.connection_alias) actual_count = collection.num_entities # 获取 metadata 中记录的文件数量 metadata_files_count = sum( 1 for file_info in milvus_kb.files_meta.values() if file_info.get("database_id") == db_id ) # 如果向量数据库中有数据但 metadata 中没有文件记录，可能存在文件缺失 if actual_count > 0 and metadata_files_count == 0: inconsistencies["missing_files"].append( { "database_id": db_id, "vector_count": actual_count, "metadata_files_count": metadata_files_count, "detected_at": utc_isoformat(), } ) logger.warning( f"发现数据库 {db_id} 在 Milvus 中有 {actual_count} 条向量数据，" "但 metadata 中没有文件记录" ) except Exception as e: logger.debug(f"检查数据库 {db_id} 的文件一致性时出错: {e}") except Exception as e: logger.error(f"检测 Milvus 数据不一致时出错: {e}") return inconsistencies def _log_inconsistencies(self, inconsistencies: dict) -> None: """将不一致检测结果输出到日志""" total_missing_collections = inconsistencies["total_missing_collections"] total_missing_files = inconsistencies["total_missing_files"] if total_missing_collections == 0 and total_missing_files == 0: logger.info("数据一致性检测完成，未发现不一致情况") return logger.warning("=" * 80) logger.warning("数据一致性检测完成，发现以下不一致情况：") logger.warning("=" * 80) # ChromaDB 不一致情况 chroma_missing = inconsistencies["chroma"]["missing_collections"] chroma_files_missing = inconsistencies["chroma"]["missing_files"] if chroma_missing or chroma_files_missing: logger.warning("ChromaDB 不一致情况：") logger.warning(f" 缺失集合数量: {len(chroma_missing)}") for collection_info in chroma_missing: logger.warning(f" - 集合: {collection_info['collection_name']}, 向量数: {collection_info['count']}") logger.warning(f" 缺失文件记录数量: {len(chroma_files_missing)}") for file_info in chroma_files_missing: logger.warning( f" - 数据库: {file_info['database_id']}, 向量数: {file_info['vector_count']}, " f"元数据文件数: {file_info['metadata_files_count']}" ) # Milvus 不一致情况 milvus_missing = inconsistencies["milvus"]["missing_collections"] milvus_files_missing = inconsistencies["milvus"]["missing_files"] if milvus_missing or milvus_files_missing: logger.warning("Milvus 不一致情况：") logger.warning(f" 缺失集合数量: {len(milvus_missing)}") for collection_info in milvus_missing: logger.warning(f" - 集合: {collection_info['collection_name']}, 实体数: {collection_info['count']}") logger.warning(f" 缺失文件记录数量: {len(milvus_files_missing)}") for file_info in milvus_files_missing: logger.warning( f" - 数据库: {file_info['database_id']}, 向量数: {file_info['vector_count']}, " f"元数据文件数: {file_info['metadata_files_count']}" ) logger.warning("=" * 80) logger.warning(f"总计：缺失集合 {total_missing_collections} 个，缺失文件记录 {total_missing_files} 个") logger.warning("建议：检查这些不一致的数据，必要时进行数据清理或元数据修复") logger.warning("=" * 80) async def manual_consistency_check(self) -> dict: """ 手动触发数据一致性检测 Returns: 检测结果字典 """ logger.info("手动触发数据一致性检测...") return await self.detect_data_inconsistencies()