理论部分
第 7 篇我们已经用"极简向量检索"跑通了 RAG 里最关键的一步：把文本向量化，并通过相似度找到语义最相关的内容。但那个实验还有两个明显的不足：
- 数据没有持久化：程序一退出，向量就没了，下次还要重新算
- 没有生成环节：只检索到了文档，还没有把"检索到的证据"交给大模型来回答
这一篇我们把 RAG 的三段式流程完整落地：
- Index（建库）：读取本地文档 → 切分成段落 → 计算 embedding → 写入向量库/索引
- Retrieve（检索）：用户提问 → 向量化 → Top-K 检索相关段落
- Generate（生成）：把相关段落拼成上下文 → 让模型基于资料回答 → 若资料没有答案就说不知道
在工程实现里，向量存储通常有两种典型选择：
- ChromaDB：更像一个"本地向量数据库"，自带集合（collection）、持久化、元数据等能力，适合快速上手
- FAISS：更像一个"高性能向量检索引擎"，速度快、依赖少，但持久化与文档映射需要我们自己处理
本项目把两种方式都实现了一遍：我们可以用同一份数据分别跑 Chroma 版和 FAISS 版，感受它们在使用体验上的差别。
实践部分
本案例做什么
我们会实现两个可运行的 RAG 问答脚本，做的事情一致：
- 第一次运行时，检测知识库是否为空；为空则读取
data/employee_handbook.txt初始化入库 - 启动交互式问答：输入问题 → 检索 Top-K 段落 → 基于段落生成答案
两个版本的差异在于"向量存储与检索"：
- Chroma 版：把段落与向量写入 ChromaDB 的 collection
- FAISS 版：把向量写入 FAISS 索引，把段落文本单独存成 pkl 文件映射
主要代码（ChromaDB 版本）
运行脚本：src/3.2_rag_chromadb.py
python
import os
import chromadb
from chromadb.utils import embedding_functions
from dotenv import load_dotenv
from openai import OpenAI
load_dotenv()
# === 配置 ===
# 使用 chromadb 的持久化存储模式
CHROMA_DATA_PATH = "chroma_db_data"
COLLECTION_NAME = "employee_handbook"
client_openai = OpenAI(
api_key=os.getenv("ZHIPUAI_API_KEY"),
base_url=os.getenv("ZHIPUAI_BASE_URL")
)
# === 1. 自定义 Embedding 函数 ===
# ChromaDB 默认使用 SentenceTransformers，但我们需要用智谱的 API
class ZhiPuEmbeddingFunction(embedding_functions.EmbeddingFunction):
def __call__(self, input):
# input 是一个字符串列表
embeddings = []
for text in input:
try:
response = client_openai.embeddings.create(
model="embedding-2",
input=text
)
embeddings.append(response.data[0].embedding)
except Exception as e:
print(f"Embedding Error: {e}")
# 出错时返回全0向量（权宜之计）
embeddings.append([0.0] * 1024)
return embeddings
# === 2. 知识库管理类 ===
class KnowledgeBase:
def __init__(self):
# 初始化 ChromaDB 客户端
self.chroma_client = chromadb.PersistentClient(path=CHROMA_DATA_PATH)
# 使用自定义的 Embedding 函数
self.embedding_fn = ZhiPuEmbeddingFunction()
# 获取或创建集合
self.collection = self.chroma_client.get_or_create_collection(
name=COLLECTION_NAME,
embedding_function=self.embedding_fn
)
def add_documents(self, file_path):
"""读取文件并存入 ChromaDB"""
if not os.path.exists(file_path):
print(f"文件不存在: {file_path}")
return
print(f"📖 读取文件: {file_path}")
with open(file_path, 'r', encoding='utf-8') as f:
text = f.read()
# 简单切分：按换行符切分段落
paragraphs = [p.strip() for p in text.split('\n\n') if p.strip()]
print(f"✂️ 切分为 {len(paragraphs)} 个段落，正在存入向量库...")
# 准备数据
ids = [f"doc_{i}" for i in range(len(paragraphs))]
metadatas = [{"source": file_path} for _ in paragraphs]
# 存入 ChromaDB (会自动调用 embedding_fn 计算向量)
self.collection.add(
documents=paragraphs,
ids=ids,
metadatas=metadatas
)
print(f"✅ 成功存入 {len(paragraphs)} 条数据！")
def query(self, question, top_k=3):
"""检索最相关的片段"""
results = self.collection.query(
query_texts=[question],
n_results=top_k
)
return results['documents'][0]
# === 3. RAG 问答函数 ===
def ask_ai(question, retrieved_docs):
context = "\n\n".join(retrieved_docs)
system_prompt = f"""
你是一个智能助手。请根据以下参考资料回答问题。
如果资料中没有答案，请直接说不知道。
【参考资料】
{context}
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": question}
]
print("🤖 AI 正在思考...")
response = client_openai.chat.completions.create(
model="glm-4-flash",
messages=messages,
temperature=0.1
)
return response.choices[0].message.content
# === 4. 主程序 ===
def main():
kb = KnowledgeBase()
data_path = os.path.join("data", "employee_handbook.txt")
if kb.collection.count() == 0:
print("⚠️ 检测到数据库为空，正在初始化数据...")
kb.add_documents(data_path)
else:
print(f"✅ 数据库已就绪，当前包含 {kb.collection.count()} 条记录。")
print("\n=== 🧠 基于 ChromaDB 的 RAG 助手 (输入 q 退出) ===")
while True:
question = input("\n👤 提问: ")
if question.lower() in ['q', 'quit']:
break
if not question.strip():
continue
docs = kb.query(question)
print(f"🔍 检索到 {len(docs)} 条相关资料")
answer = ask_ai(question, docs)
print(f"🤖 回答:\n{answer}")
if __name__ == "__main__":
main()主要代码（FAISS 版本）
运行脚本：src/3.2_rag_faiss.py
python
import os
import faiss
import numpy as np
import pickle
from dotenv import load_dotenv
from openai import OpenAI
load_dotenv()
# === 配置 ===
FAISS_DATA_PATH = "faiss_db_data"
FAISS_INDEX_PATH = os.path.join(FAISS_DATA_PATH, "faiss_index.bin")
DOC_STORE_PATH = os.path.join(FAISS_DATA_PATH, "faiss_docs.pkl")
EMBEDDING_DIM = 1024 # embedding-2 模型的维度
client_openai = OpenAI(
api_key=os.getenv("ZHIPUAI_API_KEY"),
base_url=os.getenv("ZHIPUAI_BASE_URL")
)
def get_embeddings(texts):
"""获取文本列表的向量表示"""
embeddings = []
for text in texts:
try:
response = client_openai.embeddings.create(
model="embedding-2",
input=text
)
embeddings.append(response.data[0].embedding)
except Exception as e:
print(f"Embedding Error: {e}")
embeddings.append([0.0] * EMBEDDING_DIM)
return np.array(embeddings).astype('float32')
class KnowledgeBase:
def __init__(self):
self.dimension = EMBEDDING_DIM
self.documents = []
self.index = None
if not os.path.exists(FAISS_DATA_PATH):
os.makedirs(FAISS_DATA_PATH)
self.load()
if self.index is None:
print("🆕 初始化新的 FAISS 索引...")
self.index = faiss.IndexFlatL2(self.dimension)
def add_documents(self, file_path):
"""读取文件并存入 FAISS"""
if not os.path.exists(file_path):
print(f"文件不存在: {file_path}")
return
print(f"📖 读取文件: {file_path}")
with open(file_path, 'r', encoding='utf-8') as f:
text = f.read()
paragraphs = [p.strip() for p in text.split('\n\n') if p.strip()]
if not paragraphs:
print("⚠️ 文件内容为空或无法切分")
return
print(f"✂️ 切分为 {len(paragraphs)} 个段落，正在计算向量...")
embeddings = get_embeddings(paragraphs)
self.index.add(embeddings)
self.documents.extend(paragraphs)
self.save()
print(f"✅ 成功存入 {len(paragraphs)} 条数据！")
def query(self, question, top_k=3):
"""检索最相关的片段"""
if self.index.ntotal == 0:
return []
query_vec = get_embeddings([question])
distances, indices = self.index.search(query_vec, top_k)
results = []
for idx in indices[0]:
if idx != -1 and idx < len(self.documents):
results.append(self.documents[idx])
return results
def save(self):
faiss.write_index(self.index, FAISS_INDEX_PATH)
with open(DOC_STORE_PATH, 'wb') as f:
pickle.dump(self.documents, f)
def load(self):
if os.path.exists(FAISS_INDEX_PATH) and os.path.exists(DOC_STORE_PATH):
try:
self.index = faiss.read_index(FAISS_INDEX_PATH)
with open(DOC_STORE_PATH, 'rb') as f:
self.documents = pickle.load(f)
print(f"📂 已加载现有知识库，包含 {len(self.documents)} 条记录")
except Exception as e:
print(f"⚠️ 加载失败，将重新初始化: {e}")
self.index = None
self.documents = []
else:
self.index = None
self.documents = []
def ask_ai(question, retrieved_docs):
context = "\n\n".join(retrieved_docs)
system_prompt = f"""
你是一个智能助手。请根据以下参考资料回答问题。
如果资料中没有答案，请直接说不知道。
【参考资料】
{context}
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": question}
]
print("🤖 AI 正在思考...")
response = client_openai.chat.completions.create(
model="glm-4-flash",
messages=messages,
temperature=0.1
)
return response.choices[0].message.content
def main():
kb = KnowledgeBase()
data_path = os.path.join("data", "employee_handbook.txt")
if kb.index.ntotal == 0:
print("⚠️ 检测到数据库为空，正在初始化数据...")
kb.add_documents(data_path)
else:
print(f"✅ 数据库已就绪，当前包含 {kb.index.ntotal} 条记录。")
print("\n=== 🧠 基于 FAISS 的 RAG 助手 (输入 q 退出) ===")
while True:
question = input("\n👤 提问: ")
if question.lower() in ['q', 'quit']:
break
if not question.strip():
continue
docs = kb.query(question)
print(f"🔍 检索到 {len(docs)} 条相关资料")
answer = ask_ai(question, docs)
print(f"🤖 回答:\n{answer}")
if __name__ == "__main__":
main()运行方式
先跑 Chroma 版本：
bash
python3 src/3.2_rag_chromadb.py再跑 FAISS 版本：
bash
python3 src/3.2_rag_faiss.py运行结果示例
第一次运行时，会先初始化知识库（读取文件、切分、计算向量并写入持久化存储）。随后进入交互式问答：
text
⚠️ 检测到数据库为空，正在初始化数据...
📖 读取文件: data/employee_handbook.txt
✂️ 切分为 20 个段落，正在存入向量库...
✅ 成功存入 20 条数据！
=== 🧠 基于 ChromaDB 的 RAG 助手 (输入 q 退出) ===
👤 提问: 请问公司的请假制度是什么？
🔍 检索到 3 条相关资料
🤖 AI 正在思考...
🤖 回答:
...总结
这一篇我们把 RAG 从"能检索"升级到"能回答"：把检索到的段落作为参考资料交给模型，让模型基于资料生成答案，并把向量数据持久化下来，形成可复用的知识库。