Skip to main content
Open In ColabOpen on GitHub

如何创建一个自定义的检索器

概述

许多LLM应用程序涉及使用Retriever从外部数据源检索信息。

检索器负责检索与给定用户query相关的一系列Documents

检索到的文档通常被格式化为提示,这些提示被输入到LLM中,使LLM能够使用其中的信息生成适当的响应(例如,基于知识库回答用户问题)。

接口

要创建你自己的检索器,你需要扩展BaseRetriever类并实现以下方法:

方法描述必选/可选
_get_relevant_documents获取与查询相关的文档。Required
_aget_relevant_documents实现以提供异步原生支持。Optional

_get_relevant_documents 内部的逻辑可能涉及使用请求对数据库或网络进行任意调用。

tip

通过继承BaseRetriever,你的检索器自动成为LangChain的Runnable,并且将立即获得标准的Runnable功能!

info

你可以使用RunnableLambdaRunnableGenerator来实现一个检索器。

将检索器实现为BaseRetriever而不是RunnableLambda(一个自定义的可运行函数)的主要好处是,BaseRetriever是一个众所周知的LangChain实体,因此一些用于监控的工具可能会为检索器实现专门的行为。另一个区别是,BaseRetriever在某些API中的行为与RunnableLambda略有不同;例如,astream_events API中的start事件将是on_retriever_start而不是on_chain_start

示例

让我们实现一个玩具检索器,它返回所有文本中包含用户查询文本的文档。

from typing import List

from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever


class ToyRetriever(BaseRetriever):
"""A toy retriever that contains the top k documents that contain the user query.

This retriever only implements the sync method _get_relevant_documents.

If the retriever were to involve file access or network access, it could benefit
from a native async implementation of `_aget_relevant_documents`.

As usual, with Runnables, there's a default async implementation that's provided
that delegates to the sync implementation running on another thread.
"""

documents: List[Document]
"""List of documents to retrieve from."""
k: int
"""Number of top results to return"""

def _get_relevant_documents(
self, query: str, *, run_manager: CallbackManagerForRetrieverRun
) -> List[Document]:
"""Sync implementations for retriever."""
matching_documents = []
for document in documents:
if len(matching_documents) > self.k:
return matching_documents

if query.lower() in document.page_content.lower():
matching_documents.append(document)
return matching_documents

# Optional: Provide a more efficient native implementation by overriding
# _aget_relevant_documents
# async def _aget_relevant_documents(
# self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
# ) -> List[Document]:
# """Asynchronously get documents relevant to a query.

# Args:
# query: String to find relevant documents for
# run_manager: The callbacks handler to use

# Returns:
# List of relevant documents
# """

测试它 🧪

documents = [
Document(
page_content="Dogs are great companions, known for their loyalty and friendliness.",
metadata={"type": "dog", "trait": "loyalty"},
),
Document(
page_content="Cats are independent pets that often enjoy their own space.",
metadata={"type": "cat", "trait": "independence"},
),
Document(
page_content="Goldfish are popular pets for beginners, requiring relatively simple care.",
metadata={"type": "fish", "trait": "low maintenance"},
),
Document(
page_content="Parrots are intelligent birds capable of mimicking human speech.",
metadata={"type": "bird", "trait": "intelligence"},
),
Document(
page_content="Rabbits are social animals that need plenty of space to hop around.",
metadata={"type": "rabbit", "trait": "social"},
),
]
retriever = ToyRetriever(documents=documents, k=3)
retriever.invoke("that")
[Document(page_content='Cats are independent pets that often enjoy their own space.', metadata={'type': 'cat', 'trait': 'independence'}),
Document(page_content='Rabbits are social animals that need plenty of space to hop around.', metadata={'type': 'rabbit', 'trait': 'social'})]

这是一个可运行的,所以它将受益于标准的Runnable接口!🤩

await retriever.ainvoke("that")
[Document(page_content='Cats are independent pets that often enjoy their own space.', metadata={'type': 'cat', 'trait': 'independence'}),
Document(page_content='Rabbits are social animals that need plenty of space to hop around.', metadata={'type': 'rabbit', 'trait': 'social'})]
retriever.batch(["dog", "cat"])
[[Document(page_content='Dogs are great companions, known for their loyalty and friendliness.', metadata={'type': 'dog', 'trait': 'loyalty'})],
[Document(page_content='Cats are independent pets that often enjoy their own space.', metadata={'type': 'cat', 'trait': 'independence'})]]
async for event in retriever.astream_events("bar", version="v1"):
print(event)
{'event': 'on_retriever_start', 'run_id': 'f96f268d-8383-4921-b175-ca583924d9ff', 'name': 'ToyRetriever', 'tags': [], 'metadata': {}, 'data': {'input': 'bar'}}
{'event': 'on_retriever_stream', 'run_id': 'f96f268d-8383-4921-b175-ca583924d9ff', 'tags': [], 'metadata': {}, 'name': 'ToyRetriever', 'data': {'chunk': []}}
{'event': 'on_retriever_end', 'name': 'ToyRetriever', 'run_id': 'f96f268d-8383-4921-b175-ca583924d9ff', 'tags': [], 'metadata': {}, 'data': {'output': []}}

贡献

我们感谢有趣检索器的贡献!

以下是一个清单,帮助确保您的贡献被添加到LangChain中:

文档:

  • 检索器包含所有初始化参数的文档字符串,因为这些将在API参考中展示。
  • 模型的类文档字符串包含用于检索器的任何相关API的链接(例如,如果检索器是从维基百科检索的,最好链接到维基百科API!)

测试:

  • 添加单元或集成测试以验证invokeainvoke是否正常工作。

优化:

如果检索器正在连接到外部数据源(例如,API 或文件),它几乎肯定会从异步本地优化中受益!

  • 提供_aget_relevant_documents的原生异步实现(由ainvoke使用)

这个页面有帮助吗?