[大模型实战 04] 从玩具到生产：基于 ChromaDB 打造工程级 RAG 系统

发表于 分类于 ， ， ， 阅读次数： 本文字数： 3.8k 阅读时长 ≈ 3 分钟
大模型应用开发必修课！本文通过通俗易懂的案例详解 RAG (检索增强生成) 的核心原理，对比 LlamaIndex 与 LangChain 的优劣，并手把手带咱们在 Kaggle 环境下，利用 ChromaDB 和 Qwen 模型搭建一个支持私有数据问答的企业级 RAG 系统。

[大模型实战 04] 从玩具到生产：基于 ChromaDB 打造工程级 RAG 系统

核心摘要 (TL;DR)

  • 痛点：大模型不知道最新的新闻，也不知道企业的私有文档（如员工手册）。
  • 方案：RAG (检索增强生成)。就像“开卷考试”，先去翻书找答案，再回答问题。
  • 工具链：LlamaIndex (框架) + BGE (嵌入模型) + ChromaDB (向量数据库) + Qwen2.5 (推理模型)。
  • 实战：在 Kaggle 上从零搭建一个能回答“企业内部机密”的 AI 助手。

前言

各位友人们，大家好，这里是阿尔。在上一节中，我们大概知道了大模型的构成，safetensor格式的大模型的文件组成，transformers库的基本使用。我们已经能够使用大模型去做一些简单对话应用了，它可以是上知天文，下知地理，中间还能知道人情冷暖。但是，我们需要加一个限定词，在训练数据截止日期前的。因为训练一次需要耗费很多的计算资源，时间和人力，当我们想让它知道一些新知识的时候，比如让它知道现在美国的总统是拜登，我们可以在对话中告诉他，这没问题，但是如果我们想让它知道更多，比如我的私人日记?比如我刚写的那篇博客？比如公司的员工手册, 比如自己产品的使用说明书？

这类私有数据，是大模型企业应用的痛点，毕竟大模型是基于在互联网上公开数据训练的。重新把这部分资料加进去，再训练一下模型？也不是不行，但是有点没有性价比，这时候就引出了大模型落地应用的核心技术-> RAG (Retrieval-Augmented Generation，检索增强生成)。

1. RAG（检索增强生成）

1.1 什么是RAG?

考试的时候，如果考到不会的知识，不知道各位友人们会不会头疼，如果这时候，允许我们翻书，现去书里找，我们也很有可能找得到对应的答案，哪怕我们可能完全没学过。这就是RAG的大致思路：不让模型凭空回忆，而是先给它找资料。

RAG，检索增强生成,字面上讲，就是 拿到考题->然后去翻书，通过目录之类的索引，快速翻到（检索）相关的内容->再根据这些内容（增强了的内容），回答出问题（生成回答）。

对比简单地把东西一股脑全部跟大模型说一遍，我们能清楚得发现，我们只用了检索到的那一部分内容，并没有让整本书大模型的脑子将占用, 这就是RAG的效率体现。

1.2 RAG的步骤

RAG技术的思路很简单，但是实现并非只是一个单一的技术能实现的，它有一套流水线（流水线）。 把这头"大象"放进冰箱，总共需要两步：准备好数据和让模型拿到数据。

第一个阶段：数据准备(Indexing) -> 把书装进书包

在大模型能够翻书之前，咱们得先把我们想给它看的书整理好，放进书包里。

  1. 加载 (Load)：咱们的资料可能是各种各样的格式，一般大模型是不认识这么些格式的，所以我们就需要把 PDF、Word、网页等各种格式的文件读进来，统一提取出纯文本。
  2. 切分 (Chunking)：大模型一次吃不下整本书,就和我们一眼看不完整本《三国演义》一样，它有上下文长度限制。我们需要把长文本切成一个个小的片段 (Chunks)，比如每 500 个字切一段。
  3. 向量化 (Embedding)：这是最关键的一步！
    • 计算机无法直接比较“苹果”和“iphone”是不是相关的。
    • 我们需要用一个专门的模型（Embedding Model），把每一段文字变成一串数字向量（比如 [0.1, -0.5, 0.8, ...]）,是不是有点耳熟，对这和大模型训练的Embedding是一个思路，但是我们一般会使用特制的嵌入模型来做这个专业的事情。
    • 在这个数学空间里，语义相近的词，距离就越近, 这样我们就能知道，这本书中的所有向量，哪些是和我们的问题相关的了。
  4. 存储 (Storage)：把这些向量和对应的文字，存入向量数据库 (Vector DB) 中。

第二个阶段：应用数据给大模型生成（Retrieval & Generation）-> 开卷答题

拿到书了之后，我们想要翻书，就得找到和问题有关系的内容，然后再将这些内容和我们自己的常识结合起来，对提出的问题进行答题。

  1. 问题向量化(Embedding)：要想知道用户的提问（例如“火星基地吃什么？”）和内容的相关性，我们就需要像对准备的数据一样，用同一个 Embedding 模型将问题变成向量。
  2. 检索 (Retrieval)：拿着这个“问题向量”，去向量数据库里搜, 去找到关系性高的内容。
    • 系统会计算：“哪个文档片段的向量，和问题向量的距离最近？”
    • 找出最相似的前 3-5 个片段 (Top-k)。
  3. 增强 (Augmentation)：把这 3-5 个片段拼在一起，和用户的问题组合成一个超级长的 Prompt。
    • Prompt 模板示例：

      你是一个助手。请根据以下参考资料回答问题。
      参考资料：[片段1]… [片段2]…
      用户问题：火星基地吃什么？

  4. 生成 (Generation)：把这个 Prompt 喂给大模型（LLM）。大模型阅读资料，总结并生成最终答案。

2. RAG技术选型

好了，理论我们已经懂了，现在我们撸起袖子，准备来实操一下子吧。我们打算从零开始快速搭建一个工程级的RAG系统: 私有API助手, 在我直接告诉各位友人们我们要用到的工具前，我觉得也有必要大概让各位友人们知道还有哪些别的选择，我们为什么选择了这几个。

2.1 框架: LlamaIndex vs. LangChain

  • LangChain：万能胶水，适合做复杂的 Agent（智能体），但写 RAG 代码比较啰嗦，抽象层级太碎，我们后面写智能体的时候（如果有精力做智能体的教程的话）再来使用它。
  • LlamaIndex：数据专家。专门为 RAG 也就是“索引和检索”而生。接口极度简洁，且对数据清洗（Ingestion）的处理更专业。
  • 结论：我们做RAG，直接先上LlamaIndex, 快速地实现效果。

2.2 嵌入模型 (Embedding)：BGE vs. OpenAI

  • OpenAI (text-embedding-3)：效果好，但要钱，且数据要传给 OpenAI（隐私风险）。
  • BAAI/bge-small-zh-v1.5：国货之光。中文效果霸榜，体积极小（几百 MB），完全可以在 Kaggle 本地跑。
  • 结论：为了免费和隐私，首选 BGE-Small。
  • PS: 如果是英文资料的话，建议换成 BAAI/bge-small-en-v1.5 或者 OpenAI 的 text-embedding-3-small

2.3 向量数据库：Chroma vs. Milvus vs. Pinecone

  • Pinecone：纯云端 SaaS，不可本地部署，对 Kaggle 不友好。
  • Milvus：性能强悍，适合十亿级数据，需要 Docker 部署，适合数据量大的时候使用，但是对于咱们的这个项目来说，太重了。
  • ChromaDB：轻量级王者。可以像 SQLite 一样以“本地文件”形式存在，也可以部署成服务器。
  • 结论：中小型项目，首选 ChromaDB。

3. 上手实操

项目背景：假设我们是一家名叫 “DeepStar” 的初创公司，我们有一套内部绝密的 API 文档，新来的实习生总是问重复的问题。我们要用 RAG 让他自己查。

3.1 环境配置 (Kaggle)

启动 Kaggle Notebook，确保 Internet: On，Accelerator: GPU T4 x2。

1
2
3
4
5
6
7
8
# 1. 更新transformers及其相关库
!pip install -U transformers peft accelerate bitsandbytes sentence-transformers

# 2. 安装 LlamaIndex 核心及相关插件
!pip install llama-index-core llama-index-llms-huggingface llama-index-embeddings-huggingface

# 3. 安装 ChromaDB 向量库支持
!pip install llama-index-vector-stores-chroma chromadb

下载依赖库可能会需要一点时间，之后我看看能不能在kaggle上用uv去做包管理。

3.2 造数据：模拟企业内部文档

我们创建两份文档：一份是核心接口定义，一份是错误码说明。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
import os

data_path = "/kaggle/working/data"
# 创建数据目录
os.makedirs(data_path, exist_ok=True)

# 文档 1: 核心 API 定义
api_doc = """
[机密] DeepStar 核心交易接口 v2.0
1. 创建订单 API: POST /api/v2/order/create
   - 必填参数: 'user_id' (String), 'amount' (Decimal), 'token' (X-Auth-Token)
   - 特殊逻辑: 如果 amount > 10000, 必须额外传递 'audit_code' (审计码)。
   - 频率限制: 单用户每秒最多 5 次调用。
2. 查询余额 API: GET /api/v2/balance
   - 缓存策略: 默认缓存 5 秒。传递 'no-cache=true' 可强制刷新。
"""

# 文档 2: 错误码字典
error_doc = """
[机密] DeepStar 全局错误码字典
- E1001: 签名验证失败。请检查 X-Auth-Token 是否过期。
- E2009: 余额不足。注意：冻结金额不计入可用余额。
- E5003: 审计风控拦截。大额交易未通过自动审计，请联系人工客服。
"""

with open(f"{/data_path}/api_specs.txt", "w") as f:
    f.write(api_doc)
with open(f"/{data_path}/error_codes.txt", "w") as f:
    f.write(error_doc)

print("[Success] 企业文档库已就绪！")

3.3 初始化大脑与眼睛 (Settings)

提前根据自己的情况来配置待会儿用的词嵌入模型和推理模型。

1
2
3
embedding_model ="BAAI/bge-small-zh-v1.5"
llm = "Qwen/Qwen2.5-7B-Instruct"
# 在本地服务器，可以用modelscope下载下来, 把路径配置在这儿

利用 Settings 全局配置，将默认的 OpenAI 替换为本地模型。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
import torch
from llama_index.core import Settings
from llama_index.llms.huggingface import HuggingFaceLLM
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

# 1. 设置 Embedding (眼睛)
# 使用 BGE-Small，显存占用极低，检索中文效果极佳
print("正在加载 Embedding 模型...")

Settings.embed_model = HuggingFaceEmbedding(
    model_name=embedding_model
)

# 2. 设置 LLM (大脑)
# 使用 Qwen2.5-7B-Instruct
print("正在加载 LLM 模型...")
Settings.llm = HuggingFaceLLM(
    model_name=llm,
    tokenizer_name=llm,
    context_window=30000,
    max_new_tokens=512,
    generate_kwargs={"temperature": 0.1, "do_sample": True}, # 技术文档要求严谨，温度调低
    device_map="auto",
    model_kwargs={"dtype": torch.float16, "trust_remote_code": True}
)
print("[Success] 模型加载完毕！")

3.4 核心组件：ChromaDB 持久化流水线

这是本篇最关键的代码。我们要实现：如果本地已经有数据库，就直接读；如果没有，才去解析文档。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import chromadb
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext
from llama_index.vector_stores.chroma import ChromaVectorStore

# 定义持久化路径
CHROMA_DB_PATH = "/kaggle/working/chroma_db"
COLLECTION_NAME = "deepstar_docs"

# 1. 初始化 Chroma 客户端 (PersistentClient 实现了写硬盘功能)
db_client = chromadb.PersistentClient(path=CHROMA_DB_PATH)

# 2. 创建或获取集合 (Collection)
chroma_collection = db_client.get_or_create_collection(COLLECTION_NAME)

# 3. 将 Chroma 对接给 LlamaIndex
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)

# 4. 智能加载逻辑 (幂等性设计)
if chroma_collection.count() == 0:
    print("[Info] 数据库为空，开始初始化...")
    # 读取 data 目录下的所有文件
    documents = SimpleDirectoryReader("./data").load_data()
    # 建立索引并自动存入 Chroma (Ingestion)
    index = VectorStoreIndex.from_documents(
        documents, storage_context=storage_context
    )
    print("[Success] 数据写入完成！")
else:
    print(f"[Info] 发现 {chroma_collection.count()} 条存量数据，直接加载...")
    # 直接从 Vector Store 加载，无需重新计算 Embedding
    index = VectorStoreIndex.from_vector_store(
        vector_store, storage_context=storage_context
    )
    print("[Success] 索引加载完成！")

ChromaDB流水线示意图

3.5 验收测试：复杂逻辑问答

现在，我们模拟实习生提问。注意，这个问题需要结合两个文档（接口定义 + 错误码）以及逻辑推理才能回答。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
# 创建查询引擎
query_engine = index.as_query_engine(similarity_top_k=3)

# 实习生的提问
questions = [
    "创建订单时，如果你只有 100 块钱，能传 amount=20000 吗？为什么？",
    "我收到了 E5003 错误，这是什么意思？该怎么办？"
]

print("======== 开始 RAG 问答测试 ========")

for q in questions:
    print(f"\n[Question] {q}")
    response = query_engine.query(q)
    print(f"[Answer]\n{str(response)}")

    # 打印引用源 (Debug 必备，看看它参考了哪个文件)
    source_file = response.source_nodes[0].metadata.get('file_name')
    print(f"[Source]: {source_file}")

答复如下

4. 进阶技巧：如何管理你的数据库？

既然用了 ChromaDB，我们就可以像查 SQL 一样查它。这在 Debug 时非常有用。

1
2
3
4
5
6
7
8
# 偷看数据库里的前 2 条记录
data = chroma_collection.peek(limit=2)

print("\n[Debug] 数据库抽查:")
for i, doc in enumerate(data['documents']):
    print(f"--- 片段 {i} ---")
    print(f"内容: {doc[:50]}...") # 只打印前50个字
    print(f"来源: {data['metadatas'][i]}")

5. 完整代码

完整代码可以点击kaggle笔记获取

5. 常见问题 (Q&A)

Q: 为什么不直接把所有文档都塞进 Prompt 里 (Long Context)？
A: 虽然现在很多模型支持长文本（比如 128k），但直接塞文档有三个问题：

  1. 太贵：Token 是要钱的（如果用商业 API）。
  2. 太慢：上下文越长，推理越慢。
  3. 记不住：大模型有“长上下文迷失 (Lost in the Middle)”现象，塞太多反而会忽略中间的关键细节。RAG 相当于先做了一次筛选，只给模型看最有用的，效果反而更好。

Q: LlamaIndex 和 LangChain 我该学哪个？
A:

  • 做 RAG/知识库：首选 LlamaIndex，它对数据索引、切分、向量化做了极其深度的优化，接口更简洁。
  • 做 Agent/工具调用：首选 LangChain，它的生态和工具链更丰富。
  • 结论：咱们这个项目专注于“找资料”，所以 LlamaIndex 是最佳选择。

Q: ChromaDB 的数据存在哪里了？
A: 在上面代码中，我们通过 PersistentClient 指定了路径 /kaggle/working/chroma_db。
它就像 SQLite 一样，数据就存在这个文件夹里的 .sqlite3 和 .bin 文件中。咱们可以把这个文件夹拷贝到任何电脑上，无需重新向量化就能直接使用。