RAG 向量持久化：用 ChromaDB 替换内存存储，支持 Metadata 溯源

系列导读：本系列共 6 篇，带你从零到一构建完整的 RAG + LangGraph + MCP 项目。

• 第 1 篇：最小 RAG 实现，纯 numpy，无任何 AI 框架
• 第 2 篇：接入 Ollama 本地大模型，实现真实语义检索
• 第 3 篇（本文）：接入 ChromaDB 持久化向量数据库
• 第 4 篇：用 LangChain 重构 + 多轮对话
• 第 5 篇：LangGraph 多步推理工作流
• 第 6 篇：MCP 工具调用协议集成

---

一、第 2 篇的问题：向量不持久化

第 2 篇的向量数据库存在内存里：

class VectorStore:
    def __init__(self):
        self.texts = []    # 内存，重启消失
        self.vectors = []  # 内存，重启消失

问题：每次程序重启，都要重新 embed 所有文档。

• 8 个文档：几秒
• 1000 个文档：几分钟
• 10 万个文档：几小时

这在生产环境是无法接受的。解决方案：向量持久化。

---

二、ChromaDB 是什么

ChromaDB 是专门为 AI 应用设计的向量数据库：

• 向量存储到磁盘，重启后直接加载
• 支持 metadata（来源、作者、时间等）
• 内置 HNSW 索引，百万级向量毫秒检索
• 支持 id 去重，安全多次插入

pip install chromadb

---

三、核心变化：ChromaVectorStore

用 ChromaDB 替换第 2 篇的手写 VectorStore，接口保持兼容：

import chromadb
from chromadb.utils import embedding_functions

class ChromaVectorStore:
    def __init__(self, db_path: str, collection_name: str):
        # 连接到本地 ChromaDB（自动创建目录）
        self.client = chromadb.PersistentClient(path=db_path)

        # 自定义 embedding 函数：调用 Ollama
        # ChromaDB 会在 add/query 时自动调用这个函数
        self.embed_fn = embedding_functions.OllamaEmbeddingFunction(
            url="http://localhost:11434/api/embeddings",
            model_name="nomic-embed-text",
        )

        # get_or_create：已存在则复用，不存在则新建
        # 关键：同一个 collection_name 不会重复创建
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            embedding_function=self.embed_fn,
            metadata={"hnsw:space": "cosine"},  # 使用余弦距离
        )

持久化写入

def add_documents(self, documents: list):
    # 检查哪些 id 已存在（id 去重，防止重复 embed）
    existing = set(self.collection.get()["ids"])
    new_docs = [d for d in documents if d["id"] not in existing]

    if not new_docs:
        print(f"已有 {len(existing)} 条数据，跳过建库")
        return

    self.collection.add(
        ids       = [d["id"]   for d in new_docs],
        documents = [d["text"] for d in new_docs],
        metadatas = [{"source": d["source"]} for d in new_docs],
        # embeddings 不传：ChromaDB 自动调用 embed_fn 生成
    )
    print("✅ 建库完成，数据已持久化到磁盘")

id 去重的意义：程序可以安全地多次运行，不会重复 embed 已存在的文档。第一次运行建库，之后直接复用。

带 Metadata 的检索

def search(self, query: str, top_k: int = 3) -> list:
    results = self.collection.query(
        query_texts=[query],
        n_results=top_k,
        include=["documents", "distances", "metadatas"],  # 包含 metadata
    )
    docs      = results["documents"][0]
    distances = results["distances"][0]
    metas     = results["metadatas"][0]

    # ChromaDB 返回余弦距离（越小越相似），转为相似度
    return [
        (doc, 1 - dist, meta["source"])
        for doc, dist, meta in zip(docs, distances, metas)
    ]

---

四、Metadata：知道答案来自哪里

这是 ChromaDB 相比手写 VectorStore 的重要升级------溯源能力。

DOCUMENTS = [
    {"id": "hr_1",  "text": "年假政策：员工入职满一年后享有15天年假...", "source": "HR手册第3章"},
    {"id": "hr_2",  "text": "病假政策：病假需提供医院证明...",           "source": "HR手册第3章"},
    {"id": "hr_3",  "text": "请假流程：登录OA系统...",                   "source": "HR手册第4章"},
    {"id": "fin_1", "text": "报销流程：填写费用报销单...",               "source": "财务制度第2章"},
    {"id": "fin_2", "text": "差旅标准：经济舱国内出差...",              "source": "财务制度第3章"},
]

检索时自动返回来源：

📋 检索结果：
  [1] 相似度=0.892  [HR手册第4章]  请假流程：登录OA系统，填写请假申请单...
  [2] 相似度=0.743  [HR手册第3章]  年假政策：员工入职满一年后享有15天年假...

生成时把来源附在 Prompt 里：

context_text = "\n".join(
    f"- {text}  （来源：{src}）"
    for text, _, src in contexts
)

这样大模型可以在回答中引用来源，用户知道信息出处，可信度大幅提升。

---

五、完整代码

import ollama
import chromadb
from chromadb.utils import embedding_functions

EMBED_MODEL = "nomic-embed-text"
CHAT_MODEL  = "qwen2.5:7b"
DB_PATH     = "./chroma_db"
COLLECTION  = "company_docs"

DOCUMENTS = [
    {"id": "hr_1",  "text": "年假政策：员工入职满一年后享有15天年假，满三年后20天。",     "source": "HR手册第3章"},
    {"id": "hr_2",  "text": "病假政策：病假需提供医院证明，当天请假需电话通知直属上级。", "source": "HR手册第3章"},
    {"id": "hr_3",  "text": "请假流程：登录OA系统，填写请假申请单，提前3个工作日提交。", "source": "HR手册第4章"},
    {"id": "hr_4",  "text": "加班规定：工作日加班按1.5倍计算，周末加班按2倍计算。",      "source": "HR手册第5章"},
    {"id": "fin_1", "text": "报销流程：填写费用报销单，附发票原件，提交财务部审核。",      "source": "财务制度第2章"},
    {"id": "fin_2", "text": "差旅标准：经济舱国内出差，商务舱国际出差超6小时。",         "source": "财务制度第3章"},
]

class ChromaVectorStore:
    def __init__(self, db_path, collection_name):
        self.client = chromadb.PersistentClient(path=db_path)
        self.embed_fn = embedding_functions.OllamaEmbeddingFunction(
            url="http://localhost:11434/api/embeddings",
            model_name=EMBED_MODEL,
        )
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            embedding_function=self.embed_fn,
            metadata={"hnsw:space": "cosine"},
        )

    def add_documents(self, documents):
        existing = set(self.collection.get()["ids"])
        new_docs = [d for d in documents if d["id"] not in existing]
        if not new_docs:
            print(f"已有 {len(existing)} 条，跳过建库")
            return
        self.collection.add(
            ids       = [d["id"] for d in new_docs],
            documents = [d["text"] for d in new_docs],
            metadatas = [{"source": d["source"]} for d in new_docs],
        )
        print("✅ 建库完成")

    def search(self, query, top_k=3):
        results = self.collection.query(
            query_texts=[query],
            n_results=top_k,
            include=["documents", "distances", "metadatas"],
        )
        return [
            (doc, 1 - dist, meta["source"])
            for doc, dist, meta in zip(
                results["documents"][0],
                results["distances"][0],
                results["metadatas"][0],
            )
        ]

def generate(query, contexts):
    context_text = "\n".join(f"- {text}  （来源：{src}）" for text, _, src in contexts)
    prompt = f"""你是企业知识库助手，根据资料回答问题。

【参考资料】
{context_text}

【问题】{query}
【回答】"""
    full = ""
    for chunk in ollama.generate(model=CHAT_MODEL, prompt=prompt, stream=True):
        print(chunk["response"], end="", flush=True)
        full += chunk["response"]
    print()
    return full

class ChromaRAG:
    def __init__(self):
        self.store = ChromaVectorStore(DB_PATH, COLLECTION)

    def build_index(self, documents):
        self.store.add_documents(documents)

    def query(self, question, top_k=2):
        results = self.store.search(question, top_k=top_k)
        print(f"\n🔍 问题：{question}")
        for i, (text, score, source) in enumerate(results, 1):
            print(f"  [{i}] {score:.3f}  [{source}]  {text}")
        return generate(question, results)

if __name__ == "__main__":
    rag = ChromaRAG()
    rag.build_index(DOCUMENTS)      # 第二次运行自动跳过
    rag.query("我想请假怎么申请？")
    rag.query("出差能坐商务舱吗？")

---

六、第二次运行的变化

第一次运行：

📚 新增 6 条文档...
✅ 建库完成，数据已持久化到磁盘

第二次运行：

已有 6 条，跳过建库（直接从磁盘加载）

这就是持久化的价值：embed 只做一次，之后复用。

---

七、ChromaDB 的 HNSW 索引

ChromaDB 底层使用 HNSW（Hierarchical Navigable Small World） 索引，是目前最流行的近似最近邻搜索算法：

• 第 1~2 篇手写的 matrix @ query_vec 是暴力搜索，O(N) 复杂度
• HNSW 是图结构索引，查询复杂度接近 O(log N)
• 百万级向量查询在毫秒内完成

# 配置 HNSW 参数
metadata={"hnsw:space": "cosine"}   # 余弦距离（语义检索推荐）
# 还可以配置：
# "hnsw:M": 16              # 每个节点的连接数，越大越准但占内存
# "hnsw:ef_construction": 100  # 建索引时的搜索宽度

---

八、三步对比总结

特性	第 1 篇	第 2 篇	第 3 篇
Embedding	随机向量	Ollama 语义向量	Ollama 语义向量
存储方式	内存（重启消失）	内存（重启消失）	磁盘（持久化）
检索算法	暴力搜索 O(N)	暴力搜索 O(N)	HNSW O(log N)
Metadata	无	无	支持来源追踪
重复 embed	每次都做	每次都做	只做一次

---

总结

本文核心改动是用 ChromaDB 替换手写的内存向量存储：

1. 持久化：向量存磁盘，程序重启后无需重新 embed
2. 去重：通过 id 防止重复插入，多次运行安全
3. Metadata：追踪每条文档的来源，提升可信度
4. HNSW 索引：大规模向量毫秒级检索

下一篇引入 LangChain，用标准组件替换手写代码，同时增加多轮对话记忆能力。