""" GraphRAG索引器 负责构建层次化的知识图谱索引，包括实体提取、图谱构建、社区检测和摘要生成 """ from typing import List, Dict, Optional import asyncio try: from langchain_core.documents import Document except ImportError: try: from langchain_core.documents import Document except ImportError: try: from langchain_core.documents import Document except ImportError: from langchain.schema import Document try: from langchain_text_splitters import RecursiveCharacterTextSplitter except ImportError: from langchain.text_splitter import RecursiveCharacterTextSplitter from entity_extractor import EntityExtractor, EntityDeduplicator from knowledge_graph import KnowledgeGraph, CommunitySummarizer class GraphRAGIndexer: """GraphRAG索引器 - 实现Microsoft GraphRAG的索引流程""" def __init__(self, enable_async: bool = True, async_batch_size: int = 8): """初始化GraphRAG索引器 Args: enable_async: 是否启用异步处理（默认启用） async_batch_size: 异步并发批次大小（默认5个文档并发） """ print("🚀 初始化GraphRAG索引器...") self.entity_extractor = EntityExtractor(enable_async=enable_async) self.entity_deduplicator = EntityDeduplicator() self.knowledge_graph = KnowledgeGraph() self.community_summarizer = CommunitySummarizer() self.enable_async = enable_async self.async_batch_size = async_batch_size self.indexed = False mode = "异步模式" if enable_async else "同步模式" print(f"✅ GraphRAG索引器初始化完成 ({mode}, 并发数={async_batch_size})") def index_documents(self, documents: List[Document], batch_size: int = 10, save_path: Optional[str] = None) -> KnowledgeGraph: """ 对文档集合建立GraphRAG索引 工作流程（遵循Microsoft GraphRAG）: 1. 文档分块（已在document_processor中完成） 2. 实体和关系提取 3. 实体去重和合并 4. 构建知识图谱 5. 社区检测 6. 生成社区摘要 Args: documents: 文档列表 batch_size: 批处理大小 save_path: 保存路径 Returns: 构建好的知识图谱 """ print(f"\n{'='*50}") print(f"📊 开始GraphRAG索引流程") print(f" 文档数量: {len(documents)}") print(f"{'='*50}\n") # 步骤1: 实体和关系提取 print("📍 步骤 1/5: 实体和关系提取") extraction_results = [] if self.enable_async: # 异步批量处理模式 print(f"🚀 使用异步处理模式，并发数={self.async_batch_size}") extraction_results = self._extract_async(documents) else: # 同步处理模式（原有逻辑） print("🔄 使用同步处理模式") total_batches = (len(documents) - 1) // batch_size + 1 for i in range(0, len(documents), batch_size): batch = documents[i:i+batch_size] batch_num = i // batch_size + 1 print(f"\n⚙️ === 批次 {batch_num}/{total_batches} (文档 {i+1}-{min(i+batch_size, len(documents))}) ===") for idx, doc in enumerate(batch): doc_global_index = i + idx try: result = self.entity_extractor.extract_from_document( doc.page_content, doc_index=doc_global_index ) extraction_results.append(result) except Exception as e: print(f" ❌ 文档 #{doc_global_index + 1} 处理失败: {e}") # 添加空结果以保持索引一致 extraction_results.append({"entities": [], "relations": []}) print(f"✅ 批次 {batch_num}/{total_batches} 完成") # 步骤2: 实体去重 print("\n📍 步骤 2/5: 实体去重和合并") all_entities = [] all_relations = [] for result in extraction_results: all_entities.extend(result.get("entities", [])) all_relations.extend(result.get("relations", [])) dedup_result = self.entity_deduplicator.deduplicate_entities(all_entities) unique_entities = dedup_result["entities"] entity_mapping = dedup_result["mapping"] # 更新关系中的实体名称 mapped_relations = [] for relation in all_relations: source = entity_mapping.get(relation["source"], relation["source"]) target = entity_mapping.get(relation["target"], relation["target"]) mapped_relations.append({ **relation, "source": source, "target": target }) # 步骤3: 构建知识图谱 print("\n📍 步骤 3/5: 构建知识图谱") cleaned_results = [{ "entities": unique_entities, "relations": mapped_relations }] self.knowledge_graph.build_from_extractions(cleaned_results) # 步骤4: 社区检测 print("\n📍 步骤 4/5: 社区检测") self.knowledge_graph.detect_communities(algorithm="louvain") # 步骤5: 生成社区摘要 print("\n📍 步骤 5/5: 生成社区摘要") self.community_summarizer.summarize_all_communities(self.knowledge_graph) # 保存图谱 if save_path: self.knowledge_graph.save_to_file(save_path) self.indexed = True # 打印统计信息 print(f"\n{'='*50}") print("✅ GraphRAG索引构建完成!") stats = self.knowledge_graph.get_statistics() print(f"\n📊 统计信息:") print(f" - 节点数: {stats['num_nodes']}") print(f" - 边数: {stats['num_edges']}") print(f" - 社区数: {stats['num_communities']}") print(f" - 图密度: {stats['density']:.4f}") print(f"\n 实体类型分布:") for etype, count in stats['entity_types'].items(): print(f" • {etype}: {count}") print(f"{'='*50}\n") return self.knowledge_graph def _extract_async(self, documents: List[Document]) -> List[Dict]: """异步批量提取实体和关系 Args: documents: 文档列表 Returns: 提取结果列表 """ total_docs = len(documents) extraction_results = [] # 将文档分成多个异步批次 for i in range(0, total_docs, self.async_batch_size): batch_end = min(i + self.async_batch_size, total_docs) batch_num = i // self.async_batch_size + 1 total_batches = (total_docs - 1) // self.async_batch_size + 1 print(f"\n⚡ === 异步批次 {batch_num}/{total_batches} (文档 {i+1}-{batch_end}) ===") # 准备异步批次数据 async_batch = [ (documents[idx].page_content, idx) for idx in range(i, batch_end) ] # 异步执行当前批次 try: batch_results = asyncio.run( main=self.entity_extractor.extract_batch_async(async_batch) ) extraction_results.extend(batch_results) print(f"✅ 异步批次 {batch_num}/{total_batches} 完成") except Exception as e: print(f"❌ 异步批次 {batch_num} 失败: {e}") # 添加空结果 for _ in range(len(async_batch)): extraction_results.append({"entities": [], "relations": []}) return extraction_results def get_graph(self) -> KnowledgeGraph: """获取知识图谱""" if not self.indexed: print("⚠️ 图谱尚未构建，请先调用 index_documents()") return self.knowledge_graph def load_index(self, filepath: str) -> KnowledgeGraph: """加载已有的图谱索引""" print(f"📂 从文件加载图谱索引: {filepath}") self.knowledge_graph.load_from_file(filepath) self.indexed = True return self.knowledge_graph def initialize_graph_indexer(): """初始化GraphRAG索引器""" return GraphRAGIndexer()