源码解读 - 微软GraphRAG框架

这几天，微软开源了GraphRAG——一个基于知识图谱构建的检索增强生成系统，迅速在技术圈里引起了不小的讨论。坦白说，这个框架关注了挺久，趁着热度认真啃了啃源码。这篇文章，就当作是一份阅读笔记，记录下对GraphRAG系统架构、关键概念和核心工作流的理解。

（本次分析的代码对应commit id为a22003c302bf4ffeefec76a09533acaf114ae7bb，更新于2024年7月5日。）

框架概述

在深入代码之前，有必要先搞清楚GraphRAG要解决的问题。论文里一针见血地指出了常规RAG的致命短板：

“然而，RAG在面对针对整个文本语料库的全局性问题时，比如‘这个数据集的核心主题是什么？’，就显得力不从心了。因为这本质上是一个聚焦于查询的总结性任务，而不是一个明确的检索任务。”

简单说，传统RAG玩不转这种高层次的概括性问题。GraphRAG的解决方案也很清晰，论文里给出了具体思路：用社区检测算法，把整个知识图谱拆解成一个个模块化的社区，然后让大模型自下而上地为每个社区生成摘要。最后，再通过一种“map-reduce”的方式，让每个社区并行回答查询，最终汇总成一个全局性的答案。

实现方式是什么？

和大多数RAG系统类似，GraphRAG的整个流水线也分为索引和查询两个阶段。索引阶段，用LLM来提取实体、关系和协变量等结构化信息，构建知识图谱。然后利用社区检测技术进行切分，再让LLM做摘要。到了查询阶段，就可以针对具体问题，把相关社区的摘要汇总起来，生成一个全局性的答案。

源码解析

官方文档写得很清楚了，但要理解实现细节，还得看源码。项目的整体结构长这样：

.
├── cache
├── config
├── emit
├── graph
│ ├── embedding
│ ├── extractors
│ │ ├── claims
│ │ ├── community_reports
│ │ ├── graph
│ │ └── summarize
│ ├── utils
│ └── visualization
├── input
├── llm
├── progress
├── reporting
├── storage
├── text_splitting
├── utils
├── verbs
│ ├── covariates
│ │ └── extract_covariates
│ │ └── strategies
│ │ └── graph_intelligence
│ ├── entities
│ │ ├── extraction
│ │ │ └── strategies
│ │ │ └── graph_intelligence
│ │ └── summarize
│ │ └── strategies
│ │ └── graph_intelligence
│ ├── graph
│ │ ├── clustering
│ │ │ └── strategies
│ │ ├── embed
│ │ │ └── strategies
│ │ ├── layout
│ │ │ └── methods
│ │ ├── merge
│ │ └── report
│ │ └── strategies
│ │ └── graph_intelligence
│ ├── overrides
│ └── text
│ ├── chunk
│ │ └── strategies
│ ├── embed
│ │ └── strategies
│ ├── replace
│ └── translate
│ └── strategies
└── workflows
└── v1

这里面隐藏了大量值得深挖的细节，接下来会结合具体代码来展开。

Workflows

程序运行时，控制台会打印出完整的workflow列表，这是理解整个索引流水线的关键入口：

⠹ GraphRAG Indexer 
├── Loading Input (InputFileType.text) - 1 files loaded (0 filtered) ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:00 0:00:00
├── create_base_text_units
├── create_base_extracted_entities
├── create_summarized_entities
├── create_base_entity_graph
├── create_final_entities
├── create_final_nodes
├── create_final_communities
├── join_text_units_to_entity_ids
├── create_final_relationships
├── join_text_units_to_relationship_ids
├── create_final_community_reports
├── create_final_text_units
├── create_base_documents
└── create_final_documents
All workflows completed successfully.

Index

索引阶段是整个项目的核心，流程也最为复杂。启动索引的命令很简单：

python -m graphrag.index --init --root ./ragtest
python -m graphrag.index --root ./ragtest

实际调用的是`graphrag/index/__main__.py`的主函数，通过argparse解析参数后，会转调`graphrag/index/cli.py`里的`index_cli`函数。

简单梳理一下调用链路（上图灰色标记的是核心函数）：

• cli.py::index_cli()函数首先处理用户输入，比如检查是否需要初始化目录（`--init`），这一逻辑由`cli.py::index_cli()`执行，主要看配置文件、prompt模板、.env文件等是否存在，不存在就创建。
• 真正执行索引的是内部函数`cli.py::index_cli()._run_workflow_async()`，这当中又涉及两个关键函数：`cli.py::_create_default_config()`和`run.py::run_pipeline_with_config()`。

限于篇幅，这里只讨论默认配置的运行流程。理清逻辑后，可以自行修改配置。

• 默认配置生成逻辑：`cli.py::_create_default_config()`先检查根目录及`settings.yaml`，然后调用`cli.py::_read_config_parameters()`读取系统配置，比如LLM、数据分块、缓存、存储等。接下来的操作非常关键：根据当前参数创建一个pipeline配置，具体代码在`create_pipeline_config.py::create_pipeline_config()`，这是整个项目中逻辑最复杂的模块之一。
• 深入`create_pipeline_config.py::create_pipeline_config()`的代码，会发现其核心逻辑如下：

result = PipelineConfig(
root_dir=settings.root_dir,
input=_get_pipeline_input_config(settings),
reporting=_get_reporting_config(settings),
storage=_get_storage_config(settings),
cache=_get_cache_config(settings),
workflows=[
*_document_workflows(settings, embedded_fields),
*_text_unit_workflows(settings, covariates_enabled, embedded_fields),
*_graph_workflows(settings, embedded_fields),
*_community_workflows(settings, covariates_enabled, embedded_fields),
*(_covariate_workflows(settings) if covariates_enabled else []),
],
)

这段代码的意图很清楚：根据不同的功能，生成完整的workflow序列。需要注意的是，这里并不考虑workflow之间的依赖关系，只是单纯地根据`workflows/v1`目录下的模板，生成一系列workflow。

• Pipeline执行逻辑：`run.py::run_pipeline_with_config()`先通过`load_pipeline_config.py::load_pipeline_config()`加载上一步生成的pipeline配置，然后创建目录（如cache、storage、input、output等），最后利用`run.py::run_pipeline()`依次执行每个workflow，并返回结果。

`run.py::run_pipeline()`才是真正的执行者，其核心逻辑包含两部分：

• 加载workflows：`workflows/load.py::load_workflows()`，除了常规创建工作流外，还会做拓扑排序。
• workflows/load.py::create_workflow()：利用现有模板创建不同的工作流。
• graphlib::topological_sort()：根据workflow间的依赖关系，计算DAG的拓扑排序。
• 执行过程中，还会调用`inject_workflow_data_dependencies()`进行依赖注入，`write_workflow_stats()`写入统计信息，`emit_workflow_output`保存输出。真正的`workflow.run`操作会在`write_workflow_stats()`之前执行。

Workflow

到目前为止，我们其实还没真正触及索引阶段的业务逻辑，只是理清了GraphRAG内置的这套pipeline编排系统。现在以一个具体的workflow——index/workflows/v1/create_final_entities.py——为例，看看它到底是怎么运作的。

• DataShaper：在谈workflow之前，有必要先认识一下`datashaper`，这是微软开源的一个工作流处理库，内置了许多组件（官方叫`Verb`）。通俗理解，datashaper就像一条流水线，每一步都对输入数据执行一种特定操作，比如裁剪、旋转、缩放。如果还不清楚，建议跑一下官方文档里的demo程序examples/single_verb，应该就能明白。从功能上看，它有点像Prefect。
• 知识图谱构建：对应的workflow是create_final_entities.py。翻看源码可以发现，它依赖于workflow:create_base_extracted_entities，同时定义了cluster_graph、embed_graph等操作。其中，cluster_graph采用了leiden策略，具体代码在index/verbs/graph/clustering/cluster_graph.py中：

from datashaper import TableContainer, VerbCallbacks, VerbInput, progress_iterable, verb

@verb(name="cluster_graph")
def cluster_graph(
input: VerbInput,
callbacks: VerbCallbacks,
strategy: dict[str, Any],
column: str,
to: str,
level_to: str | None = None,
**_kwargs,
) -> TableContainer:

可以看到，这里的`leiden`算法其实来自另一个图算法库`graspologic`。

• Pipeline：弄清楚了workflow的执行逻辑，再结合上节提到的编排日志或artifacts/stats.json文件，就能整理出完整的索引流水线。官方文档也给出了非常详细的说明，不过依据源码提取出来的pipeline似乎还是有些差异，有兴趣的朋友可以留言探讨。

Query

查询阶段的pipeline相对简单一些。执行查询的命令如下：

# Global search
python -m graphrag.query 
--root ./ragtest 
--method global 
"What are the top themes in this story?"

# Local search
python -m graphrag.query 
--root ./ragtest 
--method local 
"Who is Scrooge, and what are his main relationships?"

这里有Global/Local两种搜索模式。生成函数调用关系图后，整体结构会更加清晰。

graphrag/query/__main__.py中的主函数会根据参数不同，分别路由到cli::run_local_search()或cli::run_global_search()。

Global Search

cli::run_global_search()主要调用factories.py::get_global_search_engine()，返回一个GlobalSearch类的实例。这个类与LocalSearch类似，都是基于工厂模式创建的。其核心方法是structured_search/global_search/search.py::GlobalSearch.asearch()，它采用了map-reduce策略：先让大模型为每个社区的摘要并行生成答案，然后汇总所有答案，形成最终结果。作者在相关的prompts里给出了非常详细的说明。

也正是因为这种map-reduce机制，Global Search对token的消耗量相当大。

Local Search

与全局搜索类似，cli::run_local_search()主要调用factories.py::get_local_search_engine()，返回一个LocalSearch类的实例。这里的asearch()方法相对简单，直接根据上下文给出回复。这种模式更接近常规的RAG语义检索策略，token消耗也更少。

与Global Search不同，Local模式综合了nodes、community_reports、text_units、relationships、entities、covariates等多种源数据。官方文档对此有详细说明，这里就不再赘述了。

一些思考

• GraphRAG最核心的价值，在于它一定程度上解决了聚焦于查询的总结性（QFS）任务。就目前了解的情况来看，这应该是首创。之前思路最接近的项目是RAPTOR，不过后者并非基于知识图谱。QFS与多跳问答是现阶段RAG系统难以解决的痛点，但在数据分析等领域有广泛应用。尽管当下GraphRAG的使用成本还很高，但它至少提供了一种可能性。
• 相较于一般的知识图谱RAG项目，GraphRAG更令人印象深刻的是它内置的那套完整的工作流编排系统。这在其他RAG框架中并不多见，也是值得深挖的方向。基于模板定义好大部分工作流，只提供少部分的配置灵活性，让每一个环节都可控可追溯，而不是把一切都交给大模型。相比之下，常规的RAG部分（如embedding、retrieval）反而没有太多需要特别讨论的地方，尽管加了社区检测算法。
• 在知识图谱的处理颗粒度上，GraphRAG也做得很细致，比如社区检测的Leiden算法、综合多源数据的Local Search等。一个有意思的细节是：项目中的实体抽取完全采用了大模型实现，并且对三元组的schema没有设置任何约束。作者的逻辑是，反正后续要做相似性聚类，每次抽取的差异对最终结果影响不大。这一点在鲁棒性上确实是很大的提升。
• 当然，目前的GraphRAG也有不少值得改进的地方。比如，它创造了一些让人摸不着头脑的名词（如Emit、Claim、Verbs、Reporting），同时还“夹带私货”，用了一些微软自家相对小众的库，增加了理解成本。此外，模块化方面也有待加强，特别是与OpenAI相关的那部分，耦合度比较高。

综合来看，微软这次推出的GraphRAG框架确实干货满满，值得花些时间去仔细研读和思考。

Reference

• Welcome to GraphRAG (microsoft.github.io)
• GraphRAG: LLM-Derived Knowledge Graphs for RAG (youtube.com)
• GraphRAG is now on GitHub | Hacker News (ycombinator.com)
• GraphRAG & Ollama - intelligent Search of local data : r/LocalLLaMA (reddit.com)
• GraphRAG on narrative private data : r/LocalLLaMA (reddit.com)