打造专属信息检索器：一步步创建自定义Retriever-CSDN博客

本文链接：https://blog.csdn.net/akhfuiigabv/article/details/143610751

# 打造专属信息检索器：一步步创建自定义Retriever

## 引言

在许多大型语言模型（LLM）的应用中，数据检索扮演了至关重要的角色。一个有效的信息检索器（Retriever）可以从外部数据源中找到与用户查询相关的文档，并将其格式化为提示，供LLM生成响应。本篇文章将详尽介绍如何创建一个自定义的信息检索器，以便能够从多样化的数据源中高效获取信息。

## 主要内容

### 1. 什么是BaseRetriever？

在LangChain框架中，BaseRetriever是一个抽象类，专为创建自定义检索器而设计。它要求实现以下两个方法：

- `_get_relevant_documents`：同步获取与查询相关的文档。
- `_aget_relevant_documents`：异步获取与查询相关的文档（可选）。

通过扩展BaseRetriever，您可以轻松实现自定义的检索逻辑，利用数据库或网络请求等多种方式来获取数据。

### 2. 为何选择BaseRetriever？

与RunnableLambda等自定义可运行函数相比，使用BaseRetriever有以下优势：

- 标准化：BaseRetriever是LangChain生态中的一个知名实体，许多工具已经实现了对它的监控支持。
- 行为一致性：在某些API中，BaseRetriever的行为与其他可运行函数稍有不同，比如在stream_eventsAPI中，事件启动时以`on_retriever_start`而非`on_chain_start`标识。

### 3. 创建一个玩具检索器

下面是一个简单的示例，它检索包含用户查询文本的所有文档。

```python
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):
    """一个简单的检索器，根据查询文本检索文档"""

    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:
                break
            if query.lower() in document.page_content.lower():
                matching_documents.append(document)
        return matching_documents

    # 可选：提供更高效的异步实现
    # async def _aget_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"}),
    # 更多文档...
]

retriever = ToyRetriever(documents=documents, k=3)
result = retriever.invoke("that")
print(result)