如何打造一个自定义的Retriever，实现高效信息检索

自定义Retriever创建指南与优化

最新推荐文章于 2025-05-03 18:56:05 发布

原创最新推荐文章于 2025-05-03 18:56:05 发布 · 417 阅读

7 ·

CC 4.0 BY-SA版权

文章标签：

#python #windows #linux

引言

在许多大型语言模型（LLM）应用中，检索器（Retriever）非常重要。它负责从外部数据源中检索与用户查询相关的文档，这些文档通常会被格式化为提示，供LLM使用，以生成合适的响应。例如，基于知识库回答用户问题。

本文将引导您如何创建自定义Retriever，深入理解其接口及实现，并提供实用的代码示例。

主要内容

Retriever接口

要创建自己的Retriever，需要继承BaseRetriever类并实现以下方法：

_get_relevant_documents: 获取与查询相关的文档（必需）
_aget_relevant_documents: 提供异步支持（可选）

实现_get_relevant_documents的方法可能涉及对数据库或通过网络请求进行任意调用。继承自BaseRetriever可自动获得LangChain的Runnable功能，使其成为可执行的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):
    """一个简单的Retriever，返回包含用户查询的前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

代码示例

下面是如何使用ToyRetriever的示例：

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)