9.1 入门案例：本地知识库RAG

理论讲了这么多，是时候动手了。这篇手把手带你从零搭建一个​本地 RAG 知识问答系统​——加载你的 PDF/Markdown 文档，用自然语言提问，获得基于文档的准确回答。代码完整可运行，跟着做就能出效果。

📑 目录

• 项目目标与最终效果
• 环境准备
• 第一步：文档加载与分块
• 第二步：Embedding 与向量存储
• 第三步：检索与生成
• 第四步：加个 Web 界面
• 常见问题排查

项目目标与最终效果

最终效果：
你：「公司的差旅报销标准是什么？」
系统：（检索内部制度文档后回答）
「根据《费用管理制度》第 5 章：
 1. 交通费：高铁二等座以内实报实销...
 2. 住宿费：一线城市不超过 500 元/晚...
  3. 需提供正规发票和出差申请单...」
→ 回答有据可查，来源明确！

环境准备

# 创建虚拟环境
python -m venv venv
source venv/bin/activate  # Mac/Linux
# venv\Scripts\activate   # Windows

# 安装依赖
pip install langchain langchain-openai langchain-community \
    chromadb pypdf unstructured python-dotenv

# 准备 .env 文件
echo "OPENAI_API_KEY=sk-xxx" > .env
necho "OPENAI_BASE_URL=https://api.openai.com/v1" >> .env

第一步：文档加载与分块

# rag_app/loaders.py
from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
import os

def load_and_split_documents(directory: str):
    """加载指定目录下的所有 PDF/MD 文档并分块"""
    # 加载文档
    pdf_loader = DirectoryLoader(
        directory, glob="**/*.pdf", loader_cls=PyPDFLoader
    )
    md_loader = DirectoryLoader(
        directory, glob="**/*.md", loader_cls=UnstructuredMarkdownLoader
    )
    
    documents = pdf_loader.load() + md_loader.load()
    print(f"共加载 {len(documents)} 个文档")
    
    # 分块（核心参数！）
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,       # 每块最大字符数
        chunk_overlap=50,      # 块间重叠（防丢失上下文）
        separators=["\n\n", "\n", ".", "。", "！", "？", " "]
    )
    chunks = splitter.split_documents(documents)
    print(f"切分成 {len(chunks)} 个文本块")
    
    return chunks

第二步：Embedding 与向量存储

# rag_app/vectorstore.py
from langchain_openai import OpenAIEmbeddings
from langchain_chroma import Chroma
from loaders import load_and_split_documents
import os

def create_vector_store(data_dir: str, persist_dir: str = "./chroma_db"):
    """创建或加载 Chroma 向量数据库"""
    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
    
    if os.path.exists(persist_dir):
        # 已存在则直接加载（不用重新 Embedding）
        vectorstore = Chroma(
            persist_directory=persist_dir,
            embedding_function=embeddings
        )
        print(f"加载已有向量库，共 {vectorstore.count()} 条记录")
    else:
        # 首次：加载文档 → 分块 → Embedding → 存储
        chunks = load_and_split_documents(data_dir)
        vectorstore = Chroma.from_documents(
            documents=chunks,
            embedding=embeddings,
            persist_directory=persist_dir
        )
        vectorstore.persist()
        print(f"新建向量库，存储了 {len(chunks)} 条记录")
    
    return vectorstore

第三步：检索与生成

# rag_app/chain.py
from langchain_openai import ChatOpenAI
from langchain.chains import create_retrieval_chain
from langchain_core.prompts import ChatPromptTemplate
from vectorstore import create_vector_store

def build_rag_chain(data_dir: str = "./docs"):
    """构建完整的 RAG 检索链"""
    # 初始化组件
    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
    vectorstore = create_vector_store(data_dir)
    retriever = vectorstore.as_retriever(
        search_type="similarity", search_kwargs={"k": 5}
    )
    
    # 定义 Prompt 模板（关键！）
    prompt = ChatPromptTemplate.from_messages([
        ("system", "你是一个有帮助的 AI 助手。请严格基于以下参考资料来回答问题。"
         "如果参考资料中没有相关信息，请直接说『根据现有资料无法回答这个问题』，不要编造内容。"
         "如果可以回答，请在回答末尾标注信息来源。"),
        ("user", "{input}\n\n参考资料:\n{context}"),
    ])
    
    # 构建检索链
    chain = create_retrieval_chain(retriever, prompt | llm)
    return chain

# 使用
if __name__ == "__main__":
    rag_chain = build_rag_chain("./my_documents")
    
    while True:
        query = input("\n请输入你的问题（输入 quit 退出）：")
        if query.lower() == "quit":
            break
        result = rag_chain.invoke({"input": query})
        print(f"\n📝 回答：{result['answer']}")
        print(f"\n📚 引用了 {len(result['context'])} 个文档片段")

第四步：加个 Web 界面

# 安装 Streamlit
pip install streamlit

# app.py — 用 50 行代码做一个 Web UI
import streamlit as st
from chain import build_rag_chain

st.set_page_config("RAG 知识库问答", "📚")
st.title("📚 本地知识库 RAG 问答系统")

# 侧边栏配置
with st.sidebar:
    st.header("设置")
    data_dir = st.text_input("文档目录", value="./docs")
    k = st.slider("检索数量", 1, 10, 5)
    st.caption("将 PDF/Markdown 文件放入文档目录即可")

# 初始化 Chain
if "chain" not in st.session_state:
    with st.spinner("正在加载文档和建立索引..."):
        st.session_state.chain = build_rag_chain(data_dir)
    st.success(f"就绪！已加载文档")

# 对话界面
if prompt := st.chat_input("请输入问题..."):
    st.chat_message("user").write(prompt)
    
    with st.spinner("正在检索和生成回答..."):
        response = st.session_state.chain.invoke({"input": prompt})
    
    st.chat_message("assistant").write(response["answer"])
    
    with st.expander("查看引用的原始文档"):
        for i, doc in enumerate(response.get("context", [])):
            st.text(f"片段 {i+1}: {doc.page_content[:200]}...")

# 运行：streamlit run app.py

常见问题排查