4.8 基于两阶段索引的 RAG Agent 搭建
在上一节内容中,我们介绍了基于 Qwen3 Reranking 模型的原理及使用方法,并且对 Qwen3 Reranking 模型的实现代码也进行了简单的介绍。在本节内容中,我们将以第4.4节中介绍的 RAG Agent 为基础,来搭建一个包含有两阶段索引的 RAG Agent 。
4.8.1 两阶段索引实现
因为整个 RAG 框架我们已经在第4.4节内容中介绍过了,所以这里只需要再实现两阶段索引部分的流程,并进行简单改造即可。具体地,对于重排序部分的逻辑实现如下:
1 def text_rerank(query: str = None, docs: List[Document] = None, top_n=3):
2 documents = [doc.page_content for doc in docs]
3 resp = dashscope.TextReRank.call(model="qwen3-rerank", query=query,
4 documents=documents, top_n=top_n, return_documents=False,
5 instruct="Given a web search query, retrieve relevant passages that answer the query.")
6 if resp.status_code == HTTPStatus.OK:
7 results = resp.output.results
8 filtered_info = [[item.index, item.relevance_score] for item in results]
9 filtered_docs = [docs[item[0]] for item in filtered_info]
10 return filtered_docs, filtered_info
11 else:
12 return docs[:top_n], []
在上述代码中,第1行代码 docs 表示对于当前 query 来说从向量库中检索到的与之相关的候选文档,为一个列表。第2行取每个 Document 对象中的 page_content ,即文本内容。第3~5行是得到排序后的结果,并取前 top_n 个,详细内容参见第4.7.2节。第6~9行是过滤无用信息,仅返回前 top_n 个Document 对象以及对应的索引和相似性评分。
4.8.2 检索工具改造
进一步,需要对检索工具 retrieve_context 进行一定的改造,使得其能够对从向量库检索到的内容进行二次排序。具体地,实现过程如下:
1 @tool
2 def retrieve_context(query: str, config: RunnableConfig):
3 vector_store = config['configurable'].get("vector_store")
4 k = config['configurable'].get("k")
5 top_n = config['configurable'].get("top_n")
6 artifact = vector_store.similarity_search(query, k=k)
7 filtered_doc, filter_info = text_rerank(query, artifact, top_n=top_n)
8 print(f"从向量数据库检索到的{k}篇文档中,相关度前{top_n}个文档的来源为:")
9 for i in range(top_n):
10 print(f"\t{artifact[i].metadata}")
11 print(f"\n经过排序模型处理后,相关度前{top_n}个文档的来源为:")
12 for i in range(top_n):
13 print(f"\t{filtered_doc[i].metadata}")
14 if len(filtered_doc) == 0:
15 print(f"text_rerank 执行失败!")
16 print(f"一共从向量库检索到{len(artifact)}篇文章,经过reranking模型处理后,返回{len(filtered_doc)}"
17 f"篇最相似的文档,相似性评分分别为: {filter_info}")
18 content = "\n\n".join(
19 (f"Source: {doc.metadata}\nContent: {doc.page_content}")
20 for doc in filtered_doc)
21 return content
在上述代码中,第7行则是重排序后的结果。第8~17行是作为验证的相关输出信息。第18~20行则是处理返回输入到模型中的参考内容,与之前一致。
4.8.3 输出结果分析
此时,便可以运行上述代码得到类似如下结果:
# ================================ Human Message =================================
# 郭靖这个人物在这部小说当中是谁?它的出生背景是什么?
# ================================== Ai Message ==================================
# Tool Calls:
# retrieve_context (call_a80494c707914126a3b383)
# Call ID: call_a80494c707914126a3b383
# Args:
# query: 郭靖 出生背景 小说人物
# 从向量数据库检索到的20篇文档中,相关度前3个文档的来源为:
# {'level_1': '第七回 比武招亲', 'pk': 465261515437506672, 'source': 'RAGWithMe/data/jinyong/金庸-射雕英雄传精校版.txt', 'start_index': 5796}
# {'level_1': '第七回 比武招亲', 'pk': 465261515437506675, 'source': 'RAGWithMe/data/jinyong/金庸-射雕英雄传精校版.txt', 'start_index': 7025}
# {'level_1': '第三十九回 是非善恶', 'pk': 465290802323521722, 'source': 'RAGWithMe/data/jinyong/金庸-射雕英雄传精校版.txt', 'start_index': 0}
# 经过排序模型处理后,相关度前3个文档的来源为:
# {'level_1': '第三回 黄沙莽莽', 'pk': 464859367376683180, 'source': 'RAGWithMe/data/jinyong/金庸-射雕英雄传精校版.txt', 'start_index': 7398}
# {'level_1': '第七回 比武招亲', 'pk': 465261515437506672, 'source': 'RAGWithMe/data/jinyong/金庸-射雕英雄传精校版.txt', 'start_index': 5796}
# {'level_1': '第三十八回 锦囊密令', 'pk': 465290802323521715, 'source': 'RAGWithMe/data/jinyong/金庸-射雕英雄传精校版.txt', 'start_index': 15090}
# 一共从向量库检索到20篇文章,经过 reranking 模型处理后,返回3篇最相似的文档,相似性评分分别为: [[0, 0.8125869388341607], [17, 0.8091164423628996], [5, 0.7374832867736165]]
# ================================== Ai Message ==================================
# 郭靖是《射雕英雄传》中的主人公。他的母亲李萍是浙江临安人,父亲是山东好汉。郭靖出生于临安府牛家村,幼年时因战乱随母亲流落蒙古大漠。母亲李萍含辛茹苦抚养他,在大漠中以畜牧为生,并依丈夫遗言为其取名"郭靖"。他四岁才会说话,性格憨厚朴实,六岁时已能在草原上放牧牛羊。母子以临安乡音对话,但均通晓蒙古话。丘处机道长曾于其出生时赠短剑,并为其取名"郭靖",寓意不忘"靖康之耻"。
根据上面的输出结果可以看出,