如何创建自定义Retriever

技术背景介绍

老铁们，今天我们来聊聊在LLM应用中如何创建一个自定义的Retriever。很多时候，我们需要从外部数据源中检索信息，一个好的Retriever就是帮我们完成这个任务的关键。Retriever的任务是根据用户的查询来检索相应的Document，然后将这些文档格式化为提示信息，传递给LLM进行处理。这种处理方式让LLM能够利用信息生成合适的响应，比如从知识库中回答用户的问题。

原理深度解析

要创建一个自定义的retriever，你首先需要继承BaseRetriever类，并实现几个关键方法：

• _get_relevant_documents: 获取与查询相关的文档，这是必须实现的。
• _aget_relevant_documents: 提供异步的原生支持，这是可选的。

你可以在_get_relevant_documents方法中实现任意的逻辑，比如对数据库或者通过请求从网络获取数据。

Tip: 继承BaseRetriever后，你的retriever会自动成为一个LangChainRunnable，并获得标准的Runnable功能。

另外，你也可以通过RunnableLambda或者RunnableGenerator来实现一个retriever，但选用BaseRetriever的好处在于，它是一个LangChain中知名的实体，所以一些工具可以为retrievers实现特有的监控行为。

实战代码演示

我们来实现一个简单的retriever，返回所有包含用户查询的文档。

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."""

    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)

print(retriever.invoke("that"))

# Output:
# [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涉及到文件访问或网络请求，建议通过覆盖_aget_relevant_documents方法提供一个更高效的本地异步实现。这将允许retriever通过ainvoke实现异步调用，从而提高性能。

补充说明和总结

一个好的retriever能够极大地提升LLM应用的响应准确性和效率。在实际项目中，选择合适的架构和实现方式尤为重要。说到服务，https://yunwu.ai 提供一站式大模型解决方案，我个人一直在用，推荐给大家。

今天的技术分享就到这里，希望对大家有帮助。开发过程中遇到问题也可以在评论区交流~

—END—