第 4 篇:原生 Python 手写最简 RAG——5 步吃透底层原理

本文为「从零到落地:RAG 检索增强生成实战系列」第 4 篇,完整系列持续更新中。


前言

本篇不用任何框架,5 步手写一个完整 RAG 系统。代码逐行注释讲解,帮你理解框架底层到底做了什么。

前三篇我们搞懂了 RAG 的原理、全流程和技术选型。现在要动手写代码了——你可能会想,市面上有 Dify、LangGraph、LlamaIndex 这么多框架,为什么不直接用,偏要先手写一遍?

原因很简单:先理解底层,再用框架事半功倍

当你用 Dify 配一个”知识库”、在 LangGraph 里写一个”检索节点”,底层做的事情其实就是这几步:加载文档、文本分块、向量化、存入向量库、检索、拼接 Prompt、调 LLM。如果你不理解这些步骤,出了问题(检索不准、回答幻觉、分块不合理)就只能瞎调参数。手写一遍之后,再用任何框架你都知道它在封装什么。

本篇学完你将掌握

  • 用原生 Python + Chroma + Sentence-Transformers + DeepSeek API 从零搭建一个完整 RAG
  • 索引阶段(文档加载、分块、向量化、存储)的完整代码实现
  • 推理阶段(检索、拼接 Prompt、调用 LLM)的完整代码实现
  • 基础 RAG 的实际效果和局限性
  • 理解框架(Dify / LangGraph / LlamaIndex)到底帮你封装了什么

一、环境准备

1.1 技术栈

组件 选择 理由
Python 3.11 系列统一版本
Embedding 模型 bge-large-zh-v1.5(本地) 中文效果最好的开源模型,通过 sentence-transformers 加载
向量数据库 Chroma 纯 Python,零配置启动,入门首选
LLM DeepSeek API 性价比最高,兼容 OpenAI SDK

💡 提示:如果你没有 DeepSeek API Key,可以替换为任何兼容 OpenAI 接口的服务(如通义千问 API、Moonshot API),代码只需要改 base_urlapi_key 两个参数。

1.2 安装依赖

1
2
3
4
5
6
7
8
9
# 创建虚拟环境(推荐)
python -m venv rag-env
# Windows 激活
rag-env\Scripts\activate
# macOS/Linux 激活
source rag-env/bin/activate

# 安装依赖包
pip install chromadb sentence-transformers openai

⚠️ 注意sentence-transformers 会自动下载 Embedding 模型(bge-large-zh-v1.5 约 1.2GB),首次运行需要联网下载,后续使用会走本地缓存。

1.3 准备测试文档

在代码目录下创建 docs 文件夹,放入一个测试文件 company_faq.md

1
2
3
4
5
6
7
8
9
10
11
12
13
# 公司常见问题

## 年假制度
入职满一年后享受 5 天带薪年假,满三年后 8 天,满五年后 10 天。年假需提前 3 天在 OA 系统提交申请,由直属领导审批。未休完的年假可在次年 3 月底前补休,过期作废。

## 报销流程
差旅报销需在出差结束后 5 个工作日内提交,附上发票原件和出差审批单。单笔报销金额超过 5000 元需要部门总监审批。报销款一般在提交后 10 个工作日内打到工资卡。

## 技术栈
公司后端主要使用 Python 和 Go,前端使用 React + TypeScript。数据库以 PostgreSQL 为主,缓存用 Redis,消息队列用 RabbitMQ。部署在阿里云 K8s 集群上,CI/CD 用 GitLab CI。

## 产品架构
核心产品分为三个模块:数据采集层(支持 API、SDK、文件上传三种接入方式)、数据处理层(实时流处理 + 离线批处理双引擎)、数据展示层(可视化看板 + 报表导出)。

💡 提示:这里用最简单的 Markdown 文件做演示。实际场景中你的文档可能是 PDF、Word、HTML 等格式,后面会讲到不同格式的处理方式。


二、步骤 1:加载文档与文本分块

这一步做两件事:把文件读进来,然后切成小块。

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
36
37
38
39
40
41
42
43
import os

def load_documents(docs_dir: str) -> list[str]:
"""加载目录下所有 .md 文件,返回文本内容列表"""
documents = []
for filename in os.listdir(docs_dir):
if filename.endswith(".md"):
filepath = os.path.join(docs_dir, filename)
with open(filepath, "r", encoding="utf-8") as f:
documents.append(f.read())
return documents


def split_text(text: str, chunk_size: int = 300, overlap: int = 50) -> list[str]:
"""
固定长度分块:按字符数切分,相邻块之间保留重叠
- chunk_size: 每块的最大字符数
- overlap: 相邻块之间的重叠字符数
"""
chunks = []
start = 0
while start < len(text):
end = start + chunk_size
chunk = text[start:end]
chunks.append(chunk)
# 下一块的起始位置 = 当前结束位置 - 重叠长度
start = end - overlap
return chunks


# ===== 执行 =====
docs = load_documents("docs")
print(f"加载了 {len(docs)} 个文档")

all_chunks = []
for doc in docs:
chunks = split_text(doc)
all_chunks.extend(chunks)

print(f"总共切分为 {len(all_chunks)} 个文本块")
for i, chunk in enumerate(all_chunks[:3]):
print(f"\n--- 块 {i+1} ---")
print(chunk[:100] + "...")

运行结果:

1
2
3
4
5
6
7
8
加载了 1 个文档
总共切分为 7 个文本块

--- 块 1 ---
# 公司常见问题

## 年假制度
入职满一年后享受 5 天带薪年假,满三年后 8 天,满五年后 10 天。年假需提前 3 天在 OA 系统提交申请,由直属领导审批。未休完的年假可在次年 3 月底前补休,过期作废...

关键点

  • chunk_size=300 表示每块最多 300 个字符。实际项目中通常设为 500~1000,这里为了演示效果切得碎一些
  • overlap=50 表示相邻块有 50 个字符重叠,防止关键信息刚好被切断
  • 这是最简单的固定长度分块,生产环境通常会用语义分块(后面第 9 篇会讲)

三、步骤 2:调用 Embedding 生成向量

把每个文本块转成向量。这里用 sentence-transformers 加载 bge-large-zh-v1.5 模型。

1
2
3
4
5
6
7
8
9
10
11
from sentence_transformers import SentenceTransformer

# 加载 Embedding 模型(首次运行会自动下载,约 1.2GB)
embedding_model = SentenceTransformer("BAAI/bge-large-zh-v1.5")

# 批量生成向量
embeddings = embedding_model.encode(all_chunks, normalize_embeddings=True)

print(f"向量维度: {embeddings.shape}")
# 输出示例: 向量维度: (7, 1024)
# 7 个文本块,每个块 1024 维向量

normalize_embeddings=True 会把向量归一化到单位长度,这样后续用内积计算相似度等价于余弦相似度,检索更准确。

💡 提示:如果下载速度慢,可以设置 HuggingFace 镜像:

1
2
import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

四、步骤 3:存入 Chroma 向量数据库

把文本块和对应的向量一起存入 Chroma,建立索引。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
import chromadb

# 初始化 Chroma 客户端(数据存在内存中,关闭程序即丢失)
# 如果需要持久化,改为: chromadb.PersistentClient(path="./chroma_db")
client = chromadb.Client()

# 创建一个集合(Collection),类似于数据库里的一张表
collection = client.create_collection(
name="company_docs",
metadata={"description": "公司文档知识库"}
)

# 批量写入:文本块 + 向量 + 元数据
collection.add(
documents=all_chunks, # 原文
embeddings=embeddings.tolist(), # 向量(Chroma 要求 list 格式)
ids=[f"chunk_{i}" for i in range(len(all_chunks))], # 唯一 ID
metadatas=[{"source": "company_faq.md", "index": i} for i in range(len(all_chunks))] # 元数据
)

print(f"已存入 {collection.count()} 条向量数据")
# 输出: 已存入 7 条向量数据

Chroma 的核心概念

概念 说明
Client 客户端,连接 Chroma 服务
Collection 集合,存储一组相关的向量和文档
Document 原始文本内容
Embedding 文本对应的向量
Metadata 附加信息(来源文件、页码等),可用于过滤检索

💡 提示:这里用 chromadb.Client() 把数据存在内存里,适合学习和测试。生产环境用 chromadb.PersistentClient(path="./chroma_db") 持久化到磁盘。


五、步骤 4:用户提问 → 检索相似片段

当用户提问时,把问题转成向量,然后在 Chroma 中检索最相似的 Top-K 个文本块。

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
def retrieve(query: str, top_k: int = 3) -> list[dict]:
"""
检索与 query 最相似的 top_k 个文本块
返回包含文本和元数据的列表
"""
# 1. 把用户问题转成向量(必须用同一个 Embedding 模型)
query_embedding = embedding_model.encode([query], normalize_embeddings=True)

# 2. 在 Chroma 中检索
results = collection.query(
query_embeddings=query_embedding.tolist(),
n_results=top_k,
include=["documents", "metadatas", "distances"]
)

# 3. 整理检索结果
retrieved = []
for i in range(len(results["documents"][0])):
retrieved.append({
"text": results["documents"][0][i],
"metadata": results["metadatas"][0][i],
"distance": results["distances"][0][i] # 距离越小越相似
})
return retrieved


# ===== 测试检索 =====
query = "公司年假有几天?"
results = retrieve(query, top_k=3)

print(f"用户问题: {query}\n")
for i, r in enumerate(results):
print(f"[Top-{i+1}] 距离: {r['distance']:.4f}")
print(f"内容: {r['text'][:80]}...\n")

运行结果:

1
2
3
4
5
6
7
8
9
10
用户问题: 公司年假有几天?

[Top-1] 距离: 0.2103
内容: 入职满一年后享受 5 天带薪年假,满三年后 8 天,满五年后 10 天。年假需提前 3 天在 OA 系统提交...

[Top-2] 距离: 0.5841
内容: 差旅报销需在出差结束后 5 个工作日内提交,附上发票原件和出差审批单。单笔报销金额超过 5000 元需要...

[Top-3] 距离: 0.6723
内容: 核心产品分为三个模块:数据采集层(支持 API、SDK、文件上传三种接入方式)...

Top-1 精准命中了年假相关内容,距离只有 0.21;Top-2 和 Top-3 距离明显更大,说明相关性递减。向量检索确实理解了”年假”的语义,而不是简单的关键词匹配。


六、步骤 5:拼接 Prompt + 调用大模型生成回答

最后一步:把检索到的文本块拼成 Prompt,发给 LLM,让它基于这些资料回答用户的问题。

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
36
37
38
39
40
41
42
43
44
45
46
from openai import OpenAI

# 初始化 LLM 客户端(DeepSeek 兼容 OpenAI 接口)
llm_client = OpenAI(
api_key="sk-xxx", # 替换为你的 API Key
base_url="https://api.deepseek.com" # DeepSeek API 地址
)

# 如果用 OpenAI,改为:
# llm_client = OpenAI(api_key="sk-xxx")

# 如果用本地 Ollama,改为:
# llm_client = OpenAI(api_key="ollama", base_url="http://localhost:11434/v1")


def generate_answer(query: str, context_chunks: list[dict]) -> str:
"""基于检索到的上下文,调用 LLM 生成回答"""

# 1. 拼接检索到的文本块为上下文
context = "\n\n".join([f"[{i+1}] {c['text']}" for i, c in enumerate(context_chunks)])

# 2. 构造 Prompt
system_prompt = """你是一个专业的企业知识库助手。请严格根据提供的参考资料回答用户问题。
规则:
1. 只基于参考资料回答,不要编造信息
2. 如果参考资料中没有相关内容,请明确告知"根据现有资料无法回答该问题"
3. 回答简洁准确,必要时引用参考资料编号"""

user_prompt = f"""## 参考资料
{context}

## 用户问题
{query}"""

# 3. 调用 LLM
response = llm_client.chat.completions.create(
model="deepseek-chat", # DeepSeek 模型名称
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt}
],
temperature=0.3, # 低温度,减少随机性,回答更稳定
max_tokens=500
)

return response.choices[0].message.content

七、完整代码:一键运行

把上面 5 个步骤串成一个完整脚本,复制即可运行:

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
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
"""
原生 Python 实现最简 RAG —— 完整可运行版本
依赖: pip install chromadb sentence-transformers openai
"""
import os
from sentence_transformers import SentenceTransformer
import chromadb
from openai import OpenAI


# ======================== 配置区 ========================
DOCS_DIR = "docs" # 文档目录
CHUNK_SIZE = 300 # 分块大小(字符数)
CHUNK_OVERLAP = 50 # 重叠大小
TOP_K = 3 # 检索返回数量
EMBEDDING_MODEL_NAME = "BAAI/bge-large-zh-v1.5"
LLM_API_KEY = "sk-xxx" # 替换为你的 API Key
LLM_BASE_URL = "https://api.deepseek.com"
LLM_MODEL = "deepseek-chat"


# ======================== 步骤 1: 加载文档 + 分块 ========================
def load_and_chunk(docs_dir: str, chunk_size: int, overlap: int) -> list[str]:
"""加载目录下所有 .md 文件并切分为文本块"""
all_chunks = []
for filename in os.listdir(docs_dir):
if filename.endswith(".md"):
with open(os.path.join(docs_dir, filename), "r", encoding="utf-8") as f:
text = f.read()
start = 0
while start < len(text):
all_chunks.append(text[start:start + chunk_size])
start += chunk_size - overlap
return all_chunks


# ======================== 步骤 2 + 3: 向量化 + 存入 Chroma ========================
def build_index(chunks: list[str], model: SentenceTransformer) -> chromadb.Collection:
"""生成向量并存入 Chroma"""
client = chromadb.Client()
collection = client.create_collection(name="rag_docs")

embeddings = model.encode(chunks, normalize_embeddings=True)
collection.add(
documents=chunks,
embeddings=embeddings.tolist(),
ids=[f"chunk_{i}" for i in range(len(chunks))],
metadatas=[{"index": i} for i in range(len(chunks))]
)
return collection


# ======================== 步骤 4: 检索 ========================
def retrieve(collection: chromadb.Collection, model: SentenceTransformer,
query: str, top_k: int = 3) -> list[dict]:
"""检索与 query 最相似的文本块"""
query_embedding = model.encode([query], normalize_embeddings=True)
results = collection.query(
query_embeddings=query_embedding.tolist(),
n_results=top_k,
include=["documents", "distances"]
)
return [
{"text": results["documents"][0][i], "distance": results["distances"][0][i]}
for i in range(len(results["documents"][0]))
]


# ======================== 步骤 5: 生成回答 ========================
def generate(client: OpenAI, query: str, context_chunks: list[dict]) -> str:
"""拼接 Prompt 并调用 LLM 生成回答"""
context = "\n\n".join([f"[{i+1}] {c['text']}" for i, c in enumerate(context_chunks)])

response = client.chat.completions.create(
model=LLM_MODEL,
messages=[
{"role": "system", "content": (
"你是一个专业的企业知识库助手。请严格根据提供的参考资料回答用户问题。"
"如果资料中没有相关内容,请明确告知。回答简洁准确。"
)},
{"role": "user", "content": f"## 参考资料\n{context}\n\n## 用户问题\n{query}"}
],
temperature=0.3,
max_tokens=500
)
return response.choices[0].message.content


# ======================== 主流程 ========================
def main():
# 1. 加载 + 分块
print("📄 加载文档并分块...")
chunks = load_and_chunk(DOCS_DIR, CHUNK_SIZE, CHUNK_OVERLAP)
print(f" 共 {len(chunks)} 个文本块")

# 2. 加载 Embedding 模型
print("🔢 加载 Embedding 模型...")
model = SentenceTransformer(EMBEDDING_MODEL_NAME)

# 3. 构建索引
print("💾 构建向量索引...")
collection = build_index(chunks, model)
print(f" 已索引 {collection.count()} 条数据")

# 4. 初始化 LLM
llm_client = OpenAI(api_key=LLM_API_KEY, base_url=LLM_BASE_URL)

# 5. 交互问答
print("\n✅ RAG 系统就绪!输入问题开始对话,输入 quit 退出。\n")
while True:
query = input("🙋 你的问题: ").strip()
if query.lower() in ("quit", "exit", "q"):
break
if not query:
continue

# 检索
results = retrieve(collection, model, query, TOP_K)
print(f" 检索到 {len(results)} 个相关片段(最近距离: {results[0]['distance']:.4f})")

# 生成
answer = generate(llm_client, query, results)
print(f"🤖 回答: {answer}\n")


if __name__ == "__main__":
main()

运行方式:

1
2
# 确保 docs 目录下有测试文档
python main.py

八、测试效果

用三个不同类型的问题测试,观察 RAG 的表现:

测试 1:事实性问答(能命中)

1
2
3
4
5
6
🙋 你的问题: 公司年假有几天?
检索到 3 个相关片段(最近距离: 0.2103)
🤖 回答: 根据参考资料,公司年假制度如下:
- 入职满一年:5 天带薪年假
- 入职满三年:8 天带薪年假
- 入职满五年:10 天带薪年假

检索精准命中,回答完全基于资料,没有幻觉。

测试 2:技术类问答(能命中)

1
2
3
🙋 你的问题: 公司用什么数据库?
检索到 3 个相关片段(最近距离: 0.1876)
🤖 回答: 根据参考资料,公司数据库以 PostgreSQL 为主,缓存使用 Redis,消息队列使用 RabbitMQ。部署在阿里云 K8s 集群上。

测试 3:知识范围外的问题(无法命中)

1
2
3
🙋 你的问题: 公司的团建活动安排是什么?
检索到 3 个相关片段(最近距离: 0.6234)
🤖 回答: 根据现有资料无法回答该问题。参考资料中没有关于团建活动安排的相关信息。

虽然检索还是返回了 3 个结果(向量检索总会返回 Top-K),但距离明显偏大(0.62 vs 之前的 0.21)。LLM 遵守了 Prompt 约束,没有编造答案。这就是 RAG 缓解幻觉的核心机制。

💡 提示:如果你想进一步优化”拒答”效果,可以在检索后加一个距离阈值过滤——比如距离大于 0.5 的结果直接丢弃,不传给 LLM。这是后续第 9 篇优化篇会讲的技术。


九、原生方案的缺陷

跑通了基础 RAG,你会发现在实际使用中有不少问题。这些问题在用 Dify、LangGraph 等框架时都有成熟的解决方案,这也是框架存在的价值。

缺陷 具体表现 框架怎么解决
分块太粗糙 固定 300 字符切一刀,可能把一段完整的话切断 LangChain / LlamaIndex 提供语义分块、递归分块等高级策略
不支持多种文档格式 只能读 .md 文件,PDF、Word 都处理不了 框架内置多格式文档加载器(PDF、Word、HTML、CSV)
检索没有重排 只靠向量距离排序,语义理解有限 接入 Reranker 模型做二次排序,显著提升精度
没有多轮对话记忆 每次提问都是独立的,不支持追问 LangGraph 内置状态管理,LlamaIndex 有 ChatEngine
不支持混合检索 只有向量检索,缺少关键词检索互补 框架支持向量 + BM25 混合检索
没有可视化界面 纯命令行,非技术人员无法使用 Dify / RAGFlow 提供完整的 Web 管理界面
没有持久化和备份 内存模式重启就没了 Chroma 持久化 / Milvus / Qdrant 等生产级向量库

核心价值:手写这版 RAG 的意义不在于用它上生产,而在于理解框架底层做了什么。后面用 Dify 配知识库时,你知道”分块策略”在控制什么;用 LangGraph 写检索节点时,你知道向量检索和重排的关系。


踩坑记录:Chroma 和 sentence-transformers 版本冲突

现象

pip install chromadb sentence-transformers 后运行报错:

1
ImportError: cannot import name 'onnx' from 'transformers.utils'

或者 Chroma 初始化时报 ModuleNotFoundError: No module named 'onnxruntime'

原因

Chroma 默认依赖 onnxruntime 来运行内置的 Embedding 功能,但和 sentence-transformers 的依赖版本冲突。两个库各自依赖不同版本的 transformerstokenizers

解决方案

1
2
3
4
# 按顺序安装,确保版本兼容
pip install chromadb==0.5.23
pip install sentence-transformers==3.3.1
pip install openai

如果仍有问题,可以加一个 --no-deps 隔离安装:

1
2
pip install onnxruntime
pip install chromadb sentence-transformers openai

避坑建议

requirements.txt 中锁定版本号,避免后续 pip install 升级时破坏依赖关系。本篇测试通过的版本组合:chromadb==0.5.23sentence-transformers==3.3.1openai>=1.0


总结与回顾

知识点 关键要点
整体流程 加载文档 → 分块 → 向量化 → 存入 Chroma → 检索 → 拼接 Prompt → LLM 生成
分块策略 固定长度 + 重叠,chunk_size=300,overlap=50
Embedding bge-large-zh-v1.5,通过 sentence-transformers 本地加载
向量库 Chroma,内存模式,零配置启动
LLM 调用 DeepSeek API,通过 OpenAI SDK 兼容调用
核心收获 理解 RAG 底层每一步在做什么,为后续框架学习打基础
原生缺陷 分块粗糙、格式单一、无重排、无记忆、无混合检索

下篇预告

第 5 篇:Dify 社区版搭建 RAG —— 看完原生代码的“苦”,再来体验 Dify 的“甜”。Docker 一键部署社区版,零代码可视化操作,从模型接入到知识库搭建、对话测试,完整走一遍。你会发现原生手写几十行代码做的事,Dify 点几下鼠标就搞定了。


参考资料

产品 / 工具 官网 代码仓库
Chroma trychroma.com chroma-core/chroma
Sentence-Transformers sbert.net UKPLab/sentence-transformers
bge-large-zh-v1.5 HuggingFace 模型页
DeepSeek deepseek.com deepseek-ai
OpenAI Python SDK platform.openai.com openai/openai-python
Ollama ollama.com ollama/ollama

本文为「从零到落地:RAG 检索增强生成实战系列」第 4 篇,完整系列持续更新中。