您的位置：首页 > > 教程攻略 > ai资讯 >轻松上手GraphRAG源码，手把手教你怎样给GraphRAG增加流式输出

轻松上手GraphRAG源码，手把手教你怎样给GraphRAG增加流式输出

来源:互联网 更新时间:2026-08-26 21:41

在LLM项目开发中，第一个遇到的需求往往是流式输出——这个功能能大幅缩短用户的等待时间，体验提升立竿见影。但说实话，目前GraphRAG还是个演示级项目，还没支持流式输出。所以，我们不妨一起来给它加上这个能力。目标不光是新功能本身，更关键的是通过动手过程，把GraphRAG的源码吃透。

命令行添加streaming参数

GraphRAG的执行入口在 graphrag/query/__main__.py，我们需要先给它增加一个 --streaming 参数。代码如下：

parser.add_argument(
    "--streaming",
    help="Whether to output the response in a streaming format",
    action="store_true",
)

很简单，就是在命令行解析器里加一个选项。当你在命令行里带上 --streaming 时，该参数值就为 True。这样一来，程序在处理查询时就会逐步输出结果，而不是等到全部算完再一次性给你。等到所有代码改完，执行下面的命令就能体验流式输出效果：

poetry run poe query --root ./ragtest --streaming --method global '路飞有哪些伙伴？'

run_local_search 与 run_global_search 增加 streaming 参数

在 GraphRAG 源码中，Local Query 和 Global Query 实际调用的是 graphrag/query/cli.py 里的 run_local_search 和 run_global_search。这两个函数需要加上 streaming 入参：

def run_local_search(
    config_dir: str | None,
    data_dir: str | None,
    root_dir: str | None,
    community_level: int,
    response_type: str,
    streaming: bool,
    query: str,
):
    ...

def run_global_search(
    config_dir: str | None,
    data_dir: str | None,
    root_dir: str | None,
    community_level: int,
    response_type: str,
    streaming: bool,
    query: str,
):
    ...

接着在 graphrag/query/__main__.py 里把命令行参数传给这两个函数：

...
match args.method:
    case SearchType.LOCAL:
        run_local_search(
            args.config,
            args.data,
            args.root,
            args.community_level,
            args.response_type,
            args.streaming,
            args.query[0],
        )
    case SearchType.GLOBAL:
        run_global_search(
            args.config,
            args.data,
            args.root,
            args.community_level,
            args.response_type,
            args.streaming,
            args.query[0],
        )
    case _:
        raise ValueError(INVALID_METHOD_ERROR)

这两个函数的执行逻辑大体分三步： 1. 根据知识图谱的各种 parquet 文件，准备好 search 所需的数据。 2. 用这些数据构建一个查询引擎对象（称为 search_engine）。 3. 调用 search_engine 的 search 或 asearch 方法执行实际查询。其中第一步和第二步与是否流式输出无关——无论用户选不选流式，数据组织和引擎构建都一样。真正需要修改的是第三步：

result = await search_engine.asearch(query=query)
reporter.success(f"Local Search Response: {result.response}")
return result.response

如果用户选择了流式输出，就需要让这两个方法返回一个可迭代对象，逐步推送结果；同时保证没选流式时，仍然按原方式直接返回结果。为此，我们新增了一个 astream_search 方法。下面以 local_search 为例详细说明流式功能的实现，global_search 的改造完全同理——两者唯一的区别是使用的搜索引擎对象不同（GlobalSearch vs LocalSearch）。

本次修改基于 GraphRAG v0.3.0

local_search

在最新代码中，run_local_search 调用了 graphrag/query/api.py 里的 local_search 方法，这个方法封装了构建引擎和执行查询的步骤。为了实现流式输出，需要对它做一点改动：当用户选择流式时，返回一个可迭代对象；否则照常返回结果。

search_engine = get_local_search_engine(
    config=config,
    reports=read_indexer_reports(community_reports, nodes, community_level),
    text_units=read_indexer_text_units(text_units),
    entities=_entities,
    relationships=read_indexer_relationships(relationships),
    covariates={"claims": _covariates},
    description_embedding_store=description_embedding_store,
    response_type=response_type,
)

if not streaming:
    result = await search_engine.asearch(query=query)
    reporter.success(f"Global Search Response: {result.response}")
    return result.response

else:
    import sys
    full_resp = ''
    results = search_engine.astream_search(query=query)
    reporter.success(f'Global Search Response: \n')

    async for result in results:
        sys.stdout.write(result)
        sys.stdout.flush()
        full_resp += result

    sys.stdout.write('\b\n')
    return full_resp

astream_search

新加的 astream_search 方法和原有的 asearch 功能相似，关键区别在于如何调用 LLM。asearch 是等所有结果全部生成后才返回，而 astream_search 会在 LLM 处理完每一部分后立刻返回当前结果——本质上是调用了 LLM 的流式响应能力。

async def astream_search(
    self,
    query: str,
    conversation_history: ConversationHistory | None = None,
    **kwargs,
) -> SearchResult:
    """Build local search context that fits a single context window and generate answer for the user query."""
    start_time = time.time()
    search_prompt = ""

    context_text, context_records = self.context_builder.build_context(
        query=query,
        conversation_history=conversation_history,
        **kwargs,
        **self.context_builder_params,
    )
    log.info("GENERATE ANSWER: %s. QUERY: %s", start_time, query)
    try:
        search_prompt = self.system_prompt.format(
            context_data=context_text, response_type=self.response_type
        )
        search_messages = [
            {"role": "system", "content": search_prompt},
            {"role": "user", "content": query},
        ]

        response = await self.llm.astream_generate(
            messages=search_messages,
            callbacks=self.callbacks,
            **self.llm_params,
        )

        return SearchResult(
            response=response,
            context_data=context_records,
            context_text=context_text,
            completion_time=time.time() - start_time,
            llm_calls=1,
            prompt_tokens=num_tokens(search_prompt, self.token_encoder),
        )

    except Exception:
        log.exception("Exception in _asearch")
        return SearchResult(
            response="",
            context_data=context_records,
            context_text=context_text,
            completion_time=time.time() - start_time,
            llm_calls=1,
            prompt_tokens=num_tokens(search_prompt, self.token_encoder),
        )

astream_generate

流式查询的最后一环：在 graphrag/query/llm/oai/chat_openai.py 的 ChatOpenAI 类中添加 astream_generate 方法。因为 search_engine.asearch 底层调的是 self.llm.agenerate，所以要支持流式就需要在 LLM 层也增加一个流式版本。

async def astream_generate(
    self,
    messages: str | list[Any],
    callbacks: list[BaseLLMCallback] | None = None,
    **kwargs: Any,
) -> AsyncGenerator[str, None] | None:
    """Generate text asynchronously with streaming."""
    try:
        retryer = AsyncRetrying(
            stop=stop_after_attempt(self.max_retries),
            wait=wait_exponential_jitter(max=10),
            reraise=True,
            retry=retry_if_exception_type(self.retry_error_types),
        )
        async for attempt in retryer:
            with attempt:
                return self._astream_generate(
                    messages=messages,
                    callbacks=callbacks,
                    **kwargs,
                )
            except RetryError as e:
                self._reporter.error(f"Error at astream_generate(): {e}")
                return
            else:
                return
        
async def _astream_generate(
        self,
        messages: str | list[Any],
        callbacks: list[BaseLLMCallback] | None = None,
        **kwargs: Any,
    ) -> AsyncGenerator[str, None]:
        model = self.model
        if not model:
            raise ValueError(_MODEL_REQUIRED_MSG)
        response = await self.async_client.chat.completions.create(
            model=model,
            messages=messages,
            stream=True,
            **kwargs,
        )
        async for chunk in response:
            if not chunk or not chunk.choices:
                continue

            delta = (
                chunk.choices[0].delta.content
                if chunk.choices[0].delta and chunk.choices[0].delta.content
                else ""
            )

            yield delta

            if callbacks:
                for callback in callbacks:
                    callback.on_llm_new_token(delta)

测试

代码改完后，跑一下流式查询，效果流畅，完美。

poetry run poe query --root ./ragtest --streaming --method local '路飞有哪些伙伴？'

总结

整个流程走下来，我们只添加了为数不多的几段代码，就让 GraphRAG 支持了流式输出。但说实话，新功能本身不是目的——借这个机会深入理解 GraphRAG 的源码架构，才是更有价值的事。就像之前文章里提到的，GraphRAG 目前仍处于原型阶段，离生产环境还有一段路要走。不过它构建知识图谱的思路很值得借鉴。后续文章会考虑把 LangChain 或 LlamaIndex 与 GraphRAG 结合，在 GraphRAG 生成的知识图谱上利用它们进行查询，敬请期待。