在我的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的源码。
就像我在 最近爆火的GraphRAG是什么，真的能用于商业应用吗？说的那样，GraphRAG目前仍处于原型阶段，距离能够投入生产使用还有一段距离。但是，它构建知识图谱的方法值得我们借鉴学习。
来源：AI 博物院
