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

沈页

于 2024-09-29 08:30:00 发布

阅读量222

点赞数 9

文章标签：人工智能 AI大模型 ai 大模型 LLM RAG GraphRAG

本文链接：https://blog.csdn.net/Androiddddd/article/details/142620061

版权

在我的LLM项目开发过程中，第一个遇到的需求就是实现流式输出。这个功能可以减少用户的等待时间，从而提升用户体验。不过，目前作为一个演示项目，GraphRAG还没有支持流式输出。

因此，我将带领大家实现这个流式输出功能。我们的目标不仅是添加新功能，更重要的是通过这个过程，加深你对GraphRAG源码的理解。

一、命令行添加streaming参数

GraphRAG在graphrag/query/main.py文件中包含了执行Query的命令行入口代码，为了提升它的功能并且让其可以进行流式输出，我们计划增加一个名为"streaming"的参数。

这个新增的参数可以通过下面的代码段添加：

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

此代码段的作用是在命令行解析器（parser）中添加一个新的选项"–streaming"。当这个选项被设置时（即在命令行中加上–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。现在我们需要为这两个函数添加入参：

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的__main__ 中把streaming这个命令行参数传给这两个函数:

...
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)

在GraphRAG的框架中，run_local_search和run_global_search函数的执行逻辑主要包括以下几个步骤：

根据构建的知识图谱各种parquet文件，组织执行search函数所需要的数据。
利用这些数据，构建一个执行查询的对象，我们将它称为search_engine。
调用search_engine对象的search或asearch方法来执行实际的查询操作。

在这三个步骤中，第一步和第二步负责数据的准备和查询对象的构建，并不会受到我们是否引入streaming参数的影响。换句话说，无论用户是否选择了流式输出，知识图谱数据的组织和search_engine对象的构建都是相同的。

因此，我们真正需要修改的部分是第三步：调用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函数的优化也将遵循相同的方法，因为global_search和local_search函数本质上是非常相似的。它们之间唯一的区别在于所使用的搜索引擎对象不同：global_search使用的是GlobalSearch引擎，而local_search使用的是LocalSearch引擎。

本次修改的GraphRAG的版本是v0.3.0

三、local_search

在GraphRAG的最新代码版本中，run_local_search函数调用了位于graphrag/query/api.py文件中的local_search方法。此方法封装了我们之前所提的第二步（构建查询引擎对象）和第三步（执行查询操作）的逻辑。

为了实现流式输出功能，我们需要对这个local_search方法进行一些修改。具体的改动将包括在用户选择了流式输出时使其返回一个可迭代对象，从而实现逐步推送查询结果的功能。同时，我们也需要保证当用户没有选择流式输出时，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方法时，所有的查询结果会在一次性完成后返回给用户。换句话说，LLM会处理完所有的查询请求后才将结果返回。但当我们使用astream_search方法时，LLM将会在处理每一部分查询请求后就立即返回当前的结果。也就是说，astream_search调用了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

为了实现流式查询，我们需要在代码的多个部分进行修改。具体来说，在源码中search_engine.asearch实际上是调用了self.llm.agenerate方法。因此，为了支持流式查询，我们也需要在相应的位置添加一个名为astream_generate的新方法。

这个新方法将被添加到位于graphrag/query/llm/oai/chat_openai.py文件中的ChatOpenAI类中。添加astream_generate的目的是为了让ChatOpenAI类能够生成流式响应。

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),  # type: ignore
        )
        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(  # type: ignore
            model=model,
            messages=messages,  # type: ignore
            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 ""
            )  # type: ignore

            yield delta

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

六、测试

经过上面代码的新增和修改之后，下面我们就可以来测试streaming功能了，运行很流畅，pefect

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

七、总结

我们已经成功地实现了streaming功能，并且在这个过程中，只添加了为数不多的一些新的代码。然而，实现这个功能不是我们真正想要达到的目标，更重要的是通过这个过程深入理解和掌握GraphRAG的源码。

八、如何学习大模型？

学习AI大模型是一个系统的过程，需要从基础开始，逐步深入到更高级的技术。

这里给大家精心整理了一份全面的AI大模型学习资源，包括：AI大模型全套学习路线图（从入门到实战）、精品AI大模型学习书籍手册、视频教程、实战学习、面试题等，资料免费分享！

1. 成长路线图&学习规划

要学习一门新的技术，作为新手一定要先学习成长路线图，方向不对，努力白费。

这里，我们为新手和想要进一步提升的专业人士准备了一份详细的学习成长路线图和规划。可以说是最科学最系统的学习成长路线。
在这里插入图片描述

2. 大模型经典PDF书籍

书籍和学习文档资料是学习大模型过程中必不可少的，我们精选了一系列深入探讨大模型技术的书籍和学习文档，它们由领域内的顶尖专家撰写，内容全面、深入、详尽，为你学习大模型提供坚实的理论基础。（书籍含电子版PDF）

在这里插入图片描述

3. 大模型视频教程

对于很多自学或者没有基础的同学来说，书籍这些纯文字类的学习教材会觉得比较晦涩难以理解，因此，我们提供了丰富的大模型视频教程，以动态、形象的方式展示技术概念，帮助你更快、更轻松地掌握核心知识。

在这里插入图片描述

4. 大模型项目实战

学以致用 ，当你的理论知识积累到一定程度，就需要通过项目实战，在实际操作中检验和巩固你所学到的知识，同时为你找工作和职业发展打下坚实的基础。

在这里插入图片描述

5. 大模型面试题

面试不仅是技术的较量，更需要充分的准备。

在你已经掌握了大模型技术之后，就需要开始准备面试，我们将提供精心整理的大模型面试题库，涵盖当前面试中可能遇到的各种技术问题，让你在面试中游刃有余。

在这里插入图片描述

全套的AI大模型学习资源已经整理打包，有需要的小伙伴可以微信扫描下方CSDN官方认证二维码，免费领取【保证100%免费】

沈页

关注

9
点赞
踩
17

收藏

觉得还不错? 一键收藏
0
评论
复制链接

分享到 QQ

分享到新浪微博

扫一扫