从基础到进阶：基于RAG PubMed检索AI Agents (三)

这是本系列的第三部分，今天我们将继续构建一个基于RAG（检索增强生成）的AI Agents，目的是回答问题并与生物医学科研文摘进行互动。

在上一部分中，我们已经实现了根据使用者提出的自然语言问题，从PubMed数据库中获取相关的科学文摘。接下来，我们将重点讲解如何保存这些数据，如何将获取的文摘转化为向量，并将它们存储到向量数据库中（我们将以ChromaDB为例来演示）。

Chroma DB 是一个用来存储和快速查找向量数据的数据库，常用于处理机器学习中的相似度搜索。简单来说，它可以帮助你高效地找到和某个数据相似的其他数据。

再次提醒，这就是我们在整个系列中要实现的完整解决方案：

构建过程

已完成的步骤概述

如果你还没有跟随第一部分和第二部分，建议先阅读，因为接下来的内容将在这些基础上继续构建。在上一部分结束时，我们的项目结构大致如下：

.
├── app
│   ├── app.py
│   ├── backend
│   │  ├── abstract_retrieval
│   │  │   ├── interface.py
│   │  │   ├── pubmed_retriever.py
│   │  │   └── pubmed_query_simplification.py
│   │  └── data_repository
│   │  │   ├── interface.py
│   │  │   ├── local_data_store.py
│   │  │   └── models.py
│   ├── components
│   │   ├── chat_utils.py
│   │   ├── llm.py
│   │   └── prompts.py
│   └── tests
│       └── test_chat_utils.py
├── assets
│   └── pubmed-screener-logo.jpg
└── environment
    └── requirements.txt

为数据创建抽象层

• 在上一部分中，我们学习了如何使用自然语言查询检索生物医学文摘。现在，我们将构建一个逻辑，用来持久化数据，供后续在我们的Streamlit应用中使用。
• 我们将采用仓库模式来为我们的数据创建一个抽象层。如果使用不同于本教程中使用的数据存储方式（为了简单起见，我们将使用本地文件系统），这种方式对构建本地业务特别有用。例如，你可以使用Postgres数据库。使用仓库模式，你只需要插入你自己实现的通用接口即可。

构建数据仓库的步骤

• 定义pydantic模型，表示数据库记录，包含关于用户查询的详细信息（UserQueryRecord），以及检索到的文摘模型（ScientificAbstract）。
• 我们将把这些模型添加到我们现有的**models.py文件中的data_repository**模块下：

from typing import Optional
from pydantic import BaseModel


class ScientificAbstract(BaseModel):
    doi: Optional[str] = None
    title: Optional[str] = None
    authors: Optional[str] = None
    year: Optional[int] = None
    abstract_content: str

class UserQueryRecord(BaseModel):
    user_query_id: str
    user_query: str

• 在UserQueryRecord中，我们要记录用户的原始查询内容，这些内容会显示在界面上。同时，还需要一个user_query_id——这是一个唯一的整数，用来给文件命名，并通过它来查找文件。
• 在ScientificAbstract类中，我们将存储文摘的内容和相关的元数据。
• 接下来，我们需要在data_repository模块中创建一个**interface.py**文件：

from abc import ABC, abstractmethod
from typing import List, Dict
from langchain_core.documents.base import Document
from backend.data_repository.models import ScientificAbstract


class UserQueryDataStore(ABC):
    """与文摘数据库交互的仓储类"""

    @abstractmethod
    def save_dataset(self, abstracts_data: List[ScientificAbstract], user_query: str) -> str:
        """
        将文摘数据和查询细节保存到数据存储中。
        返回一个字符串，表示新分配的查询ID。
        """
        raise NotImplementedError

    @abstractmethod
    def read_dataset(self, query_id: str) -> List[ScientificAbstract]:
        """
        从数据存储中检索指定查询ID的文摘数据。
        """
        raise NotImplementedError

    @abstractmethod
    def delete_dataset(self, query_id: str) -> None:
        """
        从数据库中删除指定查询ID的所有数据。
        """
        raise NotImplementedError

    @abstractmethod 
    def get_list_of_queries(self) -> Dict[str, str]:
        """
        检索查询ID和用户查询的字典。用于在UI上显示查询列表，并供查找使用。
        """
        raise NotImplementedError

    def create_document_list(self, abstracts_data: List[ScientificAbstract]) -> List[Document]:
        """
        将文摘数据转换为 LangChain 文档对象的列表。
        每个文档包含文摘内容及其元数据。
        """
        return [
            Document(
                page_content=entry.abstract_content, metadata={
                    "source": entry.doi, "title": entry.title, 
                    "authors": entry.authors, "year_of_publication": entry.year
                }
            )
            for entry in abstracts_data
        ]

    def read_documents(self, query_id: str) -> List[Document]:
        """ 
        读取数据集并将其转换为所需的 List[Document] 类型。
        """
        query_record = self.read_dataset(query_id)
        return self.create_document_list(query_record)

• 这个接口定义了我们与数据库交互的主要方法，用来处理用户查询相关的文摘。
• 有三个抽象方法需要实现，还有两个通用方法会被继承——create_document_list，它将文摘列表转换为LangChainDocuments格式，和read_documents，这是read_dataset和create_document_list的封装，方便我们调用。

请注意，我们会直接使用文摘的完整内容来创建Documents，因为文摘本身很简短。如果是较长的文本，我们需要先将文本分块处理，再加载到LangChain Documents中。

• 接下来，在data_repository模块下创建一个local_data_store.py文件，实现接口中的抽象方法。这个文件将负责把数据以JSON文件的形式保存在本地（当然，也可以根据需求使用其他存储方式）。

import json
import os
import shutil
from typing import Dict, List
from backend.data_repository.models import UserQueryRecord, ScientificAbstract
from backend.data_repository.interface import UserQueryDataStore
from config.logging_config import get_logger


class LocalJSONStore(UserQueryDataStore):
    """ 
    用于本地测试，通过本地JSON文件模拟数据库存储。
    """

    def __init__(self, storage_folder_path: str):
        self.storage_folder_path = storage_folder_path
        self.index_file_path = os.path.join(storage_folder_path, 'index.json')
        self.logger = get_logger(__name__)
        self.metadata_index = None

    def get_new_query_id(self) -> str:
        """
        通过递增上一个查询ID的整数后缀来生成新的查询ID。
        """
        try:
            with open(self.index_file_path, 'r') as file:
                data = json.load(file)
        except (FileNotFoundError, json.JSONDecodeError):
            data = {}
        keys = [k for k in data.keys() if k.startswith('query_')]
        if not keys:
            return 'query_1'
        numbers = [int(k.split('_')[-1]) for k in keys]
        max_number = max(numbers)
        return f'query_{max_number + 1}'

    def read_dataset(self, query_id: str) -> List[ScientificAbstract]:
        """ 
        从本地存储读取包含文摘的数据集。
        """
        try:
            with open(f'{self.storage_folder_path}/{query_id}/abstracts.json', 'r') as file:
                data = json.load(file)
                return [ScientificAbstract(**abstract_record) for abstract_record in data]
        except FileNotFoundError:
            self.logger.error(f'未找到查询 {query_id} 的JSON文件。')
            raise FileNotFoundError('未找到JSON文件。')

    def save_dataset(self, abstracts_data: List[ScientificAbstract], user_query: str) -> str:
        """ 
        将文摘数据集和查询元数据保存到本地存储，重建索引，并返回查询ID。
        """
        try:
            query_id = self.get_new_query_id()
            user_query_details = UserQueryRecord(
                user_query_id=query_id, 
                user_query=user_query
            )

            os.makedirs(f'{self.storage_folder_path}/{query_id}', exist_ok=True)
            
            with open(f"{self.storage_folder_path}/{query_id}/abstracts.json", "w") as file:
                list_of_abstracts = [model.model_dump() for model in abstracts_data]
                json.dump(list_of_abstracts, file, indent=4)

            with open(f"{self.storage_folder_path}/{query_id}/query_details.json", "w") as file:
                json_data = user_query_details.model_dump_json(indent=4)
                file.write(json_data)

            self.logger.info(f"查询ID {query_id} 的数据保存成功。")
            self._rebuild_index()  # 数据保存后重建索引

            return query_id

        except Exception as e:
            self.logger.error(f"保存查询ID {query_id} 数据集失败: {e}")
            raise RuntimeError(f"由于错误导致保存数据集失败: {e}")
        
    def delete_dataset(self, query_id: str) -> None:
        """ 
        从本地存储中删除文摘数据集和查询元数据。
        """
        path_to_data = f'{self.storage_folder_path}/{query_id}'
        if os.path.exists(path_to_data):
            shutil.rmtree(path_to_data)
            self.logger.info(f"目录 '{path_to_data}' 已删除。")
            self._rebuild_index()  # 删除数据后重建索引
        else:
            self.logger.warning(f"目录 '{path_to_data}' 不存在，无法删除。")

    def get_list_of_queries(self) -> Dict[str, str]:
        """ 
        从索引中获取包含查询ID（作为键）和原始用户查询（作为值）的字典。
        """
        return self.metadata_index

    def _rebuild_index(self) -> Dict[str, str]:
        """ 
        从所有查询详情文件重建索引，供查询使用。
        """
        index = {}
        query_data_paths = [os.path.join(self.storage_folder_path, name) for name in os.listdir(self.storage_folder_path)
                            if os.path.isdir(os.path.join(self.storage_folder_path, name))]
        
        for query_data_path in query_data_paths:
            metadata_path = os.path.join(query_data_path, 'query_details.json')
            if os.path.exists(metadata_path):
                with open(metadata_path, 'r') as file:
                    metadata = json.load(file)
                    index[metadata['user_query_id']] = metadata['user_query']
            else:
                self.logger.warning(f"在 {query_data_path} 中未找到 query_details.json 文件")
        
        with open(self.index_file_path, 'w') as file:
            json.dump(index, file, indent=4)
        self.metadata_index = index
        return index

• 我们刚刚实现了将检索到的科研文摘保存在本地文件中。每个用户的查询都会生成一个专属文件夹，这个文件夹包含两个文件：

  • abstracts.json：保存用户查询的相关文摘内容。
  • query_details.json：保存用户查询的详细信息，比如查询的原始文本。

• 除此之外，所有查询的索引信息还会保存在一个名为index.json的文件中，里面记录了每个查询ID和对应的用户查询内容。

这部分的工作已经完成了，现在我们可以继续进入vector DB部分。

构建RAG 工作流（RAG Pipeline）

在这个教程系列中，我们将从一个简单的RAG场景入手，虽然它有些局限，但后续会有更高级的技术讲解，比如GraphRAG。你可以在后面的文章中看到这些内容。

• 要构建这个RAG 工作流，我们需要做两个选择：

  1. 选择一个向量嵌入模型，它的作用是把文本（比如文摘）转化成一组数字，这样计算机就能“理解”文本内容，并能进行相似度比较。
  2. 选择一个向量存储，用于存储这些“理解过”的数字，方便后续的快速查找和比对。

• 在这个例子中，我选择了OpenAI向量嵌入模型和Chroma向量存储。不过，你也可以根据需要，选择不同的模型和存储方式。

简单总结

• 向量嵌入模型：将文本内容转化为数字向量（简单来说，就是把文本变成计算机能理解的数据）。
• 向量存储：就是保存这些数字向量的地方，方便后续查找。

接下来，我们将继续构建这个Pipeline，让模型能够从向量存储中找到相关的文摘内容，用来回答生物医学问题。

构建RAG工作流的步骤

1. 安装Chromadb驱动：

   • 你需要安装chromadb驱动，确保版本为0.4.24，可以通过以下命令来安装：

`pip install chromadb==`0.6.3

2. 更新**requirements.txt**文件：

*   安装完依赖后，`requirements.txt`文件的内容应该是这样的：

streamlit==1.41.1
langchain==0.3.14
langchain-community==0.3.14
langchain-core==0.3.30
langchain-openai==0.3.1
langchain-text-splitters==0.3.5
python-dotenv==1.0.1 
pydantic==2.10.5
pydantic-settings==2.7.1
pydantic_core==2.27.2
metapub==0.5.12  
chromadb==0.6.3

3. 创建RAG工作流模块：

*   在**backend**文件夹下创建一个新的子模块，命名为**rag\_pipeline**，用于存放RAG工作流的代码。

4. 定义RAG工作流接口：

*   在刚创建的**rag\_pipeline**模块下，创建一个新的文件，命名为\*\*[interface.py](http://interface.py)\*\*。这个文件将定义RAG工作流的接口。

from typing import List
from abc import ABC, abstractmethod
from langchain_core.documents.base import Document
from langchain_core.embeddings import Embeddings
from langchain.vectorstores import VectorStore


class RagWorkflow(ABC):
    """ 
    RAG工作流的接口
    """

    def __init__(self, embeddings: Embeddings):
        # 初始化时传入用于向量嵌入的对象
        self.embeddings = embeddings
    
    @abstractmethod
    def create_vector_index_for_user_query(self, documents: List[Document], query_id: str) -> VectorStore:
        """ 
        为特定用户查询创建基于文档的向量索引
        """
        raise NotImplementedError
    
    @abstractmethod
    def get_vector_index_by_user_query(self, documents: List[Document], query_id: str) -> VectorStore:
        """ 
        根据查询ID获取现有的向量索引
        """
        raise NotImplementedError

• 这个接口包含两个抽象方法——一个用于为给定的查询ID创建向量索引，另一个用于通过查询ID获取现有的向量索引。
• 接下来，我们来看一下ChromaDB的实现。在rag_pipeline文件夹下创建一个新文件，命名为chromadb_rag.py：

from typing import List
import chromadb
from langchain.vectorstores import VectorStore
from langchain_community.vectorstores import Chroma
from langchain_core.embeddings import Embeddings
from langchain_core.documents.base import Document
from backend.rag_pipeline.interface import RagWorkflow
from config.logging_config import get_logger


class ChromaDbRag(RagWorkflow):
    """ 
    使用 Chroma 作为向量存储的简单 RAG 工作流 
    """

    def __init__(self, persist_directory: str, embeddings: Embeddings):
        self.persist_directory = persist_directory  # 持久化存储的目录路径
        self.embeddings = embeddings  # 用于生成向量的嵌入模型
        self.client = self._create_chromadb_client()  # 创建 Chroma 数据库客户端
        self.logger = get_logger(__name__)  # 配置日志记录器
    
    def _create_chromadb_client(self):
        """ 创建并返回一个持久化的 Chroma 客户端 """
        return chromadb.PersistentClient(path=self.persist_directory)
    
    def create_vector_index_for_user_query(self, documents: List[Document], query_id: str) -> VectorStore:
        """
        为用户查询创建 Chroma 向量索引，并将查询 ID 作为集合名称。
        """
        self.logger.info(f'为查询 {query_id} 创建向量索引')
        try:
            # 使用文档和嵌入模型创建向量索引
            index = Chroma.from_documents(
                documents, 
                self.embeddings,
                client=self.client, 
                collection_name=query_id  # 将查询 ID 作为集合名称
            )
            return index
        except Exception as e:
            self.logger.error(f'创建查询 {query_id} 的向量索引时发生错误。错误信息：{e}')
            raise
    
    def get_vector_index_by_user_query(self, query_id: str) -> VectorStore:
        """
        根据查询 ID 加载已有的 Chroma 向量索引。
        """
        self.logger.info(f'加载查询 {query_id} 的向量索引')
        try:
            # 根据查询 ID 加载现有的向量索引
            index = Chroma(
                client=self.client,
                collection_name=query_id,
                embedding_function=self.embeddings,
            )
            return index
        except Exception as e:
            self.logger.error(f'加载查询 {query_id} 的向量索引时发生错误。错误信息：{e}')
            raise

• 这个实现负责创建新的向量索引，并从ChromaDB中获取现有的索引。
• 在rag_pipeline模块中创建一个新文件embeddings.py，并在其中初始化你的嵌入实例：

import os
from langchain_openai import OpenAIEmbeddings
from dotenv import load_dotenv

load_dotenv()

#  OpenAI 的 API Key
OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')

# 初始化 OpenAI 嵌入模型
embeddings = OpenAIEmbeddings(
    openai_api_key=OPENAI_API_KEY
)

测试客户端代码！

• 现在我们准备好进行客户端代码的测试了。测试的内容包括我们在本系列的上一部分中写的代码（PubMedRetriever），这个代码可以帮助我们从PubMed获取文章。
• 接下来，我们要创建一个测试脚本，命名为test_rag_pipeline.py，并在里面写下以下的客户端代码：

from metapub import PubMedFetcher
from backend.rag_pipeline.embeddings import embeddings
from backend.abstract_retrieval.pubmed_retriever import PubMedAbstractRetriever
from backend.data_repository.local_storage import LocalJSONStore
from backend.rag_pipeline.chromadb_rag import ChromaDbRag

query = "Does abamectin cause cancer?"

# 第一步：使用 PubMedAbstractRetriever 根据查询“Does abamectin cause cancer?” 获取文献摘要数据
pubmed_fetcher = PubMedFetcher()  # 创建 PubMedFetcher 实例，用于从 PubMed 获取数据
abstract_retriever = PubMedAbstractRetriever(pubmed_fetcher)  # 创建 PubMedAbstractRetriever 实例，用于检索文献摘要
abstracts = abstract_retriever.get_abstract_data(query)  # 使用 query 获取相关文献摘要

# 第二步：使用 LocalJSONStorage 将检索到的摘要数据保存在本地存储中
storage_folder_path = "backend/data"  # 定义本地存储文件夹路径
store = LocalJSONStore(storage_folder_path)  # 创建 LocalJSONStore 实例，用于管理本地存储
query_id = store.save_dataset(abstracts, query)  # 保存摘要数据到本地，并返回查询 ID

# 第三步：使用 ChromaDbRag 创建向量索引，利用从 LocalJSONStorage 中读取的文档
persist_directory = "backend/chromadb_storage"  # 定义 Chroma 数据存储目录
rag_workflow = ChromaDbRag(persist_directory, embeddings)  # 创建 ChromaDbRag 实例，用于处理向量存储
documents = store.read_documents(query_id)  # 读取存储中的文档数据
vector_index = rag_workflow.create_vector_index_for_user_query(documents, query_id)  # 创建向量索引

# 使用原始用户查询在新创建的向量索引上执行相似性搜索：
print(vector_index.similarity_search(query))  # 执行相似性搜索，并打印结果

• 客户端代码使用PubMedAbstractRetriever的get_abstract_data方法来检索文摘，使用LocalJSONStore将文摘保存到本地文件系统，接着使用ChromaDbRag创建向量索引。最后，在代码的最后一行，使用原始用户查询对该索引进行相似度搜索。
• 输出日志如下所示：

• 在日志的下面，你会看到一个打印出来的文档列表，这些文档在使用初始查询进行相似度搜索时得分最高。

总结

• 我们已经完成了基于PubMed的AI Agents应用后端部分的构建，包括数据持久化和向量化的实现。
• 我们测试了到目前为止构建的所有后端组件的协作情况。
• 我们展示了如何使用JSON文件和本地文件系统存储数据的示例，以及如何使用ChromaDB作为本地向量存储。
• 我们设计了接口，方便轻松替换自己喜欢的向量存储和数据仓库实现。

以上就是本部分的全部内容！在系列的最后一部分中，我们将结合相似度搜索和大语言模型（LLM），为使用人员的初始问题生成答案，并使用我们在Streamlit中构建的聊天界面进一步探索这些数据。期待与你再次相见！