用于与多个数据库聊天的智能 SQL 代理问答和 RAG 系统（2） —— 从 PDF 文档生成矢量数据库 (VectorDB)，然后存储文本的嵌入向量

从 PDF 文档生成矢量数据库 (VectorDB)，然后存储文本的嵌入向量。它结合了多个工具和库（如 langchain, PyPDFLoader, 和 Chroma）来处理文档，并将其转换为嵌入形式，以便用于基于向量的检索任务。以下是代码逻辑的详细解释：

首先定义一个 PrepareVectorDB 类，

类与方法

PrepareVectorDB 类

• 作用：

  • 加载 PDF 文档。
  • 将文档分块并生成文本嵌入。
  • 将生成的嵌入存储在一个持久化的向量数据库目录中。

• 构造函数 (__init__) 参数：

  1. doc_dir：文档所在的目录路径。
  2. chunk_size：分块时每块的最大字符数。
  3. chunk_overlap：分块时块之间的字符重叠数。
  4. embedding_model：生成嵌入向量的模型。
  5. vectordb_dir：向量数据库存储目录。
  6. collection_name：向量数据库中的集合名称。

path_maker 方法

• 作用：生成文件的完整路径。
• 实现：
  • 使用 os.path.join 将目录路径和文件名组合为完整路径。

run 方法

• 作用：

  • 检查 vectordb_dir 是否存在：
    • 如果不存在：创建目录，加载文档，分块生成嵌入，并存储到向量数据库中。
    • 如果存在：跳过向量数据库的生成过程，提示目录已经存在。

• 步骤：

  1. 创建目录：检查目录是否存在，若不存在则创建。
  2. 加载文档：使用 PyPDFLoader 加载目录下所有 PDF 文档并提取其内容。
  3. 分块：使用 RecursiveCharacterTextSplitter 将文档分块，支持分块大小及重叠。
  4. 生成向量嵌入：通过 OpenAIEmbeddings 模型生成文本块的嵌入。
  5. 存储向量数据库：通过 Chroma 保存嵌入向量到本地持久化数据库。
  6. 打印状态：提示数据库是否已生成，以及向量数量。

主要逻辑

1. 加载环境变量：

   load_dotenv()
   os.environ['OPENAI_API_KEY'] = os.getenv("OPENAI_API_KEY")
   

   • 使用 .env 文件中的 OPENAI_API_KEY 设置 OpenAI API 密钥。

2. 加载配置文件：

   with open(here("configs/tools_config.yml")) as cfg:
       app_config = yaml.load(cfg, Loader=yaml.FullLoader)
   

   • 从 tools_config.yml 文件加载配置参数。
   • 配置文件包含了不同任务（如 swiss_airline_policy_rag 和 stories_rag）的相关参数。

3. 执行向量数据库生成：

   • 读取配置中 swiss_airline_policy_rag 和 stories_rag 的参数。
   • 依次创建对应的 PrepareVectorDB 实例，调用 run 方法生成和保存向量数据库。

代码运行示例

输入

假设tools_config.yml文件配置文件包含以下内容：

swiss_airline_policy_rag:
  chunk_size: 500
  chunk_overlap: 50
  embedding_model: "openai"
  vectordb: "swiss_airline_vectordb"
  collection_name: "swiss_policies"
  unstructured_docs: "data/swiss_policies"

stories_rag:
  chunk_size: 400
  chunk_overlap: 40
  embedding_model: "openai"
  vectordb: "stories_vectordb"
  collection_name: "stories"
  unstructured_docs: "data/stories"

运行后：

1. 如果目录 swiss_airline_vectordb 不存在：
   • 加载 data/swiss_policies 下的 PDF 文件。
   • 分块并生成嵌入。
   • 保存嵌入到 swiss_airline_vectordb。
2. 同理处理 stories_vectordb。

代码运行结果：

代码依赖

1. 主要依赖库：
   • os, yaml: 用于文件和配置管理。
   • dotenv: 加载环境变量。
   • langchain_chroma: 用于持久化向量数据库。
   • langchain_community.document_loaders: 加载文档。
   • langchain_openai: 嵌入向量生成。
   • pyprojroot: 管理项目路径。
2. 配置文件：需要一个 tools_config.yml 文件定义参数。

总结

• 这段代码是一个自动化流程，用于将 PDF 文档处理为可用于检索的矢量嵌入，并持久化存储。
• 提供了灵活性，允许根据不同任务加载不同的文档和配置。
• 适合用于构建文档检索或问答系统等应用场景。