用 Python 实现 RAG:从文档加载到语义检索全流程

检索增强生成(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生成回答

关键要点:

  1. 分块质量决定检索上限 ------ 注意 chunk_size 和 overlap 的调参
  2. Embedding 模型 决定语义理解质量 ------ 中文场景推荐 bge 系列
  3. 检索策略影响召回率 ------ 推荐混合检索 + 重排序
  4. Prompt 工程影响最终输出质量 ------ 明确指令,约束幻觉
相关推荐
AI_Claude_code2 小时前
安全与合规核心:匿名化、日志策略与法律风险规避
网络·爬虫·python·tcp/ip·安全·http·网络爬虫
大任视点2 小时前
深耕AI短剧赛道!聿潇娱乐签约鹤砚声工作室 加速精品内容布局
人工智能
chao1898442 小时前
基于改进二进制粒子群算法的含需求响应机组组合问题MATLAB实现
开发语言·算法·matlab
lcj25112 小时前
字符函数,字符串函数,内存函数
c语言·开发语言·c++·windows
独特的螺狮粉2 小时前
古诗词飞花令随机出题小助手:鸿蒙Flutter框架 实现的古诗词游戏应用
开发语言·flutter·游戏·华为·架构·开源·harmonyos
bryant_meng2 小时前
【Reading Notes】(8.11)Favorite Articles from 2025 November
人工智能·深度学习·业界资讯
Eiceblue2 小时前
Python 如何实现 Excel 数据分列?一列拆分为多列
python·microsoft·excel
不是株2 小时前
FastAPI
python·fastapi
Spliceㅤ2 小时前
Transformer
人工智能·深度学习·transformer