检索增强生成(Retrieval-Augmented Generation, RAG)是当下最热门的 AI 应用架构之一。本文将带你从零开始,用 Python 完整实现一个 RAG 系统,涵盖文档加载、文本分块、向量嵌入、语义检索与生成回答的完整链路。
一、什么是 RAG?
RAG 的核心思想:让大语言模型(LLM)在回答问题时,先从外部知识库中检索相关内容,再基于检索结果生成回答。 这有效解决了 LLM 的幻觉问题和知识时效性问题。
RAG vs 纯 LLM 对比
| 维度 | 纯 LLM | RAG |
|---|---|---|
| 知识来源 | 训练数据(静态) | 外部知识库(动态) |
| 幻觉问题 | 严重 | 显著降低 |
| 数据更新 | 需重新训练 | 增量更新索引即可 |
| 私有数据 | 无法使用 | 完美支持 |
二、RAG 系统架构总览
┌─────────────────────────────────────────────────────────┐
│ RAG 系统架构 │
├─────────────────────────────────────────────────────────┤
│ │
│ ┌──────────┐ ┌──────────┐ ┌──────────────────┐ │
│ │ 文档加载 │───▶│ 文本分块 │───▶│ 向量嵌入(Embed) │ │
│ └──────────┘ └──────────┘ └────────┬─────────┘ │
│ │ │
│ ▼ │
│ ┌──────────────┐ │
│ 离线阶段 │ 向量数据库 │ │
│ ───────────────────────────── │ (Vector DB) │ │
│ 在线阶段 └──────┬───────┘ │
│ │ │
│ ┌──────────┐ ┌──────────┐ │ │
│ │ 用户提问 │───▶ │ Query │───▶ 检索相似文档 │ │
│ └──────────┘ │ Embedding│ │ │
│ └──────────┘ │ │
│ ▼ │
│ ┌──────────┐ ┌──────────────┐ │
│ │ LLM 生成 │◀───│ 拼接 Prompt │ │
│ └─────┬────┘ └──────────────┘ │
│ │ │
│ ▼ │
│ ┌──────────┐ │
│ │ 最终回答 │ │
│ └──────────┘ │
└─────────────────────────────────────────────────────────┘
三、环境准备
3.1 安装依赖
bash
pip install langchain langchain-community \
chromadb sentence-transformers \
pypdf unstructured \
openai tiktoken
3.2 项目结构
rag-project/
├── data/ # 存放原始文档(PDF、TXT、MD 等)
├── vector_db/ # 向量数据库持久化目录
├── main.py # 主程序
└── config.py # 配置文件
四、Step 1 ------ 文档加载
文档加载是 RAG 的起点。实际场景中,知识库可能包含 PDF、Word、Markdown、网页等多种格式。
4.1 加载 PDF 文件
python
from langchain_community.document_loaders import PyPDFLoader
def load_pdf(file_path: str):
"""加载 PDF 文件,返回 Document 列表"""
loader = PyPDFLoader(file_path)
pages = loader.load()
print(f"成功加载 {len(pages)} 页")
return pages
# 每个Document对象包含:
# - page_content: 文本内容
# - metadata: 元数据(页码、来源等)
4.2 加载 Markdown 文件
python
from langchain_community.document_loaders import UnstructuredMarkdownLoader
def load_markdown(file_path: str):
"""加载 Markdown 文件"""
loader = UnstructuredMarkdownLoader(file_path)
docs = loader.load()
print(f"成功加载 Markdown: {len(docs)} 段")
return docs
4.3 批量加载目录下所有文档
python
from langchain_community.document_loaders import DirectoryLoader
def load_directory(dir_path: str, glob_pattern: str = "**/*.pdf"):
"""批量加载目录中的文档"""
loader = DirectoryLoader(
dir_path,
glob=glob_pattern,
show_progress=True,
use_multithreading=True
)
documents = loader.load()
print(f"从 {dir_path} 加载了 {len(documents)} 个文档")
return documents
五、Step 2 ------ 文本分块(Chunking)
大文档不能整篇丢给模型,需要切分成合适大小的片段。
5.1 为什么需要分块?
原始文档(可能 100+ 页)
│
▼ 切分
┌──────┐┌──────┐┌──────┐┌──────┐
│Chunk1││Chunk2││Chunk3││Chunk4│ ... 每块 500~1000 tokens
└──────┘└──────┘└──────┘└──────┘
│ │ │ │
▼ ▼ ▼ ▼
向量化 向量化 向量化 向量化 → 存入向量数据库
5.2 递归字符分块器(推荐)
python
from langchain.text_splitter import RecursiveCharacterTextSplitter
def split_documents(documents, chunk_size=500, chunk_overlap=50):
"""
递归字符分块器 ------ 按段落、句子边界智能切分
参数:
chunk_size: 每块最大字符数
chunk_overlap: 相邻块重叠字符数(保持上下文连贯)
"""
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
separators=["\n\n", "\n", "。", "!", "?", ".", " ", ""],
length_function=len
)
chunks = text_splitter.split_documents(documents)
print(f"文档被切分为 {len(chunks)} 个文本块")
return chunks
5.3 分块策略选择指南
┌─────────────────────────────────────────────────┐
│ 分块策略选择 │
├─────────────────┬───────────────────────────────┤
│ 策略 │ 适用场景 │
├─────────────────┼───────────────────────────────┤
│ 固定大小分块 │ 通用场景,简单快速 │
│ 递归字符分块 │ ✅ 推荐,兼顾语义与效率 │
│ 按语义分块 │ 对语义完整性要求高 │
│ 按文档结构分块 │ Markdown / HTML 等结构化文档 │
└─────────────────┴───────────────────────────────┘
六、Step 3 ------ 向量嵌入(Embedding)
将文本块转换为高维向量,使其可被计算机进行相似度计算。
6.1 使用本地开源模型(免费)
python
from langchain_community.embeddings import HuggingFaceEmbeddings
def get_embedding_model():
"""使用本地 Sentence Transformer 模型"""
model = HuggingFaceEmbeddings(
model_name="BAAI/bge-small-zh-v1.5", # 中文优秀模型
model_kwargs={"device": "cpu"},
encode_kwargs={"normalize_embeddings": True}
)
return model
# 测试嵌入效果
model = get_embedding_model()
vector = model.embed_query("什么是机器学习?")
print(f"向量维度: {len(vector)}") # 输出: 向量维度: 384
6.2 使用 OpenAI Embedding(付费,效果好)
python
from langchain_community.embeddings import OpenAIEmbeddings
def get_openai_embedding():
"""使用 OpenAI text-embedding-3-small"""
return OpenAIEmbeddings(
model="text-embedding-3-small",
openai_api_key="your-api-key"
)
七、Step 4 ------ 构建向量数据库
7.1 使用 Chroma(轻量级本地向量数据库)
python
from langchain_community.vectorstores import Chroma
def build_vector_store(chunks, embedding_model, persist_directory="./vector_db"):
"""
构建向量数据库并持久化
流程: 文本块 → Embedding → 存入 Chroma
"""
vectorstore = Chroma.from_documents(
documents=chunks,
embedding=embedding_model,
persist_directory=persist_directory
)
vectorstore.persist()
print(f"向量数据库构建完成,共 {vectorstore._collection.count()} 条记录")
return vectorstore
def load_vector_store(embedding_model, persist_directory="./vector_db"):
"""加载已有的向量数据库"""
vectorstore = Chroma(
persist_directory=persist_directory,
embedding_function=embedding_model
)
return vectorstore
7.2 向量数据库选型对比
┌────────────┬──────────┬───────────┬──────────────────────┐
│ 数据库 │ 部署方式 │ 适合场景 │ 特点 │
├────────────┼──────────┼───────────┼──────────────────────┤
│ Chroma │ 本地 │ 开发/小项目 │ 轻量,Python 原生 │
│ FAISS │ 本地 │ 高性能检索 │ Meta 开源,速度快 │
│ Milvus │ 分布式 │ 生产环境 │ 可扩展,支持亿级向量 │
│ Pinecone │ 云服务 │ 免运维 │ 全托管,按量付费 │
│ Qdrant │ 独立部署 │ 中大型项目 │ Rust 编写,性能优秀 │
└────────────┴──────────┴───────────┴──────────────────────┘
八、Step 5 ------ 语义检索
8.1 基础相似度检索
python
def similarity_search(vectorstore, query: str, k: int = 4):
"""
基础语义检索 ------ 返回最相似的 k 个文本块
原理: 将 query 向量化 → 与数据库中所有向量计算余弦相似度 → 返回 Top-K
"""
results = vectorstore.similarity_search(
query=query,
k=k
)
for i, doc in enumerate(results, 1):
print(f"\n--- 检索结果 {i} ---")
print(f"内容: {doc.page_content[:200]}...")
print(f"来源: {doc.metadata.get('source', '未知')}")
return results
8.2 带分数的相似度检索(MMR 多样性检索)
python
def mmr_search(vectorstore, query: str, k: int = 4, fetch_k: int = 20):
"""
MMR(最大边际相关性)检索
在保证相关性的同时,尽量减少结果之间的冗余
相比普通检索:
- 普通检索可能返回内容高度重复的多个结果
- MMR 检索能返回既相关又多样化的结果
"""
results = vectorstore.max_marginal_relevance_search(
query=query,
k=k,
fetch_k=fetch_k
)
return results
8.3 检索策略对比
python
# ===== 三种检索方式对比 =====
# 1. 纯相似度检索 ------ 最简单,可能冗余
docs1 = vectorstore.similarity_search("什么是深度学习?", k=4)
# 2. 相似度 + 分数 ------ 可按阈值过滤
docs2 = vectorstore.similarity_search_with_relevance_scores(
"什么是深度学习?", k=4
)
for doc, score in docs2:
print(f"分数: {score:.4f} | {doc.page_content[:80]}")
# 3. MMR 检索 ------ 相关性 + 多样性兼顾(推荐)
docs3 = vectorstore.max_marginal_relevance_search(
"什么是深度学习?", k=4, fetch_k=20
)
九、Step 6 ------ 拼接 Prompt 并调用 LLM 生成回答
9.1 构建 RAG Prompt 模板
python
from langchain.prompts import ChatPromptTemplate
RAG_PROMPT_TEMPLATE = """你是一个专业的知识助手。请根据以下检索到的参考资料来回答用户的问题。
要求:
- 只基于参考资料回答,不要编造信息
- 如果参考资料中没有相关内容,请诚实说明
- 回答要条理清晰,必要时使用列表或分点说明
参考资料:
{context}
用户问题:{question}
请给出你的回答:"""
def build_rag_prompt(query: str, retrieved_docs: list) -> str:
"""将检索到的文档拼接为 context,构建完整 Prompt"""
context = "\n\n".join([
f"[来源 {i+1}]: {doc.page_content}"
for i, doc in enumerate(retrieved_docs)
])
prompt = ChatPromptTemplate.from_template(RAG_PROMPT_TEMPLATE)
return prompt.format(context=context, question=query)
9.2 调用 LLM 生成回答
python
from langchain_community.chat_models import ChatOpenAI
def generate_answer(query: str, retrieved_docs: list, llm=None):
"""调用 LLM 生成最终回答"""
if llm is None:
llm = ChatOpenAI(
model="gpt-4o-mini",
temperature=0,
openai_api_key="your-api-key"
)
prompt = build_rag_prompt(query, retrieved_docs)
response = llm.invoke(prompt)
return response.content
十、完整 RAG 流程整合
python
"""
完整的 RAG 系统 ------ 从文档到问答
"""
from langchain_community.document_loaders import DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.prompts import ChatPromptTemplate
from langchain_community.chat_models import ChatOpenAI
class SimpleRAG:
"""一个简单但完整的 RAG 系统"""
def __init__(self, docs_dir: str = "./data", db_dir: str = "./vector_db"):
self.docs_dir = docs_dir
self.db_dir = db_dir
self.embedding_model = HuggingFaceEmbeddings(
model_name="BAAI/bge-small-zh-v1.5",
encode_kwargs={"normalize_embeddings": True}
)
self.vectorstore = None
self.llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
# ---------- 离线阶段:构建知识库 ----------
def build_index(self):
"""Step 1~4: 加载文档 → 分块 → 嵌入 → 存入向量数据库"""
print("=" * 50)
print("Step 1: 加载文档...")
loader = DirectoryLoader(self.docs_dir, glob="**/*.pdf", show_progress=True)
documents = loader.load()
print(f" 加载了 {len(documents)} 个文档")
print("Step 2: 文本分块...")
splitter = RecursiveCharacterTextSplitter(
chunk_size=500, chunk_overlap=50,
separators=["\n\n", "\n", "。", ".", " ", ""]
)
chunks = splitter.split_documents(documents)
print(f" 切分为 {len(chunks)} 个文本块")
print("Step 3~4: 向量嵌入并存储...")
self.vectorstore = Chroma.from_documents(
documents=chunks,
embedding=self.embedding_model,
persist_directory=self.db_dir
)
self.vectorstore.persist()
print(f" 向量数据库构建完成!")
print("=" * 50)
def load_index(self):
"""加载已有的向量数据库"""
self.vectorstore = Chroma(
persist_directory=self.db_dir,
embedding_function=self.embedding_model
)
print("已加载向量数据库")
# ---------- 在线阶段:检索 + 生成 ----------
def query(self, question: str, k: int = 4) -> str:
"""完整 RAG 问答流程"""
# Step 5: 语义检索
retrieved = self.vectorstore.max_marginal_relevance_search(
query=question, k=k, fetch_k=k*5
)
# Step 6: 拼接 Prompt + LLM 生成
context = "\n\n".join([
f"[来源 {i+1}]: {doc.page_content}"
for i, doc in enumerate(retrieved)
])
prompt = f"""基于以下参考资料回答问题。如果资料中没有答案,请说明。
参考资料:
{context}
问题:{question}"""
response = self.llm.invoke(prompt)
return response.content
# ===== 使用示例 =====
if __name__ == "__main__":
rag = SimpleRAG(docs_dir="./data", db_dir="./vector_db")
# 首次运行:构建索引
# rag.build_index()
# 后续运行:直接加载
rag.load_index()
# 提问
answer = rag.query("什么是深度学习?它有哪些主要应用?")
print(f"\n回答:\n{answer}")
十一、完整数据流图
用户提问: "什么是深度学习?"
│
▼
┌──────────────────┐
│ Query Embedding │ "什么是深度学习?" → [0.023, -0.045, 0.078, ...]
└────────┬─────────┘
│
▼
┌──────────────────────────────────────────────┐
│ 向量数据库 (Chroma) │
│ │
│ Doc1: "深度学习是机器学习的分支..." → [0.02, -0.04, 0.08, ...] ✅ 相似度 0.92
│ Doc2: "神经网络通过反向传播..." → [0.01, -0.03, 0.06, ...] ✅ 相似度 0.87
│ Doc3: "CNN 在图像识别中..." → [0.03, -0.02, 0.05, ...] ✅ 相似度 0.84
│ Doc4: "Python 是一种编程语言..." → [0.01, 0.05, -0.02, ...] ❌ 相似度 0.31
│ ... │
└──────────────────────────────────────────────┘
│
│ Top-K 检索结果 (k=3)
▼
┌──────────────────────────────────────────┐
│ 拼接 Prompt │
│ │
│ System: 你是一个知识助手... │
│ │
│ Context: │
│ [来源1]: 深度学习是机器学习的分支... │
│ [来源2]: 神经网络通过反向传播... │
│ [来源3]: CNN 在图像识别中... │
│ │
│ Question: 什么是深度学习? │
└────────────────┬─────────────────────────┘
│
▼
┌────────────────────────┐
│ LLM (GPT-4o-mini) │
└────────────────┬───────┘
│
▼
┌────────────────────────────────────────────────────┐
│ 回答: │
│ 深度学习是机器学习的一个重要分支,基于人工神经网络...│
│ 主要应用包括: │
│ 1. 图像识别(CNN) │
│ 2. 自然语言处理(Transformer) │
│ 3. 语音识别 ... │
└────────────────────────────────────────────────────┘
十二、进阶优化方向
构建完基础 RAG 后,可以从以下方向持续优化:
┌─────────────────────────────────────────────────────┐
│ RAG 优化路线图 │
│ │
│ 基础 RAG ──▶ 优化检索 ──▶ 优化生成 ──▶ 高级架构 │
│ │
│ 📄 文档处理优化 │
│ ├── 更智能的分块策略(语义分块) │
│ ├── 表格 / 图片内容提取 │
│ └── OCR 处理扫描件 │
│ │
│ 🔍 检索优化 │
│ ├── 混合检索(向量 + 关键词 BM25) │
│ ├── 重排序(Reranker / Cross-Encoder) │
│ ├── Query 改写与扩展 │
│ └── 元数据过滤 │
│ │
│ 🤖 生成优化 │
│ ├── 更好的 Prompt 工程 │
│ ├── 引用溯源(标注来源段落) │
│ └── 自适应温度参数 │
│ │
│ 🏗️ 高级架构 │
│ ├── Agentic RAG(带工具调用的智能体) │
│ ├── Multi-modal RAG(图文混合) │
│ ├── GraphRAG(知识图谱增强) │
│ └── 评估框架(RAGAS) │
└─────────────────────────────────────────────────────┘
12.1 混合检索示例
python
from langchain.retrievers import BM25Retriever, EnsembleRetriever
def build_hybrid_retriever(chunks, vectorstore, k=4):
"""
混合检索 = 向量检索(语义) + BM25检索(关键词)
取长补短,效果优于单一检索
"""
# 关键词检索(BM25)
bm25_retriever = BM25Retriever.from_documents(chunks)
bm25_retriever.k = k
# 语义检索(向量)
vector_retriever = vectorstore.as_retriever(
search_kwargs={"k": k}
)
# 混合检索器(各占 50% 权重)
ensemble_retriever = EnsembleRetriever(
retrievers=[bm25_retriever, vector_retriever],
weights=[0.5, 0.5]
)
return ensemble_retriever
12.2 检索结果重排序
python
from sentence_transformers import CrossEncoder
def rerank_results(query: str, documents: list, top_k: int = 4) -> list:
"""
使用 Cross-Encoder 对检索结果重排序
Cross-Encoder 比 Bi-Encoder 更精确,但速度较慢
适合对 Top-K 候选做精排
"""
reranker = CrossEncoder("BAAI/bge-reranker-base")
pairs = [[query, doc.page_content] for doc in documents]
scores = reranker.predict(pairs)
# 按分数降序排列
ranked = sorted(
zip(documents, scores),
key=lambda x: x[1],
reverse=True
)
return [doc for doc, score in ranked[:top_k]]
十三、总结
本文完整实现了一个 RAG 系统,核心流程回顾:
📄 文档加载 → ✂️ 文本分块 → 🔢 向量嵌入 → 💾 存入向量库
│
🎤 用户提问 → 🔍 语义检索 → 📝 拼接Prompt → 🤖 LLM生成回答
关键要点:
- 分块质量决定检索上限 ------ 注意 chunk_size 和 overlap 的调参
- Embedding 模型 决定语义理解质量 ------ 中文场景推荐
bge系列 - 检索策略影响召回率 ------ 推荐混合检索 + 重排序
- Prompt 工程影响最终输出质量 ------ 明确指令,约束幻觉
