如何创建自定义检索器
概述
许多 LLM 应用涉及使用 检索器 从外部数据源检索信息。
检索器负责检索与给定用户 query 相关的 文档 列表。
检索到的文档通常会被格式化为传递给 LLM 的提示,使 LLM 能够利用知识库中的信息来生成适当的响应(例如,根据知识库回答用户问题)。
接口
要创建自己的检索器,您需要继承 BaseRetriever 类并实现以下方法:
| 方法 | 描述 | 必需/可选 |
|---|---|---|
_get_relevant_documents | 获取与查询相关的文档。 | 必需 |
_aget_relevant_documents | 实现此方法可提供异步原生支持。 | 可选 |
_get_relevant_documents 中的逻辑可以涉及使用 requests 调用数据库或 Web。
通过继承 BaseRetriever,您的检索器将自动成为 LangChain 的 可运行组件,并立即获得标准的 Runnable 功能!
您可以使用 RunnableLambda 或 RunnableGenerator 来实现检索器。
与 RunnableLambda(自定义 可运行函数)相比,将检索器实现为 BaseRetriever 的主要好处是 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 self.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,因此它将受益于标准的 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 中:
文档:
- 检索器包含所有初始化参数的 doc-strings,因为这些将呈现在API 参考中。
- 模型的类 doc-string 应包含指向检索器使用的任何相关 API 的链接(例如,如果检索器从维基百科检索,最好链接到维基百科 API!)。
测试:
- 添加单元测试或集成测试来验证
invoke和ainvoke是否正常工作。
优化:
如果检索器连接到外部数据源(例如 API 或文件),它几乎肯定会受益于异步原生优化!
- 提供
_aget_relevant_documents(由ainvoke使用)的原生异步实现。