RAG 数据检索

数据检索

我们前面已经学习完了知识库的构建，那我们接下来来看一下数据检索是如何实现的。

在大模型回答问题之前，先从私有知识库 / 文档库里，快速找出和用户问题最相关、最匹配的一段 / 多段真实原文内容，喂给大模型，再让模型基于这些真实资料 + 自身知识来回答。

简单一句话： 先找资料，再答题；不让模型凭空瞎编。

数据检索这一部分，我会从以下 4 个部分给大家展开讲解。 这部分整体是比较简单的。大家回想一下，我们前期完成了知识库的构建，也体验过向量数据库的创建，如何往里面插入数据，以及如何根据文本、根据用户查询，检索得到相似度较高的文本内容。

那么在本章节，我们要实现的场景就是：已经建好知识库的前提下，如何处理用户输入的查询内容，经过向量化转换，把它作为查询条件，在向量数据库中完成检索，获取我们需要的内容，主要分为以下四个部分。

用户查询

第一部分就是用户查询 ，也就是我们的query。 在使用 RAG 的任意场景下，都必须要有用户输入，也就是用户的查询条件。

用户查询（User Query） 是用户在检索系统或对话系统中输入的原始问题、指令或关键词，用于表达其信息需求。在 RAG 系统中，它是整个检索与生成流程的起点。简单来说就是用户向应用发起提问。

比如我们问一个 AI 应用的内容就是用户查询，例如我们问公司的 AI 助理，"公司新发布的休假政策是什么？"。

查询向量化

第二部分是查询向量化 。 简单来说，就是把用户输入的query，通过embedding model嵌入模型，转换成向量数组，也就是浮点数组。

⚠️重点注意 ：这里使用的模型，必须和构建知识库时使用的模型保持完全一致。 比如我们使用 BGE‑M3 模型，在构建知识库、把分块文本向量化入库的时候用的是 BGE‑M3；做查询的时候，也必须使用同一个 BGE‑M3 模型。 如果使用不同的模型，生成的向量会处在不同的语义空间，检索出来的结果就会出现错误，所以模型一致性非常关键。

用户输入问题后，系统使用与索引阶段相同的嵌入模型将查询文本转换为向量表示，使其与知识库中的文档向量处于同一语义空间。这个步骤和我们将语句块转为向量化的过程一样，我们不再重复赘述。

向量检索

第三部分是向量检索 。 把用户query经过嵌入模型转换成向量之后，拿到向量数据库中执行查询。 将查询向量和向量库内存储的全部向量做相似度比对，同时可以通过limit参数控制返回结果数量，比如设置返回 10 条、20 条相似度最高的数据，这就是我们之前讲过的limit参数的作用。

将用户查询转换为向量后，系统需要在向量数据库中快速找到语义最接近的文本块。这个过程本质上是向量相似度检索，用搜索算法检索查询向量与库中所有向量的 "距离"，距离越近，语义越相似。距离之前我们已经介绍过了，这里就不重复介绍了。

代码实现

第四部分就是以上三个流程的代码实现。

大家要注意，数据检索整体逻辑比较简单。 我们的重点工作其实都在：数据加载、数据清洗、各类文本分块方案，以及分块之后的文本向量化、入库。 数据检索本质只是查询环节，所以本章的实现，完全是基于前面已经完成的工作。

讲完本篇整体的目录，我们就先来学习用户查询。

用户查询也很好理解。我们使用 RAG 的时候，首先接收用户查询，在知识库中检索出相关内容，再把检索到的内容连同用户的原始问题，一起送入大模型。 在提示词环节做内容限定，划定模型回答的边界，避免模型凭空生成内容，也就是解决大模型的幻觉问题。当知识库中没有对应资料时，模型容易胡乱编造答案，RAG 就是通过检索真实知识库内容来规避这个问题，这就是用户查询在 RAG 里的作用。

接下来是查询向量化 。 复用构建知识库时使用的那套嵌入模型，对用户输入的查询做向量表示。把用户输入的query传入嵌入模型，得到对应的vector向量。

然后就是向量检索。 把用户查询转换得到向量之后，直接在向量数据库执行相似度比对，取出语义相近的文本块即可。

概念部分都比较简单，接下来我们着手编写代码实现。

---

我们打开编辑器，新建一个 Python 文件，命名为vectorSearch.py，用来实现向量检索功能。

**回顾向量检索的完整流程：**首先对用户查询做向量化；同时需要构建向量数据库，提升检索效率。 所以第一步，我们先构建向量数据库。

首先声明数据库名称DB_NAME，我们命名为vectorSearch。 数据库定义完成之后，接着创建集合collection。

定义好数据库名、集合名之后，我们编写init()初始化函数，完成向量数据库的初始化，这里可以复用我们之前写好的milvus模块内的工具方法。

初始化第一步，先判断数据库是否已经存在。 Milvus 本地数据库本质是磁盘上的文件夹。我们构造db_path路径：config.base_path拼接上{DB_NAME}.db，这个就是数据库存放的目录。判断这个路径是否存在，如果已经存在，就调用shutil.rmtree()把旧数据库文件夹删除，并且加上time.sleep(0.5)等待，保证删除操作完成，再执行后续逻辑。

删除旧库完成之后，第二步创建数据库。 调用milvus.create_db()，传入我们定义好的DB_NAME。 数据库创建完毕，再基于该数据库创建集合，调用milvus.create_collection()，第一个参数传入集合名称COLLECTION；因为我们使用 BGE‑M3 模型，向量维度是 1024，我们把向量维度定义为常量DIMENSION = 1024，作为第二个参数传入。

到这里，向量数据库的初始化就完成了。

数据库初始化完成之后，下一步：文本分块，将分块内容向量化之后入库。 拿到原始文档，使用我们之前学过的分块方案：句子分块、语义分块、结构分块，把文档切分成文本块；再将每一个文本块向量化，插入向量数据库。

这里我们复用之前写好的分块逻辑，选择句子分块方案。 我们新建函数index_to_milvus()，返回值为int类型。

调用sentence_chunk_documents()做句子分块，传入config.structure_path文档路径，设置最大句子数参数为1，也就是一个句子作为一个文本块。

分块完成之后做校验：如果chunks分块结果为空，直接抛出ValueError异常，提示：分块结果为空，请检测数据集内容。

接下来准备入库数据，定义空列表rows。 循环遍历每一个分块chunk： 调用milvus.embed_model.get_text_embedding(chunk.text)，把分块文本转为向量vec。 把每一条数据组装成字典，追加到rows列表： 字典内包含三个字段：

• id：i+1，作为主键 id
• vector：刚才生成的向量vec
• text：原始分块文本chunk.text

所有分块全部处理完毕，调用milvus.insert()，传入集合名COLLECTION和rows，完成批量入库。 打印日志提示入库完成，输出入库记录条数和向量维度，函数返回入库记录数量。

# 导入需要用到的标准库
import shutil
import time

# 导入项目自定义模块
import config
import milvus
from sentenceChunk import sentence_chunk_documents

# ===================== 向量数据库配置常量 =====================
# 数据库名称
DB_NAME = "vectorSearch"
# 集合(表)名称
COLLECTION = "vectorSearch"
# 向量维度，和embedding模型输出维度保持一致
DIMENSION = 1024


def init():
    """
    初始化Milvus向量数据库
    功能：
    1. 如果旧的数据库文件存在，先删除旧库，清理残留数据
    2. 创建新的数据库
    3. 创建指定维度的集合(数据表)
    """
    # 拼接本地数据库文件路径
    dp_path = config.base_path / f"{DB_NAME}.db"

    # 判断数据库文件是否已经存在
    if dp_path.exists():
        # 递归删除整个旧数据库目录，清理历史数据
        shutil.rmtree(dp_path)
        # 短暂休眠，确保文件系统删除操作完成，避免文件占用报错
        time.sleep(0.5)

    # 创建Milvus数据库
    milvus.create_db(DB_NAME)
    # 在数据库内创建集合，指定向量维度
    milvus.create_collection(COLLECTION, DIMENSION)


def index_to_milvus() -> int:
    """
    文档分块 -> 生成embedding向量 -> 批量插入Milvus向量库
    :return: 成功插入的文档块数量
    :raises ValueError: 如果分块结果为空，抛出异常
    """
    # 读取文档路径，执行句子级分块，第二个参数为分块参数
    chunks = sentence_chunk_documents(config.structure_path, 1)

    # 校验分块结果，如果没有得到任何文本块，抛出异常终止流程
    if not chunks:
        raise ValueError("分块结果为空，请检查数据集内容")

    # 用于存放待插入Milvus的数据行
    rows = []

    # 遍历所有分块，生成向量，组装插入数据
    for i, chunk in enumerate(chunks):
        # 使用嵌入模型，把文本转换成向量数组
        vec = milvus.embed_model.get_text_embedding(chunk.text)
        # 组装单条记录：id、向量、原始文本
        rows.append(
            {
                "id": i + 1,          # 主键id，从1开始计数
                "vector": vec,        # embedding向量
                "text": chunk.text,   # 原始分块文本
            }
        )

    # 将组装好的全部数据批量插入Milvus集合
    milvus.insert(COLLECTION, rows)

    # 打印入库统计信息
    print(f"已经完成入库, {len(rows)} 个记录, 向量维度{DIMENSION}")

    # 返回插入的记录总数
    return len(rows)

以上就是知识库构建、分块、向量化、入库的完整代码流程，接下来就可以实现用户查询、查询向量化、向量检索的业务逻辑。

完成向量入库之后，接下来我们来实现查询功能，也就是文本检索功能。

我们定义一个函数search_similar，这个函数的第一个入参就是query，也就是用户输入的查询内容。第二个参数我们声明为limit，用来控制返回结果的条数，这里给它设置默认值为2。

接下来我们来实现函数内部逻辑。 首先需要加载集合，调用milvus.get_client().load_collection()，传入我们定义好的集合名称COLLECTION。 集合加载完成之后，执行查询操作，调用milvus.search_by_text()，这就是我们之前向量数据库里封装好的向量查询方法。

传入参数：

• 第一个参数：集合名COLLECTION
• text=query：传入用户的查询文本
• limit=limit：设置返回结果数量
• fields=["id", "text"]：指定需要返回的字段

拿到查询返回结果之后，我们要对返回的数据做简单的整理处理。 我们观察原始返回结果，里面包含id、text还有distance相似度距离，原始的distance小数位数过长，可读性不好，所以我们要对输出做格式化处理。

我们用变量hits接收结果：

hits = results[0] if results else []

做判断，如果hits为空，说明没有检索到匹配内容，打印未检索到相关结果，然后直接return退出函数。

如果检索到结果，我们先打印用户的原始查询：

print(f"\n用户查询:{query}")

接着打印检索到结果的总条数：

print(f"共{len(hits)}条相似的结果")

然后遍历hits列表，使用enumerate循环，设置start=1，拿到每条结果的序号rank和结果对象hit。 从hit里面取出entity实体对象，再从entity中拿到text原始文本。

注意：数据结构里，原始文本存放在entity内部。

接下来格式化打印输出内容：

1. 打印排名rank，打印相似度距离distance。原始distance小数位数很多，我们做格式化，保留小数点后 4 位。
2. 打印id，这里要注意，外层结果对象hit有id，entity内部也有id，两个都打印出来方便对照。
3. 打印检索出来的原始文本text。

def search_similar(query: str, limit: int = 2):
    """
    根据用户查询文本，执行向量相似度检索
    :param query: 用户输入查询问题
    :param limit: 返回相似度最高的结果条数，默认2条
    """
    # 加载集合到内存，保证可以执行检索
    milvus.get_client().load_collection(collection_name=COLLECTION)

    # 执行文本向量检索，返回匹配结果
    results = milvus.search_by_text(
        COLLECTION,
        text=query,
        limit=limit,
        fields=["id", "text"],
    )

    # 取出检索命中结果列表
    hits = results[0] if results else []

    # 无命中结果处理
    if not hits:
        print("未检索到相关结果")
        return

    # 打印查询与结果统计
    print(f"\n用户查询:{query}")
    print(f"共{len(hits)}条相似的结果")

    # 遍历打印每条命中结果，输出排名、距离、id、原文
    for rank, hit in enumerate(hits, start=1):
        text = hit.get("text")
        entity = hit.get("entity")
        print(f"\n【rank】{rank} 距离 distance={hit.get('distance'):.4f}")
        print(f"【id】{entity.get('id')}")
        print(f"【text】{text}")

到此，search_similar查询函数就编写完成。

接下来我们要实现交互逻辑：程序运行之后，可以循环接收用户输入查询，持续检索知识库内容。 我们新建一个search()函数，实现循环交互。

函数内部先打印提示信息：向量检索系统，请输入用户查询问题。 开启while True无限循环：

• 使用input()获取用户输入，调用.strip()去除首尾空格，赋值给query。
• 判断，如果用户输入q，代表退出程序，打印已经退出，执行break跳出循环。
• 如果用户输入为空，执行continue，继续等待下一次输入。
• 输入合法查询文本，调用search_similar(query)，使用默认limit=2，执行检索。

最后编写程序入口if __name__ == "__main__": 程序执行的顺序分为三步：

1. 调用init()，完成向量数据库初始化，创建数据库和集合；
2. 调用index_to_milvus()，完成文档分块、向量化、数据入库；
3. 调用search()，启动交互循环，接收用户查询，执行检索。

def search():
    """
    交互式检索循环：持续接收用户输入查询，调用相似度检索
    输入 q / quit 即可退出程序
    """
    print("向量检索系统，请输入用户查询问题")
    while True:
        # 获取用户输入，去除首尾空白字符
        query = input("请输入查询问题：").strip()

        # 退出条件：输入q或者quit
        if query in ["q", "quit"]:
            print("已经退出")
            break

        # 空输入直接跳过
        if not query:
            continue

        # 执行相似度检索
        search_similar(query)


if __name__ == "__main__":
    # 程序入口：初始化库 -> 文档入库 -> 开启交互式检索
    init()
    index_to_milvus()
    search()

写到这里，我们就完整实现了用户查询、向量检索的整套流程。

那我们把这个代码写完了之后，我们来测一下它这个效果。这块我们先来执行代码，看看它能否按照我们的预期运行。

我们执行程序，观察两件事：第一，能否把这个数据创建出来；第二，能否正常进行查询。 首先程序会先把模型先加载到内存。

好，可以看到已经完成入库，一共16个记录，向量维度是1024维。因为我们原始数据集采用按句子分块的策略，每一个句子单独定为一个块，最终得到 16 条记录。

Loading weights: 100%|██████████| 391/391 [00:00<00:00, 48135.39it/s]
已经完成入库,16个记录，向量维度1024
向量检索系统，请输入用户查询问题
请输入查询问题：

接下来我们输入问题，测试检索效果，我们选取数据集里面存在的业务场景来提问：PDF 文件乱码的原因是什么？ 执行程序，观察向量数据库返回结果，一共查出来了两条。

这里要注意： 返回的数值并不是相似度，而是距离 。距离数值越小代表向量越接近。 第一条距离为0.2，距离比较近，返回内容就是 RAG 数据集里面原始文本，完整原样返回。 第二条对应内容是内存溢出相关，距离达到0.49，向量距离就相对更远。

我们再换一个问题测试：文本分块有哪几种策略。 执行检索，可以成功命中对应的原始句子，文档 ID 也正常返回。 第二条返回结果是检索增强策略，因为关键词 "策略" 发生匹配，向量距离相对更近。

到这里，本章数据检索的全部内容就实现完毕。整体难度不高，大部分底层能力，在前面章节知识库构建阶段就已经铺垫完成。