AI 应用开发实战（4）：从零实现 RAG 系统

前言

前几篇我们搭建了 AI 应用的基础骨架、Prompt 工程体系和多轮对话管理。这一篇进入 RAG（检索增强生成）——让 AI 拥有自己的知识库，能基于私有文档回答问题。

市面上讲 RAG 的文章很多，但要么太理论（只讲概念），要么太碎片（只讲某个环节）。这篇的目标是：从零实现一个完整的、可运行的 RAG 系统。

最终效果：上传 PDF/文档，自动切分、向量化、存入向量库，用户提问时检索相关片段并生成回答。

1. RAG 整体架构

用户问题
│
▼
┌──────────────┐ ┌────────────────┐
│ Query 改写 │───▶│ Embedding 模型 │
│ 改写/扩展 │ │ 向量化问题 │
└──────────────┘ └───────┬────────┘
│
▼
┌──────────────────────┐
│ 向量数据库 │
│ (Qdrant/Milvus) │
│ → Top-K 检索 │
│ → 混合检索(可选) │
└───────┬──────────────┘
│
▼
┌──────────────────────┐
│ Re-rank │
│ 重排序 Top-K 结果 │
└───────┬──────────────┘
│
▼
┌──────────────────────┐
│ LLM + Prompt │
│ 基于资料生成回答 │
└───────┬──────────────┘
│
▼
最终答案

2. 文档解析与切分

2.1 文档解析

先支持最常见的 PDF 和 TXT 格式：

import os
from pathlib import Path

def parse_document(file_path: str) –> str:
"""解析文档，返回纯文本内容。"""
ext = Path(file_path).suffix.lower()

if ext == ".txt":
with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
return f.read()

elif ext == ".pdf":
return parse_pdf(file_path)

elif ext == ".md":
with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
return f.read()

else:
raise ValueError(f"Unsupported format: {ext}")

def parse_pdf(file_path: str) –> str:
"""解析 PDF 文件。使用 PyMuPDF（fitz）实现。"""
try:
import fitz # PyMuPDF
except ImportError:
raise ImportError("Please install PyMuPDF: pip install pymupdf")

doc = fitz.open(file_path)
text = []
for page in doc:
text.append(page.get_text())
doc.close()
return "\\n\\n".join(text)

2.2 文本切分（Chunking）

切分是 RAG 效果的基础。三种常用策略：

策略 1：递归字符切分（推荐）

from typing import List

class RecursiveCharacterTextSplitter:
"""递归字符文本切分器。

先按段落切，段落太长再按句子切，句子太长再按固定长度切。
"""

def __init__(
self,
chunk_size: int = 512,
chunk_overlap: int = 128,
separators: List[str] = None,
):
self.chunk_size = chunk_size
self.chunk_overlap = chunk_overlap
self.separators = separators or ["\\n\\n", "\\n", "。", "。", "，", " ", ""]

def split_text(self, text: str) –> List[str]:
"""递归切分文本。"""
return self._split(text, self.separators, 0)

def _split(self, text: str, separators: List[str], level: int) –> List[str]:
if level >= len(separators):
# 没有更多分隔符了，按固定长度切
return self._split_fixed_length(text)

separator = separators[level]
if not separator:
return self._split_fixed_length(text)

chunks = []
# 用当前分隔符切分
for segment in text.split(separator):
segment = segment.strip()
if not segment:
continue

if len(segment) <= self.chunk_size:
chunks.append(segment)
else:
# 太长，递归到下一级分隔符
sub_chunks = self._split(segment, separators, level + 1)
chunks.extend(sub_chunks)

# 合并太小的 chunk，处理 overlap
return self._merge_chunks(chunks)

def _split_fixed_length(self, text: str) –> List[str]:
"""固定长度切分（按字符数）。"""
chunks = []
start = 0
while start < len(text):
end = min(start + self.chunk_size, len(text))
if end < len(text):
# 在边界附近找最后一个换行或句号
cut = text.rfind("\\n", start, end)
if cut == –1 or cut < start + self.chunk_size // 2:
cut = text.rfind("。", start, end)
if cut == –1 or cut < start + self.chunk_size // 2:
cut = end
else:
cut += 1
else:
cut += 1
else:
cut = end

chunks.append(text[start:cut].strip())
start = cut – self.chunk_overlap # overlap

return [c for c in chunks if c]

def _merge_chunks(self, chunks: List[str]) –> List[str]:
"""合并过小的 chunks。"""
merged = []
buffer = ""
for c in chunks:
if len(buffer) + len(c) <= self.chunk_size:
buffer = (buffer + "\\n\\n" + c).strip() if buffer else c
else:
if buffer:
merged.append(buffer)
buffer = c
if buffer:
merged.append(buffer)
return merged

经验值：

• 中文场景：chunk_size=512, overlap=128
• 英文场景：chunk_size=1024, overlap=256
• 代码文档：chunk_size=256, overlap=64

策略 2：语义切分

按段落、标题等语义边界切分，适合结构化文档：

def split_by_semantic(text: str) –> List[str]:
"""按 Markdown 标题或空行切分。"""
import re
# 先按标题切
sections = re.split(r"\\n#{1,3}\\s+", text)
result = []
for s in sections:
s = s.strip()
if not s:
continue
if len(s) > 1500:
# 太长的再按空行切
paras = s.split("\\n\\n")
for p in paras:
p = p.strip()
if p:
result.append(p)
else:
result.append(s)
return result

2.3 生成元数据

每个 chunk 需要附带元数据，方便溯源和过滤：

def create_chunks_with_metadata(
text: str,
source: str,
chunk_size: int = 512,
chunk_overlap: int = 128,
) –> List[dict]:
"""切分文本并生成带元数据的 chunks。"""
splitter = RecursiveCharacterTextSplitter(
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
)
texts = splitter.split_text(text)

chunks = []
for i, content in enumerate(texts):
chunks.append({
"id": f"{source}#chunk{i}",
"text": content,
"metadata": {
"source": source,
"chunk_index": i,
"total_chunks": len(texts),
},
})
return chunks

3. Embedding 与向量存储

3.1 Embedding 封装

import numpy as np
from typing import List

class EmbeddingModel:
"""Embedding 模型封装，支持多种后端。"""

def __init__(self, model_type: str = "local", model_name: str = None):
self.model_type = model_type
self.model_name = model_name
self._model = None

if model_type == "local":
self._init_local()
elif model_type == "api":
self._init_api()

def _init_local(self):
"""使用本地 Sentence-Transformers 模型。"""
try:
from sentence_transformers import SentenceTransformer
name = self.model_name or "BAAI/bge-small-zh-v1.5"
self._model = SentenceTransformer(name)
self.dimension = self._model.get_sentence_embedding_dimension()
print(f"Loaded local model: {name} (dim={self.dimension})")
except ImportError:
raise ImportError("Install: pip install sentence-transformers")

def _init_api(self):
"""使用 API 版 Embedding 模型。"""
from openai import OpenAI
import os
self._client = OpenAI(
api_key=os.getenv("LLM_API_KEY"),
base_url=os.getenv("LLM_BASE_URL"),
)
self.dimension = 1024 # bge-m3 dimension

def encode(self, texts: List[str]) –> List[List[float]]:
"""将文本列表编码为向量。"""
if self.model_type == "local":
emb = self._model.encode(texts, normalize_embeddings=True)
return emb.tolist()
else:
response = self._client.embeddings.create(
model=self.model_name or "BAAI/bge-m3",
input=texts,
)
return [d.embedding for d in response.data]

def encode_query(self, text: str) –> List[float]:
"""编码查询（单条）。"""
if self.model_type == "local":
emb = self._model.encode([text], normalize_embeddings=True)
return emb[0].tolist()
else:
return self.encode([text])[0]

Embedding 选型建议：

3.2 向量存储（内存版，适合原型）

import numpy as np
from typing import List, Dict, Optional
import json

class VectorStore:
"""内存向量存储（原型用，生产环境用 Qdrant/Milvus）。"""

def __init__(self, dimension: int = 512):
self.dimension = dimension
self.vectors = [] # List[np.array]
self.documents = [] # List[dict]
self.index = None # FAISS index

def add_documents(self, documents: List[dict]):
"""添加文档到向量库。"""
for doc in documents:
self.documents.append({
"id": doc["id"],
"text": doc["text"],
"metadata": doc.get("metadata", {}),
})

def add_vectors(self, vectors: List[List[float]], documents: List[dict]):
"""添加向量和文档。"""
for vec, doc in zip(vectors, documents):
self.vectors.append(np.array(vec, dtype=np.float32))
self.documents.append({
"id": doc["id"],
"text": doc["text"],
"metadata": doc.get("metadata", {}),
})

def build_index(self):
"""构建 FAISS 索引加速检索。"""
try:
import faiss
except ImportError:
raise ImportError("Install faiss: pip install faiss-cpu")

if not self.vectors:
return

matrix = np.array(self.vectors).astype(np.float32)
self.index = faiss.IndexFlatIP(self.dimension) # 内积 = 余弦相似度
faiss.normalize_L2(matrix)
self.index.add(matrix)
print(f"Index built: {self.index.ntotal} vectors")

def search(
self,
query_vector: List[float],
k: int = 5,
score_threshold: float = 0.0,
filter_expr: Optional[dict] = None,
) –> List[dict]:
"""检索 Top-K 相似文档。"""
query = np.array([query_vector]).astype(np.float32)
faiss.normalize_L2(query)

if self.index is None:
self.build_index()

if self.index is None:
return []

scores, indices = self.index.search(query, k * 2)

results = []
for score, idx in zip(scores[0], indices[0]):
if idx < 0 or idx >= len(self.documents):
continue
if score < score_threshold:
continue

doc = dict(self.documents[idx])
doc["score"] = float(score)

# 应用过滤器
if filter_expr:
meta = doc.get("metadata", {})
match = all(
meta.get(k) == v for k, v in filter_expr.items()
)
if not match:
continue

results.append(doc)
if len(results) >= k:
break

return results

4. Query 改写与检索增强

4.1 Query 改写

用户的原始问题往往不适合直接检索。改写可以大幅提升检索效果：

class QueryRewriter:
"""Query 改写引擎。"""

def __init__(self, llm_func):
self.llm_func = llm_func # chat_sync or similar
self._prompt = """你是一个检索专家。用户的问题可能不够清晰，
你需要将其改写成更适合向量检索的形式。

要求：
– 提取核心关键词和实体
– 补充同义词或相关术语
– 保持简洁，不超过 50 个字
– 如果问题已经清晰，直接返回原问题

用户问题：{question}

改写后的检索查询："""

def rewrite(self, question: str) –> str:
"""改写问题。"""
reply = self.llm_func([
{"role": "system", "content": self._prompt.format(question=question)}
])
return reply.strip()

def expand(self, question: str) –> list:
"""生成多个检索查询（HyDE 方法）。"""
prompt = """用户提问：{question}

请生成 3 个不同的检索查询，每个一行。
这些查询应该从不同角度覆盖用户的问题。
"""
reply = self.llm_func([
{"role": "system", "content": prompt.format(question=question)}
])
queries = [q.strip() for q in reply.split("\\n") if q.strip()]
return queries[:3]

4.2 混合检索

class HybridRetriever:
"""混合检索：向量检索 + 关键词检索（BM25）。"""

def __init__(self, vector_store: VectorStore, alpha: float = 0.5):
self.vector_store = vector_store
self.alpha = alpha # 向量检索权重，1-alpha 为 BM25 权重
self._bm25 = None

def _build_bm25(self, texts: List[str]):
"""构建 BM25 索引。"""
try:
from rank_bm25 import BM25Okapi
except ImportError:
raise ImportError("Install: pip install rank-bm25")

tokenized = [t.split() for t in texts]
self._bm25 = BM25Okapi(tokenized)

def search(
self,
query_vector: List[float],
query_text: str,
k: int = 5,
) –> List[dict]:
"""混合检索。"""
# 向量检索
vec_results = self.vector_store.search(query_vector, k=k * 2)

# BM25 检索
if self._bm25:
tokenized_query = query_text.split()
bm25_scores = self._bm25.get_scores(tokenized_query)
docs = self.vector_store.documents
bm25_results = sorted(
[
(i, bm25_scores[i])
for i in range(len(bm25_scores))
],
key=lambda x: –x[1],
)[:k]

# RRF 合并
rrf_scores = {}
for rank, r in enumerate(vec_results):
doc_id = r.get("id", str(rank))
rrf_scores[doc_id] = rrf_scores.get(doc_id, 0) + 1 / (60 + rank + 1)

if self._bm25:
for rank, (idx, score) in enumerate(bm25_results):
doc_id = docs[idx].get("id", str(idx))
rrf_scores[doc_id] = rrf_scores.get(doc_id, 0) + 1 / (60 + rank + 1)

sorted_results = sorted(rrf_scores.items(), key=lambda x: –x[1])
top_ids = [doc_id for doc_id, _ in sorted_results[:k]]

# 返回结果
doc_map = {d.get("id", str(i)): d for i, d in enumerate(self.vector_store.documents)}
return [doc_map[doc_id] for doc_id in top_ids if doc_id in doc_map]

5. Re-ranking

检索到的 Top-K 结果需要用 Cross-Encoder 重排序，能大幅提升最终回答质量：

class ReRanker:
"""用 Cross-Encoder 重排序检索结果。"""

def __init__(self, model_name: str = "BAAI/bge-reranker-v2-m3"):
try:
from transformers import AutoModelForSequenceClassification, AutoTokenizer
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
self.model.eval()
except Exception as e:
print(f"Re-ranker not available: {e}")
self.model = None

def rerank(self, query: str, documents: List[dict], top_k: int = 5) –> List[dict]:
"""重排序并返回 Top-K。"""
if not self.model or not documents:
return documents[:top_k]

pairs = [[query, doc["text"][:512]] for doc in documents]
inputs = self.tokenizer(
pairs, padding=True, truncation=True,
max_length=512, return_tensors="pt",
)

import torch
with torch.no_grad():
scores = self.model(**inputs).logits.squeeze(–1).tolist()

if isinstance(scores, float):
scores = [scores]

for doc, score in zip(documents, scores):
doc["rerank_score"] = score

documents.sort(key=lambda x: x.get("rerank_score", 0), reverse=True)
return documents[:top_k]

为什么需要 Re-rank？

Bi-Encoder 把文本独立编码为向量，速度快但精度有限。Cross-Encoder 把 query 和 doc 一起送入模型计算相似度，精度高但慢。所以标准做法是：先用 Bi-Encoder 检索 Top-50，再用 Cross-Encoder 重排取 Top-5。

6. 完整 RAG Pipeline

class RAGPipeline:
"""完整的 RAG Pipeline。"""

def __init__(
self,
embedding_model: EmbeddingModel,
llm_func,
vector_store: VectorStore = None,
):
self.embedder = embedding_model
self.llm = llm_func
self.vector_store = vector_store or VectorStore(embedding_model.dimension)
self.rewriter = QueryRewriter(llm_func)
self.reranker = ReRanker()
self._rag_prompt = """你是一个基于知识库的 AI 助手。

请根据以下检索到的资料来回答用户的问题。
如果资料不足以回答问题，请直接说"资料中没有相关信息"，不要编造。

## 检索资料

{context}

## 用户问题

{question}

请基于以上资料回答，并在回答末尾注明引用的资料来源。"""

def ingest(self, file_path: str):
"""导入文档到知识库。"""
# 1. 解析
text = parse_document(file_path)
source = os.path.basename(file_path)

# 2. 切分
chunks = create_chunks_with_metadata(text, source)

# 3. 向量化
texts = [c["text"] for c in chunks]
vectors = self.embedder.encode(texts)

# 4. 存入向量库
self.vector_store.add_vectors(vectors, chunks)
print(f"Ingested {len(chunks)} chunks from {source}")

def query(self, question: str, k: int = 5) –> dict:
"""回答用户问题。"""
# 1. Query 改写
rewritten = self.rewriter.rewrite(question)
print(f" Rewritten: {rewritten}")

# 2. 向量化问题
query_vector = self.embedder.encode_query(rewritten)

# 3. 检索
results = self.vector_store.search(
query_vector, k=k * 3,
)

# 4. Re-rank
results = self.reranker.rerank(rewritten, results, top_k=k)

# 5. 构建 Prompt
context = "\\n\\n—\\n\\n".join([
f"[来源 {r.get('metadata', {}).get('source', 'unknown')}] {r['text']}"
for r in results
])

prompt = self._rag_prompt.format(context=context, question=question)

# 6. 生成回答
answer = self.llm([
{"role": "system", "content": prompt}
])

return {
"answer": answer,
"sources": [
{
"source": r.get("metadata", {}).get("source", "unknown"),
"text": r["text"][:200],
"score": r.get("score", 0),
}
for r in results
],
}

7. 集成到之前的 AI 应用

在系列第一篇的 ai-chat-app 中集成 RAG：

# rag_app.py – 在 AI Chat 中添加 RAG 知识库功能

from flask import Flask, request, jsonify
from rag_pipeline import RAGPipeline
from llm import chat_sync

app = Flask(__name__)
rag = RAGPipeline(
embedding_model=EmbeddingModel(model_type="local", model_name="BAAI/bge-small-zh-v1.5"),
llm_func=chat_sync,
)

UPLOAD_FOLDER = "knowledge_base"
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

@app.route("/api/rag/upload", methods=["POST"])
def upload_document():
"""上传文档到知识库。"""
if "file" not in request.files:
return jsonify({"error": "No file"}), 400

file = request.files["file"]
file_path = os.path.join(UPLOAD_FOLDER, file.filename)
file.save(file_path)

try:
rag.ingest(file_path)
return jsonify({"status": "ok", "message": f"Imported {file.filename}"})
except Exception as e:
return jsonify({"error": str(e)}), 500

@app.route("/api/rag/query", methods=["POST"])
def rag_query():
"""基于知识库回答问题。"""
data = request.json
question = data.get("question", "")

if not question:
return jsonify({"error": "Question required"}), 400

result = rag.query(question)

return jsonify({
"answer": result["answer"],
"sources": result["sources"],
"mode": "rag",
})

@app.route("/api/rag/status", methods=["GET"])
def get_status():
"""查看知识库状态。"""
doc_count = len(rag.vector_store.documents)
return jsonify({
"document_count": doc_count,
"status": "ready" if doc_count > 0 else "empty",
})

8. 效果调优清单

当 RAG 效果不好时，按这个清单排查：

回答不准确？
├─ 检索不到相关内容？
│ ├─ Chunk 太大/太小？→ 尝试 chunk_size=256/512/1024
│ ├─ Embedding 模型不适合中文？→ 用 bge-small-zh/bge-m3
│ ├─ Query 太模糊？→ 加 Query Rewriter
│ └─ 需要混合检索？→ 加 BM25
├─ 检索到了但答案不对？
│ ├─ Re-rank 了吗？→ 加 Cross-Encoder 重排序
│ ├─ Prompt 不够好？→ 明确要求"基于资料回答"
│ └─ 上下文太长？→ 限制使用 Top-3 而不是 Top-5
└─ 答案在胡说？
├─ 模型幻觉？→ Prompt 加"不知道就说不知道"
└─ 资料质量差？→ 检查源文档质量

9. 生产环境注意事项

rag_metrics = {
"retrieval_latency_ms": 150, # 检索耗时
"retrieval_recall@5": 0.85, # 前5命中率
"generation_latency_ms": 1200, # 生成耗时
"total_chunks": 15000, # 知识库总量
}

总结

RAG 系统的核心链路：

PDF/TXT → 切分 Chunk → Embedding → 向量库
用户问题 → Query 改写 → 向量检索 → Re-rank → LLM 回答

每个环节都有优化空间，但先把全链路跑通，再根据实际效果逐步调优，是最高效的方式。

本文是 《AI 应用开发实战》系列 的第 4 篇。 系列目录：

本文由 Zyentor（智元界） 原创发布

本文发布于 Zyentor（智元界） —— AI 开发者社区 原文链接：https://www.zyentor.com/news/3324