链载Ai

标题: 高级 RAG 检索策略之查询路由 [打印本页]

作者: 链载Ai    时间: 昨天 11:00
标题: 高级 RAG 检索策略之查询路由


之前介绍 Self-RAG 的时候提到了其中的按需检索功能,就是根据用户的问题来判断是否需要进行文档检索,如果不需要检索的话则直接返回 LLM(大语言模型)生成的结果,这样不仅可以提升系统的性能,还可以提高用户的体验。在 Self-RAG 中按需检索是通过特殊训练后的 LLM 来实现的,但是在高级 RAG(Retrieval Augmented Generation)检索中我们可以使用查询路由来实现这个功能,借助查询路由我们可以轻松实现类似代码中的 If/Else 功能。今天我们就来介绍查询路由的原理以及实现方式,并通过代码示例来了解查询路由在实际项目中的使用。

查询路由

查询路由是 RAG 中的一种智能查询分发功能,能够根据用户输入的语义内容,从多个选项中选择最合适的处理方式或数据源。查询路由能够显著提高 RAG 检索的相关性和效率,适用于各种复杂的信息检索场景,如将用户查询分发到不同的知识库。查询路由的灵活性和智能性使其成为构建高效 RAG 系统的关键组件。



查询路由的类型

根据查询路由的实现原理我们可以将其分为两种类型:

下面我们就来了解这两种查询路由具体的实现原理。

LLM Router

使用 LLM 来判断用户的意图目前是 RAG 中一种常见的路由方法,首先在提示词中列出问题的所有类别,然后让 LLM 将问题进行分类,最后根据分类结果来选择相应的处理方式。

LLM 应用框架 LlamaIndex[1] 使用的就是 LLM Router。在 LlamaIndex 中有几种查询路由的实现,比如路由检索器 RouterRetriever、路由查询引擎 RouterQueryEngine、流水线路由模块 RouterComponent,它们的实现原理基本一致,初始化时需要一个选择器和一个工具组件列表,通过选择器来得到工具组件序号,然后根据序号来选择相应的工具组件,最后执行工具组件的处理逻辑。以 RouterQueryEngine 为例,其示例代码如下:

from llama_index.core.query_engine import RouterQueryEngine
from llama_index.core.selectors import LLMSingleSelector
from llama_index.core.tools import QueryEngineTool

# initialize tools
list_tool = QueryEngineTool.from_defaults(
query_engine=list_query_engine,
description="Useful for summarization questions related to the data source",
)
vector_tool = QueryEngineTool.from_defaults(
query_engine=vector_query_engine,
description="Useful for retrieving specific context related to the data source",
)

# initialize router query engine (single selection, llm)
query_engine = RouterQueryEngine(
selector=LLMSingleSelector.from_defaults(),
query_engine_tools=[
list_tool,
vector_tool,
],
)
query_engine.query("<query>")

下面是 LlamaIndex Router 的流程图:



在 LlamaIndex 中选择器有 4 种,如下图所示:



这 4 种选择器都是通过 LLM 来判断用户问题的意图,按选择结果可以分为单个结果选择器和多个结果选择器,单个结果选择器只返回一个选择结果,多个结果选择器返回多个选择结果,然后会将多个结果合并为一个最终结果。

按解析结果可以分为文本结果选择器和对象结果选择器,文本结果选择器使用的是 LLM 的 completion API 来生成文本类型的选择结果,格式为:<index>. <reason>index为选择结果的序号,reason为选择结果的原因,对象结果选择器使用的是 LLM 的 Function Calling API,将选择结果解析成一个 Python 对象,默认的对象为 SingleSelection,其定义如下:

class SingleSelection(BaseModel):
"""A single selection of a choice."""

index: int
reason: str

2 种解析结果示例如下所示:

# Text selector
2. Useful for questions related to oranges

# Object selector
SingleSelection(index=2, reason="Useful for questions related to oranges")

使用文本结果选择器得到选择结果后,还需要进行额外处理,比如提取出结果中的序号,而使用对象结果选择器则不需要额外处理,可以直接使用对象的属性得到结果。

我们再来看下 Selector 的提示词模板:

DEFAULT_SINGLE_SELECT_PROMPT_TMPL = (
"Some choices are given below. It is provided in a numbered list "
"(1 to {num_choices}), "
"where each item in the list corresponds to a summary.\n"
"---------------------\n"
"{context_list}"
"\n---------------------\n"
"Using only the choices above and not prior knowledge, return "
"the choice that is most relevant to the question: '{query_str}'\n"
)

使用 LLM Router 的一个关键就是构建有效的提示词,如果使用的 LLM 足够强大,那么提示词不用很清晰也能达到我们想要的效果,但如果 LLM 不够强大,那么提示词需要不断调整才能得到满意的结果。笔者在使用 LlamaIndex Router 的过程中发现,在选择 OpenAI gpt-3.5-turbo 模型的情况下,使用 LLMSingleSelector 选择器时偶尔会出现解析失败的情况,而使用 PydanticSingleSelector 选择器则比较稳定。

最后得到选择结果的序号后就可以通过该序号来选择工具组件了,下面是 RouterQueryEngine 的代码片段:

class RouterQueryEngine(BaseQueryEngine):
def _query(self, query_bundle: QueryBundle) -> RESPONSE_TYPE:
......
result = self._selector.select(self._metadatas, query_bundle)
selected_query_engine = self._query_engines[result.ind]
final_response = selected_query_engine.query(query_bundle)
......

优缺点

Embedding Router

查询路由的另外一种实现方式是使用 Embedding 模型将用户问题进行向量化,然后通过向量相似性来将用户问题进行分类,得到分类结果后再选择相应的处理方式。

Semantic Router[2] 是基于该原理实现的一个路由工具,它旨在提供超快的 AI 决策能力,通过语义向量进行快速决策,以提高 LLM 应用和 AI Agent 的效率。Semantic Router 使用非常简单,示例代码如下:

import os
from semantic_router import Route
from semantic_router.encoders import CohereEncoder, OpenAIEncoder
from semantic_router.layer import RouteLayer

# we could use this as a guide for our chatbot to avoid political conversations
politics = Route(
name="politics",
utterances=[
"isn't politics the best thing ever",
"why don't you tell me about your political opinions",
"don't you just love the president",
"they're going to destroy this country!",
"they will save the country!",
],
)

# this could be used as an indicator to our chatbot to switch to a more
# conversational prompt
chitchat = Route(
name="chitchat",
utterances=[
"how's the weather today?",
"how are things going?",
"lovely weather today",
"the weather is horrendous",
"let's go to the chippy",
],
)

# we place both of our decisions together into single list
routes = [politics, chitchat]

# OpenAI Encoder
os.environ["OPENAI_API_KEY"] = "<YOUR_API_KEY>"
encoder = OpenAIEncoder()

rl = RouteLayer(encoder=encoder, routes=routes)

rl("don't you love politics?").name
# politics
rl("how's the weather today?").name
# chitchat

OpenAI Encoder 默认使用的是 text-embedding-3-small Embedding 模型,它比 OpenAI 之前的 text-embedding-ada-002 Embedding 模型效果更好且价格更便宜。同时 Semantic Router 还支持其他 Encoder,比如 Huggingface Encoder,它默认使用的是 sentence-transformers/all-MiniLM-L6-v2[3] Embedding 模型,这是一个句子转换模型,它将句子和段落映射到一个 384 维度的向量空间,可用于分类或语义搜索等任务。

优缺点

查询路由实践

下面我们结合 LlamaIndex 和 Semantic Router 来实现一个查询路由,该路由会将用户的问题分发到不同的工具组件中,这些工具组件包括:使用 LLM 和用户进行闲聊,使用 RAG 流程检索文档并生成答案,以及使用 Bing 搜索引擎进行网络搜索。



首先我们定义一个与 LLM 闲聊的工具组件,这里我们使用 LlamaIndex 的 Pipeline[4] 功能来构建一个查询流水线,更多的查询流水线功能可以参考我之前的这篇文章,示例代码如下:

from llama_index.llms.openai import OpenAI
from llama_index.core.query_pipeline import QueryPipeline, InputComponent

llm = OpenAI(model="gpt-3.5-turbo", system_prompt="You are a helpful assistant.")
chitchat_p = QueryPipeline(verbose=True)
chitchat_p.add_modules(
{
"input": InputComponent(),
"llm": llm,
}
)
chitchat_p.add_link("input", "llm")
output = chitchat_p.run(input="hello")
print(f"Output: {output}")

# 显示结果
Output: assistant: Hello! How can I assist you today?

然后我们再添加一个普通 RAG 的工具组件,同样是创建一个查询流水线,这里的测试文档还是用维基百科上的复仇者联盟[6]电影剧情,示例代码如下:

from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
from llama_index.core.response_synthesizers.tree_summarize import TreeSummarize

documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(documents)
retriever = index.as_retriever(similarity_top_k=2)
rag_p = QueryPipeline(verbose=True)
rag_p.add_modules(
{
"input": InputComponent(),
"retriever": retriever,
"output": TreeSummarize(),
}
)

rag_p.add_link("input", "retriever")
rag_p.add_link("input", "output", dest_key="query_str")
rag_p.add_link("retriever", "output", dest_key="nodes")
output = rag_p.run(input="Which two members of the Avengers created Ultron?")
print(f"Output: {output}")

# 显示结果
Output: Tony Stark and Bruce Banner.

接下来我们再添加一个使用 Bing 搜索引擎的工具组件,同样我们使用查询流水线来进行创建,但这一次需要用到自定义模块,示例代码如下:

web_p = QueryPipeline(verbose=True)
web_p.add_modules(
{
"input": InputComponent(),
"web_search": WebSearchComponent(),
}
)
web_p.add_link("input", "web_search")

在实现这个自定义模块之前,我们需要先在 Azure 上创建一个 Bing 搜索服务,然后获取 API Key,具体操作可以参考微软的官方文档[7]。然后安装 LlamaIndex 的 Bing 查询工具库:pip install llama-index-tools-bing-search,然后就可以开始实现自定义模块了,示例代码如下:

import os
from typing import Dict, Any
from llama_index.core.query_pipeline import CustomQueryComponent
from llama_index.tools.bing_search import BingSearchToolSpec
from llama_index.agent.openai import OpenAIAgent

class WebSearchComponent(CustomQueryComponent):
"""Web search component."""

def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
"""Validate component inputs during run_component."""
assert "input" in input, "input is required"
return input

@property
def _input_keys(self) -> set:
"""Input keys dict."""
return {"input"}

@property
def _output_keys(self) -> set:
return {"output"}

def _run_component(self, **kwargs) -> Dict[str, Any]:
"""Run the component."""
tool_spec = BingSearchToolSpec(api_key=os.getenv("BING_SEARCH_API_KEY"))
agent = OpenAIAgent.from_tools(tool_spec.to_tool_list())
question = kwargs["input"]
result = agent.chat(question)
return {"output": result}

3 个工具组件创建之后,我们需要创建一个路由模块,我们使用 Semantic Router 来实现这个路由模块,我们先定义 Semantic Router 的几个 Route,示例代码如下:

chitchat = Route(
name="chitchat",
utterances=[
"how's the weather today?",
"how are things going?",
"lovely weather today",
"the weather is horrendous",
"let's go to the chippy",
],
)

rag = Route(
name="rag",
utterances=[
"What mysterious object did Loki use in his attempt to conquer Earth?",
"Which two members of the Avengers created Ultron?",
"How did Thanos achieve his plan of exterminating half of all life in the universe?",
"What method did the Avengers use to reverse Thanos' actions?",
"Which member of the Avengers sacrificed themselves to defeat Thanos?",
],
)

web = Route(
name="web",
utterances=[
"Search online for the top three countries in the 2024 Paris Olympics medal table.",
"Find the latest news about the U.S. presidential election.",
"Look up the current updates on NVIDIA’s stock performance today.",
"Search for what Musk said on X last month.",
"Find the latest AI news.",
],
)

接下来我们创建一个自定义的路由模块,使用 Semantic Router 来实现查询路由,示例代码如下:

from llama_index.core.base.query_pipeline.query import (
QueryComponent,
QUERY_COMPONENT_TYPE,
)
from llama_index.core.bridge.pydantic import Field

class SemanticRouterComponent(CustomQueryComponent):
"""Semantic router component."""

components: Dict[str, QueryComponent] = Field(
..., description="Components (must correspond to choices)"
)

def __init__(self, components: Dict[str, QUERY_COMPONENT_TYPE]) -> None:
"""Init."""
super().__init__(components=components)

def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
"""Validate component inputs during run_component."""
return input

@property
def _input_keys(self) -> set:
"""Input keys dict."""
return {"input"}

@property
def _output_keys(self) -> set:
return {"output", "selection"}

def _run_component(self, **kwargs) -> Dict[str, Any]:
"""Run the component."""
if len(self.components) < 1:
raise ValueError("No components")
if chitchat.name not in self.components.keys():
raise ValueError("No chitchat component")

routes = [chitchat, rag, web]
encoder = OpenAIEncoder()
rl = RouteLayer(encoder=encoder, routes=routes)
question = kwargs["input"]
selection = rl(question).name
if selection is not None:
output = self.components[selection].run_component(input=question)
else:
output = self.components["chitchat"].run_component(input=question)
return {"output": output, "selection": selection}

最后我们将所有的工具组件和路由模块添加到一个单独的查询流水线中,示例代码如下:

p = QueryPipeline(verbose=True)
p.add_modules(
{
"router": SemanticRouterComponent(
components={
"chitchat": chitchat_p,
"rag": rag_p,
"web": web_p,
}
),
}
)

下面我们来执行一下这个流水线,看看效果如何:

output = p.run(input="hello")
# Selection: chitchat
# Output: assistant: Hello! How can I assist you today?

output = p.run(input="Which two members of the Avengers created Ultron?")
# Selection: rag
# Output: Tony Stark and Bruce Banner.

output = p.run(input="Search online for the top three countries in the 2024 Paris Olympics medal table.")
# Selection: web
# Output: The top three countries in the latest medal table for the 2024 Paris Olympics are as follows:
# 1. United States
# 2. China
# 3. Great Britain

可以看到我们的查询路由工作的很好,根据用户问题的不同意图选择了不同的工具组件,并得到了相应的结果。

总结

今天我们介绍了 RAG 检索策略中的查询路由,并介绍了 LLM Router 和 Embedding Router 两种查询路由的实现原理,最后通过一个实战项目了解了查询路由在实际项目中的使用。但目前的查询路由还有很多不确定性,因此我们无法保证查询路由总能做出完全准确的决策,需要经过精心测试才能得到更加可靠的 RAG 应用程序。







欢迎光临 链载Ai (https://www.lianzai.com/) Powered by Discuz! X3.5