自定义Retriever的实现方法

技术背景介绍

在许多大语言模型（LLM）应用中，检索器（Retriever）用于从外部数据源获取信息。检索器的任务是根据用户查询检索相关的文档，这些文档通常被格式化为提示，供LLM使用，从而生成适当的响应，例如，根据知识库回答用户问题。

核心原理解析

要实现自定义的检索器，需要继承BaseRetriever类，并实现以下方法：

• _get_relevant_documents：获取与查询相关的文档，必需实现。
• _aget_relevant_documents：提供异步支持的实现，可选。

通过继承BaseRetriever，您的检索器将自动成为一个LangChainRunnable，并获得标准的Runnable功能。

代码实现演示(重点)

以下是一个简单的检索器示例，它返回所有文本中包含用户查询的文档：

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个文档。"""

    documents: List[Document]
    """要检索的文档列表"""
    k: int
    """返回的结果数量"""

    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

测试它

以下代码展示了如何使用该检索器：

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)

应用场景分析

这个简单的检索器示例展示了如何在定义好的文档集合中基于文本内容进行检索。此类应用对于知识管理系统、搜索引擎优化以及客户服务聊天机器人等场景非常有用。

实践建议

• 在涉及文件访问或网络访问时，考虑实现异步版本以提高性能。
• 确保为所有初始化参数编写文档字符串，以便更好地生成API引用。
• 编写单元或集成测试来验证invoke和ainvoke的工作。
• 提供处理外部数据源的优化，例如通过异步实现。

如果遇到问题欢迎在评论区交流。

—END—