GraphRAG-02-Query-查询引擎

模块概览

Query 模块提供多种查询模式，在构建好的知识图谱上执行智能检索和问答。

核心查询模式：

Global Search（全局搜索）：Map-Reduce 模式，适合宏观问题
Local Search（局部搜索）：基于实体关系的上下文检索
DRIFT Search（混合搜索）：结合局部和全局策略
Basic Search（基础搜索）：向量相似度检索

查询引擎架构图

flowchart TB
    QueryAPI[Query API]
    Factory[Search Factory]
    
    subgraph SearchEngines["查询引擎"]
        Global[Global Search]
        Local[Local Search]
        DRIFT[DRIFT Search]
        Basic[Basic Search]
    end
    
    subgraph ContextBuilders["上下文构建器"]
        GlobalContext[Global Context Builder]
        LocalContext[Local Context Builder]
        DRIFTContext[DRIFT Context Builder]
        BasicContext[Basic Context Builder]
    end
    
    subgraph DataSources["数据源"]
        Storage[Storage: Parquet Files]
        VectorStore[Vector Store]
    end
    
    LLM[LLM Service]
    
    QueryAPI --> Factory
    Factory --> SearchEngines
    SearchEngines --> ContextBuilders
    ContextBuilders --> Storage
    ContextBuilders --> VectorStore
    SearchEngines --> LLM

Local Search（局部搜索）

API 规格

函数签名：

async def local_search(
    config: GraphRagConfig,
    entities: pd.DataFrame,
    communities: pd.DataFrame,
    community_reports: pd.DataFrame,
    text_units: pd.DataFrame,
    relationships: pd.DataFrame,
    covariates: pd.DataFrame | None,
    community_level: int,
    response_type: str,
    query: str,
    callbacks: list[QueryCallbacks] | None = None,
) -> tuple[str, dict]:

参数说明：

参数	类型	必填	说明
config	GraphRagConfig	是	配置对象
entities	pd.DataFrame	是	实体表
community_reports	pd.DataFrame	是	社区报告表
text_units	pd.DataFrame	是	文本块表
relationships	pd.DataFrame	是	关系表
covariates	pd.DataFrame	否	协变量表
community_level	int	是	社区层级（0-N）
response_type	str	是	响应类型（如 “Multiple Paragraphs”）
query	str	是	用户查询

返回值：

字段	类型	说明
response	str	LLM 生成的答案文本
context_data	dict	上下文数据（包含相关实体、关系、文本单元）

核心代码

async def local_search(...):
    # 1. 数据模型转换
    entities_ = read_indexer_entities(entities, communities, community_level)
    reports = read_indexer_reports(community_reports, communities, community_level)
    text_units_ = read_indexer_text_units(text_units)
    relationships_ = read_indexer_relationships(relationships)
    
    # 2. 创建向量存储
    description_embedding_store = get_embedding_store(
        config_args=vector_store_args,
        embedding_name="entity_description_embedding",
    )
    
    # 3. 创建搜索引擎
    search_engine = get_local_search_engine(
        config=config,
        reports=reports,
        text_units=text_units_,
        entities=entities_,
        relationships=relationships_,
        covariates={"claims": covariates_},
        description_embedding_store=description_embedding_store,
        response_type=response_type,
    )
    
    # 4. 执行搜索
    result = await search_engine.search(query=query)
    
    return result.response, result.context_data

执行流程时序图

sequenceDiagram
    autonumber
    participant User
    participant API as local_search API
    participant Engine as LocalSearch Engine
    participant Context as LocalContextBuilder
    participant Vector as Vector Store
    participant LLM
    
    User->>API: local_search(query, entities, ...)
    API->>API: 数据模型转换
    API->>Engine: 创建 LocalSearch 引擎
    
    Engine->>Context: build_context(query)
    Context->>Context: 生成查询嵌入
    Context->>Vector: 向量相似度搜索 (top-k entities)
    Vector-->>Context: 匹配实体列表
    
    Context->>Context: 扩展相关关系和文本块
    Context->>Context: Token 预算分配
    Context-->>Engine: ContextResult(text, entities, ...)
    
    Engine->>Engine: 构建完整提示词
    Engine->>LLM: 调用 LLM 生成答案
    LLM-->>Engine: 答案文本
    
    Engine-->>API: SearchResult(response, context)
    API-->>User: (response, context_data)

上下文构建策略

Token 预算分配：

max_context_tokens：5000（总限制）
text_unit_prop：0.5（分配给文本单元）
community_prop：0.1（分配给社区报告）

实体选择：

向量检索 top-k 实体（默认 10）
扩展相关关系（top-k relationships）
获取实体关联的文本单元
获取实体所属社区的报告

Global Search（全局搜索）

API 规格

函数签名：

async def global_search(
    config: GraphRagConfig,
    entities: pd.DataFrame,
    communities: pd.DataFrame,
    community_reports: pd.DataFrame,
    community_level: int | None,
    dynamic_community_selection: bool,
    response_type: str,
    query: str,
) -> tuple[str, dict]:

核心代码

class GlobalSearch:
    async def search(self, query: str) -> GlobalSearchResult:
        # Step 1: 构建上下文（选择社区报告）
        context_result = await self.context_builder.build_context(
            query=query,
            **self.context_builder_params,
        )
        
        # Step 2: Map - 并行处理每个社区报告
        map_responses = await asyncio.gather(*[
            self._map_response_single_batch(
                context_data=data,
                query=query,
                max_length=self.map_max_length,
            )
            for data in context_result.context_chunks
        ])
        
        # Step 3: Reduce - 合并中间答案
        final_response = await self._reduce_responses(
            query=query,
            map_responses=map_responses,
            max_length=self.reduce_max_length,
        )
        
        return GlobalSearchResult(
            response=final_response,
            context_data=context_result.context_records,
        )

Map-Reduce 时序图

sequenceDiagram
    autonumber
    participant Engine as GlobalSearch
    participant Context as GlobalContextBuilder
    participant LLM
    
    Engine->>Context: build_context(query)
    Context->>Context: 选择社区报告（按层级或动态选择）
    Context->>Context: 分批（每批 token 不超过限制）
    Context-->>Engine: context_chunks (多个批次)
    
    par Map 阶段（并行）
        Engine->>LLM: Map Prompt + Chunk 1
        LLM-->>Engine: Intermediate Answer 1
    and
        Engine->>LLM: Map Prompt + Chunk 2
        LLM-->>Engine: Intermediate Answer 2
    and
        Engine->>LLM: Map Prompt + Chunk N
        LLM-->>Engine: Intermediate Answer N
    end
    
    Engine->>Engine: 收集所有中间答案
    Engine->>LLM: Reduce Prompt + All Answers
    LLM-->>Engine: Final Answer

动态社区选择

当 dynamic_community_selection=True 时：

从高层级社区开始
LLM 判断哪些社区与查询相关
递归展开相关社区的子社区
直到达到最大层级或 token 限制

DRIFT Search（混合搜索）

核心特性

DRIFT 结合了局部搜索的精准性和全局搜索的全面性：

局部阶段：基于实体向量检索找到相关实体
Follower Follow：根据实体所属社区，选择相关社区报告
全局 Reduce：对选中的社区报告执行 Map-Reduce

执行流程

flowchart LR
    Query[用户查询]
    Local[局部搜索: 找到 top-k 实体]
    Follow[Follow: 找到实体所属社区]
    Global[全局搜索: 对选中社区执行 Map-Reduce]
    Result[最终答案]
    
    Query --> Local
    Local --> Follow
    Follow --> Global
    Global --> Result

Basic Search（基础搜索）

核心特性

最简单的查询模式，仅基于文本单元的向量相似度：

查询嵌入生成
向量相似度搜索（top-k 文本单元）
直接将文本单元作为上下文
LLM 生成答案

配置参数对比

参数	Local	Global	DRIFT	Basic
max_context_tokens	5000	8000	5000	5000
top_k_entities	10	N/A	10	N/A
top_k_relationships	10	N/A	10	N/A
text_unit_prop	0.5	N/A	0.5	1.0
community_prop	0.1	1.0	动态	N/A
map_max_length	N/A	1000	1000	N/A
reduce_max_length	N/A	2000	2000	N/A

查询模式选择指南

查询类型	推荐模式	原因
特定实体的详细信息	Local	精准检索相关上下文
主题性宏观问题	Global	覆盖所有社区的综合视角
需要细节和概览的问题	DRIFT	结合局部精准和全局全面
简单文本检索	Basic	快速响应，无图结构依赖

性能优化

并发控制：

Local/DRIFT：单次 LLM 调用
Global：Map 阶段并发度 = concurrent_coroutines（默认 32）

缓存策略：

查询嵌入缓存
向量检索结果缓存（TTL 1小时）

流式生成：

所有查询模式支持流式输出
降低首 token 延迟

本文档详细介绍了 GraphRAG 查询引擎的四种模式、API、执行流程和最佳实践。通过合理选择查询模式和调优参数，可以在不同场景下实现最优的查询效果。

GraphRAG-02-Query-查询引擎#

模块概览#

查询引擎架构图#

Local Search（局部搜索）#

API 规格#

核心代码#

执行流程时序图#

上下文构建策略#

Global Search（全局搜索）#

API 规格#

核心代码#

Map-Reduce 时序图#

动态社区选择#

DRIFT Search（混合搜索）#

核心特性#

执行流程#

Basic Search（基础搜索）#

核心特性#

配置参数对比#

查询模式选择指南#

性能优化#

GraphRAG-02-Query-查询引擎

模块概览

查询引擎架构图

Local Search（局部搜索）

API 规格

核心代码

执行流程时序图

上下文构建策略

Global Search（全局搜索）

API 规格

核心代码

Map-Reduce 时序图

动态社区选择

DRIFT Search（混合搜索）

核心特性

执行流程

Basic Search（基础搜索）

核心特性

配置参数对比

查询模式选择指南

性能优化