微信扫码
与创始人交个朋友
我要投稿
在我的LLM项目开发过程中,第一个遇到的需求就是实现流式输出。这个功能可以减少用户的等待时间,从而提升用户体验。不过,目前作为一个演示项目,GraphRAG还没有支持流式输出。
因此,我将带领大家实现这个流式输出功能。我们的目标不仅是添加新功能,更重要的是通过这个过程,加深你对GraphRAG源码的理解。
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 '路飞有哪些伙伴?'
在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函数的执行逻辑主要包括以下几个步骤:
在这三个步骤中,第一步和第二步负责数据的准备和查询对象的构建,并不会受到我们是否引入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
在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
的方法。这个方法与原有的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),
)
为了实现流式查询,我们需要在代码的多个部分进行修改。具体来说,在源码中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目前仍处于原型阶段,距离能够投入生产使用还有一段距离。但是,它构建知识图谱的方法值得我们借鉴学习。在下一篇文章中,我会把Langchain或者llamaindex和GraphRAG结合起来,在GraphRAG生成的知识图谱上使用Langchain或llamaindex进行查询,欢迎点赞关注,获取实时更新
53AI,企业落地应用大模型首选服务商
产品:大模型应用平台+智能体定制开发+落地咨询服务
承诺:先做场景POC验证,看到效果再签署服务协议。零风险落地应用大模型,已交付160+中大型企业
2024-10-06
Anthropic RAG: 上下文检索技术
2024-10-06
AI检索黑马获2000万美元融资,推进RAG系统精准化,破解AI幻觉难题
2024-10-06
OpenRAG:全面增强RAG推理,超越Self-RAG、RAG 2.0、Command R+
2024-10-06
如何引入chunk层级优化RAG性能:从树结构RAPTOR到GEM-RAG加权图方案解读
2024-10-06
KAG:RAG已经不够了,知识增强生成才是王道,提升朴素RAG一倍性能
2024-10-05
打造自己的RAG解析大模型:Windows下部署OCR应用服务(可商业应用)
2024-10-05
基于RAG的to B智能体(Agent)应用实践
2024-10-03
RAG实战篇:将用户输入转换为精确的数据库查询语言
2024-07-18
2024-07-09
2024-07-09
2024-07-08
2024-07-09
2024-06-20
2024-07-07
2024-05-05
2024-07-07
2024-05-19
2024-09-30
2024-09-26
2024-09-26
2024-09-20
2024-09-16
2024-09-12
2024-09-11
2024-09-10