用于与多个数据库聊天的智能 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. 打印状态：提示数据库是否已生成，以及向量数量。

主要逻辑

加载环境变量：
python 复制代码
```
load_dotenv()
os.environ['OPENAI_API_KEY'] = os.getenv("OPENAI_API_KEY")
```
- 使用 .env 文件中的 OPENAI_API_KEY 设置 OpenAI API 密钥。
加载配置文件：
python 复制代码
```
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）的相关参数。
执行向量数据库生成：
- 读取配置中 swiss_airline_policy_rag 和 stories_rag 的参数。
- 依次创建对应的 PrepareVectorDB 实例，调用 run 方法生成和保存向量数据库。

代码运行示例

输入

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

yaml 复制代码

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"

运行后：

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

代码运行结果：

代码依赖

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

总结

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