使用 MultiQueryRetriever 提升向量数据库检索的效果

最新推荐文章于 2025-10-03 19:59:37 发布

原创最新推荐文章于 2025-10-03 19:59:37 发布 · 464 阅读

CC 4.0 BY-SA版权

文章标签：

在 AI 驱动的数据检索领域，使用向量数据库进行基于距离的检索是一个常见的方法。该方法通过在高维空间中嵌入查询，并基于距离度量找到相似的嵌入文档。然而，微小的查询措辞变化可能导致不同的结果，同时如果嵌入无法很好地捕捉数据的语义，这也可能影响检索效果。通常我们会使用提示工程来手动解决这些问题，但这个过程可能相当繁琐。

MultiQueryRetriever 通过使用大型语言模型（LLM）为给定的用户输入查询生成多个不同视角的查询来自动化提示调整过程。对于每个生成的查询，它会检索一组相关文档，并在所有查询之间取唯一的并集，以获得更大且潜在相关的文档集合。通过对同一问题生成多个视角，MultiQueryRetriever 能够缓解基于距离检索的一些局限性，并获得更丰富的结果。

构建示例向量数据库

我们将使用 LangChain 库来构建示例向量数据库，导入相关模块：

from langchain_chroma import Chroma
from langchain_community.document_loaders import WebBaseLoader
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 加载博客文章
loader = WebBaseLoader("https://lilianweng.github.io/posts/2023-06-23-agent/")
data = loader.load()

# 文本拆分
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=0)
splits = text_splitter.split_documents(data)

# 向量数据库
embedding = OpenAIEmbeddings()
vectordb = Chroma.from_documents(documents=splits, embedding=embedding)

# 提供日志信息以跟踪生成的查询
import logging
logging.basicConfig()
logging.getLogger("langchain.retrievers.multi_query").setLevel(logging.INFO)

# 生成示例查询
from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain_openai import ChatOpenAI

question = "What are the approaches to Task Decomposition?"
llm = ChatOpenAI(temperature=0)
retriever_from_llm = MultiQueryRetriever.from_llm(
    retriever=vectordb.as_retriever(), llm=llm
)

# 获取文档
unique_docs = retriever_from_llm.invoke(question)
print(len(unique_docs))

此代码会从指定的博客文章加载数据，将文章内容分割为多个片段，然后将这些片段存储到一个向量数据库中。接着，它会设置一个多查询检索工具来生成不同视角的查询，并返回相关文档。

提供自定义提示模板

MultiQueryRetriever 允许用户定制生成查询的提示模板。我们可以创建一个 PromptTemplate 并配套一个输出解析器来将结果分割为多个查询列表：

from typing import List
from langchain_core.output_parsers import BaseOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field

# 输出解析器
class LineListOutputParser(BaseOutputParser[List[str]]):
    """将 LLM 的输出分割为查询列表。"""
    def parse(self, text: str) -> List[str]:
        lines = text.strip().split("\n")
        return list(filter(None, lines))  # 移除空行

output_parser = LineListOutputParser()

QUERY_PROMPT = PromptTemplate(
    input_variables=["question"],
    template="""你是一个 AI 语言模型助手。你的任务是为给定的用户问题生成五种不同版本以从向量数据库中检索相关文档。通过为问题提供多个视角，你的目标是帮助用户克服基于距离相似性搜索的一些限制。将这些不同的查询通过换行分隔。
    原始问题: {question}""",
)
llm = ChatOpenAI(temperature=0)

# 链设置
llm_chain = QUERY_PROMPT | llm | output_parser

# 执行检索
retriever = MultiQueryRetriever(
    retriever=vectordb.as_retriever(), llm_chain=llm_chain, parser_key="lines"
)

# 获取文档
unique_docs = retriever.invoke("What does the course say about regression?")
print(len(unique_docs))