GraphRAG Local Search(局部搜索)完全指南:基于实体的知识图谱推理与源码级解析
【免费下载链接】graphragA modular graph-based Retrieval-Augmented Generation (RAG) system项目地址: https://gitcode.com/GitHub_Trending/gr/graphrag
本指南系统讲解 GraphRAG 查询引擎中的 Local Search(局部搜索)方法——一种以知识图谱实体为核心入口、融合结构化图数据与原始文档文本块来回答问题的实体级检索生成方案。读者将掌握 Local Search 的完整方法论与数据流、LocalSearch与LocalSearchMixedContext的全部可调参数及默认值、内置 prompt 的引用规范,并能够结合源码在已完成索引的 GraphRAG 数据上独立实现与调优一个局部搜索引擎。
Local Search 是什么:基于实体的推理(Entity-based Reasoning)
Local Search 是 GraphRAG 查询引擎(Query Engine)提供的检索方式之一,它在查询时将知识图谱中的结构化数据与输入文档中的非结构化数据相结合,为 LLM 上下文补充与问题相关的实体信息。它尤其适合回答需要对输入文档中具体提到的某个(或某几个)实体有深入理解的问题,例如:“甘菊(chamomile)有哪些治疗功效?”这类问题关心的是特定实体的属性、关联关系与证据来源。
与 Global Search(面向整库的社区报告 map-reduce 摘要)和 DRIFT Search(结合社区洞察的扩展版局部搜索)不同,Local Search 的出发点是「把用户问题映射到一小批实体」,再以这些实体为图上的入口向下挖掘细节,因此计算开销相对可控,回答颗粒度更细、证据更贴近原文。GraphRAG 还内置了一版基础的向量检索(Basic Search)用于对照,便于开发者根据问题类型比较不同检索策略的输出差异(参见 Query Engine 总览)。
Methodology:从查询到答案的完整数据流
Local Search 的整体数据流可以概括为「实体映射 → 多路候选召回 → 排序过滤 → 单上下文窗口组装 → LLM 生成」。官方文档给出的数据流图如下:
结合源码实现,该流程图对应的实际步骤为:
- 输入组装:给定用户查询(
query)与可选的对话历史(conversation_history)。若存在对话历史,Local Search 会先取出最近若干轮的用户提问(由conversation_history_max_turns控制),拼接在当前查询之后一起参与实体映射,从而把上下文中的指代带入本次检索。见 mixed_context.py。 - 查询到实体的语义映射:用文本嵌入模型将「拼接后的查询」编码,在与索引阶段生成的entity description embedding向量库中做相似度检索,映射出语义相关的实体集合(
selected_entities)。映射逻辑map_query_to_entities位于 entity_extraction.py。 - 多路候选召回与排序过滤:以上述实体为图入口,通过「实体—文本单元」「实体—社区报告」「实体—实体关系」「实体—协变量」等多条映射,召回五类候选数据,再分别按相关性打分、截断:
- Prioritized Text Units:候选原始文本块,先按所属实体的命中顺序、再按块内关系数量降序排列后装入上下文(mixed_context.py);
- Prioritized Community Reports:把命中实体按「关联实体命中数」与社区
rank双重排序后选取的社区报告(同上文件 L224-L304); - Prioritized Entities / Relationships / Covariates:由
_build_local_context逐步把实体及其关系、协变量累加进上下文直到 token 预算耗尽(同上文件 L377-L493)。
- 组装上下文:将「实体上下文 + 关系/协变量上下文 + 社区上下文 + 文本单元上下文 + 会话历史」拼接成一个字符串(
context_chunks),并保证其总 token 数不超出预设的单上下文窗口预算。 - 生成答案:把拼接结果作为数据表塞进 system prompt 的
{context_data}槽位,与用户问题一起送入 LLM,生成带引用标注的回答。
整个过程源码位于 search.py 的LocalSearch.search(),第 2~4 步全部收敛在context_builder.build_context()一次调用内完成。
上下文预算机制:三类数据的占比分配
Local Search 的上下文预算分配是其核心工程细节,直接在LocalSearchMixedContext.build_context()中实现:
- 三个比例参数共同约束上下文:
community_prop(社区报告占比)、text_unit_prop(文本单元占比),其余local_prop = 1 - community_prop - text_unit_prop留给实体/关系/协变量(mixed_context.py)。约束条件:community_prop + text_unit_prop必须 ≤ 1,否则build_context会直接抛出ValueError(提示 "The sum of community_prop and text_unit_prop should not exceed 1.")。因此真正进入 context 的还有「local」这一隐藏通道,这是 DRIFT 等场景调参时最容易踩坑的点。 - 每个通道的实际 token 预算为
int(max_context_tokens * prop),其中max_context_tokens是单个上下文窗口的总体上限。例如默认community_prop=0.15、text_unit_prop=0.5时,约 65% 的窗口预算会留给原始文本单元,15% 留给社区报告,剩余约 35% 用于实体—关系—协变量,而完整的窗口还包含会话历史占用的空间——会话历史上下文会先从max_context_tokens中扣除(见 mixed_context.py)。 - 若
return_candidate_context=True,context builder 会额外返回所有候选记录(不限于被选入窗口的),并为每条记录附加in_context布尔标记,方便事后分析哪些数据被真正送入了 prompt——这是开发者在做召回质量评估时的有力工具。
LocalSearch 类配置参数详解
官方文档给出LocalSearch的核心参数清单,对照当前仓库 search.py 的实现,逐一说明如下:
| 参数 | 说明 | 仓库默认值 / 依据 |
|---|---|---|
model | 用于生成回答的 LLM chat completion 对象 | 必填,类型为LLMCompletion(graphrag_llm.completion中定义),由工厂方法创建 |
context_builder | 负责从知识模型对象集合中准备上下文数据的对象 | 必填,类型为LocalContextBuilder;Local Search 的标准实现是LocalSearchMixedContext |
tokenizer | 用于 token 计数的分词器 | 可选,默认取model.tokenizer,见 base.py |
system_prompt | 生成回答所用的 prompt 模板 | 可选,缺省使用 local_search_system_prompt.py 中的LOCAL_SEARCH_SYSTEM_PROMPT |
response_type | 期望回答类型与格式的自由文本描述 | 缺省为"multiple paragraphs";官方建议示例:Multiple Paragraphs、Multi-Page Report |
model_params | 传给 LLM 调用的额外参数(temperature、max_tokens 等) | 可选 dict,最终以**self.model_params展开传给completion_async。注意:官方文档写作llm_params,但当前仓库__init__中实际参数名为model_params,二者语义一致,对接时以源码签名为准 |
context_builder_params | 传给context_builder.build_context()的额外参数字典 | 可选 dict,见下文「核心参数与默认值」一节 |
callbacks | 可选回调函数集合,用于自定义 LLM 流式输出的on_llm_new_token等事件处理 | 可选,可为空列表 |
核心context_builder_params与默认值
上下文构建的绝大多数调优点都通过context_builder_params传入。下表汇总了参数、作用与当前仓库使用的默认配置(工厂方法在 factory.py 中的实际接线,也即LocalSearchConfig各字段的落地位置):
| 参数 | 作用 | 默认值 |
|---|---|---|
max_context_tokens | 单个上下文窗口的最大 token 数,应结合所选模型的上下文上限设定(如 8k 模型建议设为约 5000) | 12_000 |
text_unit_prop | 文本单元在窗口中的占比 | 0.5 |
community_prop | 社区报告在窗口中的占比(与text_unit_prop之和必须 ≤ 1) | 0.15 |
conversation_history_max_turns | 参与检索/生成的对话历史最大轮数 | 5 |
conversation_history_user_turns_only | 仅取用户提问轮作为对话历史上下文 | True(工厂固定传入) |
top_k_mapped_entities | 映射出的候选实体数量上限(映射时先按oversample_scaler=2超采样再精选去重) | 10 |
top_k_relationships | 每个实体可纳入的关系数量上限 | 10 |
include_entity_rank | 是否在实体表输出中附带 rank | 工厂设为True |
include_relationship_weight | 是否在关系表输出中附带权重 | 工厂设为True |
include_community_rank | 是否在社区报告输出中附带 rank | 工厂设为False |
return_candidate_context | 是否返回全部候选数据并附in_context标记 | False |
embedding_vectorstore_key | 实体向量库的 key 类型(EntityVectorStoreKey.ID或.TITLE,取决于向量库以实体 id 还是 title 为键) | EntityVectorStoreKey.ID |
关于response_type:它不会改变检索流程,而是被直接写入 system prompt 的---Target response length and format---槽位(见 search.py),因此你可以在不改 prompt 的情况下通过它精确控制输出的篇幅与体裁。
settings.yaml 中的 Local Search 配置与默认值
在使用 CLI 或配置文件驱动查询时,以上参数通过配置文件的local_search段暴露。LocalSearchConfig(local_search_config.py)定义了以下可配置字段,其默认值集中定义在 defaults.py 的LocalSearchDefaults中:
| settings.yaml 字段 | 说明 | 默认值 |
|---|---|---|
prompt | 自定义 local search prompt(覆盖内置模板) | None |
completion_model_id | 用于生成回答的模型 ID(指向 models 段中的配置) | DEFAULT_COMPLETION_MODEL_ID |
embedding_model_id | 用于实体/查询嵌入的模型 ID(需与索引阶段使用的嵌入模型保持一致) | DEFAULT_EMBEDDING_MODEL_ID |
text_unit_prop | 文本单元占比 | 0.5 |
community_prop | 社区报告占比 | 0.15 |
conversation_history_max_turns | 对话历史最大轮数 | 5 |
top_k_entities | 映射出的 top-k 实体数 | 10 |
top_k_relationships | top-k 关系数 | 10 |
max_context_tokens | 上下文窗口 token 上限 | 12_000 |
示例如下(字段均可省略,省略时使用上表默认值):
local_search: prompt: "" # 留空使用内置 LOCAL_SEARCH_SYSTEM_PROMPT text_unit_prop: 0.5 community_prop: 0.15 top_k_entities: 10 top_k_relationships: 10 max_context_tokens: 12000对配置文件的整体结构与加载方式感兴趣的读者,可进一步参考 Query 配置说明 与 models.md。
运行时调用链与流式接口
从源码看,一次 Local Search 的完整调用链为:
LocalSearch.search(query, conversation_history, ...) └─ context_builder.build_context(...) # LocalSearchMixedContext,纯检索侧 ├─ map_query_to_entities(...) # 实体映射(含嵌入检索) ├─ conversation_history.build_context() # 会话历史占用 token 计算 ├─ _build_community_context() # 社区报告通道 ├─ _build_local_context() # 实体/关系/协变量通道 └─ _build_text_unit_context() # 原始文本块通道 └─ system_prompt.format(context_data=..., response_type=...) # prompt 组装 └─ model.completion_async(messages, stream=True, **model_params) # LLM 生成 └─ 汇总为 SearchResult(含 token / 耗时 / 调用次数统计)LocalSearch.search()返回的是SearchResult(定义于 base.py),除response外还携带结构化指标:completion_time、总llm_calls、prompt_tokens、output_tokens,以及按阶段拆分的llm_calls_categories/prompt_tokens_categories/output_tokens_categories——其中build_context阶段与response阶段的 token 消耗被分别记录,方便你精确核算「上下文构建开销」与「生成开销」;异常时也会返回一条response=""的统计型SearchResult而非直接中断。
LocalSearch还实现了stream_search()(search.py),通过 async generator 逐 token 产出文本,便于做打字机式流式输出;配合callbacks中的on_llm_new_token/on_context钩子,可以构建自定义事件处理(如打印中间上下文)。
内置 System Prompt 与数据引用规范
默认的LOCAL_SEARCH_SYSTEM_PROMPT(local_search_system_prompt.py)为模型设定了三条重要行为准则:
- 只基于输入数据表作答,不知道就直说,禁止编造;
- 每条陈述必须附带数据引用,格式为
[Data: <dataset name> (record ids); ...],例如[Data: Sources (15, 16), Reports (1), Entities (5, 7); Relationships (23); Claims (2, 7, 34, 46, 64, +more)],其中 id 是数据记录的真实 id(而非下标); - 单条引用最多列 5 个记录 id,超出部分用
+more表示,避免引用过长撑爆输出。
模板中的{response_type}与{context_data}两个槽位分别由构造函数参数与build_context()结果填充;数据表以多段文本(Sources、Reports、Entities、Relationships、Claims 等命名)拼接而成,这些命名与mixed_context.py中的context_name(如"Sources"、"Reports")一一对应。若你的业务需要自定义引用格式或输出风格,可通过LocalSearchConfig.prompt或get_local_search_engine(system_prompt=...)整体替换。
How to Use:如何在代码中运行 Local Search
方式一:官方 Notebook 逐行体验
GraphRAG 提供了完整可运行的示例 Notebook:local_search.ipynb,其中同时演示了 Local Search 的上下文构建与问答。其输入数据(operation dulce数据集)与 Lancedb/Parquet 索引产物均位于 examples_notebooks/inputs 目录,便于直接对照索引输出理解各数据表来源。相关配套代码可参考 api_overview.ipynb 与 local_search.ipynb 中查询引擎的初始化部分。
方式二:通过工厂方法编程接入
在高阶 API 中,仓库提供了统一的引擎工厂 factory.py:
from graphrag.query.factory import get_local_search_engine engine = get_local_search_engine( config=config, # GraphRagConfig(含 local_search / models 段) reports=reports, # list[CommunityReport] text_units=text_units, # list[TextUnit] entities=entities, # list[Entity] relationships=relationships, # list[Relationship] covariates=covariates, # dict[str, list[Covariate]] response_type="Multiple Paragraphs", description_embedding_store=embedding_store, # 实体描述嵌入向量库 system_prompt=None, # 留空使用内置模板 ) result = await engine.search("What are the healing properties of chamomile?") print(result.response)工厂内部按config.local_search段自动完成:根据completion_model_id/embedding_model_id创建 LLM 与嵌入模型、从model_settings.call_args提取model_params、实例化LocalSearchMixedContext并把各配置项组装成context_builder_params。前置条件是数据(entities、relationships、community reports、text_units、covariates)已从索引中读出、且 entity description embedding 向量库已就绪——这部分在 Query Engine 中由query.indexer_adapters与各类 input retrieval 模块负责(参见 query/indexer_adapters.py),这也是为何 Local Search 只能运行在已完成 GraphRAG 索引的数据之上。API 层的用法可参考 query.py。
从源码看 Local Search 的适用边界与调优建议
结合 local_search.md 与源码实现,可总结出如下工程要点:
- 适用场景定位:Local Search 面向「理解特定实体」型问题。若问题需要纵观全局(如"这些数据中最显著的主题是什么"),应改用 Global Search;若希望在局部搜索基础上引入社区洞察以扩大事实覆盖广度,应评估 DRIFT Search——后者内部正是复用
LocalSearchMixedContext来初始化局部上下文(见 drift_context.py)。连 Question Generation(问题生成)也复用同一套上下文构建机制来生成候选追问(见 question_gen/local_gen.py),三者的检索底座是相通的。 - 实体映射质量决定上限:查询到实体的映射依赖「实体描述嵌入」向量库,因此检索所用的
embedding_model_id必须与索引阶段一致,否则会出现语义空间不匹配导致的召回偏差。 - 上下文预算再确认:
community_prop与text_unit_prop之和不得大于 1,三者间的比例直接决定答案是更"贴近原文证据"(提高text_unit_prop)还是更"高层概括"(提高community_prop);同时务必为会话历史留出 token 空间。 - 与模型窗口对齐:
max_context_tokens默认12_000,使用 8k 等小窗口模型时必须下调(源码注释建议约 5000),并可通过response_type控制输出长度,防止总 token 超限。 - 证据可追溯:内置 prompt 强制模型输出
[Data: ...]引用,且SearchResult中保留了context_records(结构化记录)与context_text(实际进入窗口的文本),可用于核对答案的每一条论断是否有据可依。
Local Search 是 GraphRAG 知识图谱能力中最贴近"实体级问答"的检索组件,理解其上下文预算模型与参数语义,是将其应用于事实密集型问答、可溯源分析等真实业务的前提。
【免费下载链接】graphragA modular graph-based Retrieval-Augmented Generation (RAG) system项目地址: https://gitcode.com/GitHub_Trending/gr/graphrag
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考