📚 课程目录第2章实现一个基于RAG的文档问答系统

2. 实现一个基于RAG的文档问答系统

任务目标

构建一个最小可运行的RAG问答系统：给定一组本地文本文档，用户提问后，系统先从文档中检索最相关的片段，再将片段与问题一起送入大模型生成回答。本案例聚焦RAG核心链路，不涉及分布式向量库与生产级部署。

环境与依赖

pip install langchain langchain-community chromadb sentence-transformers flask

使用开源嵌入模型 BAAI/bge-small-zh-v1.5 与本地向量库 Chroma，LLM 部分通过兼容接口调用（可用 base_url 指向本地 vLLM/Ollama 服务）。

完整实现代码

import os
from langchain_community.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain_ import ChatOpenAI

# ========== 1. 文档加载与切分 ==========
def load_and_split(doc_dir: str):
    docs = []
    for fname in os.listdir(doc_dir):
        if fname.endswith(".txt"):
            loader = TextLoader(os.path.join(doc_dir, fname), encoding="utf-8")
            docs.extend(loader.load())
    splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=50)
    return splitter.split_documents(docs)

# ========== 2. 构建向量库 ==========
def build_vectorstore(doc_dir: str, persist_dir: str):
    embedding = HuggingFaceEmbeddings(model_name="BAAI/bge-small-zh-v1.5")
    if os.path.exists(persist_dir) and os.listdir(persist_dir):
        return Chroma(persist_directory=persist_dir, embedding_function=embedding)
    chunks = load_and_split(doc_dir)
    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embedding,
        persist_directory=persist_dir
    )
    vectorstore.persist()
    return vectorstore

# ========== 3. 创建问答链 ==========
def create_qa_chain(vectorstore):
    llm = ChatOpenAI(
        model="qwen2.5-7b-instruct",
        temperature=0.1,
        base_url="http://localhost:8000/v1",  # 指向本地推理服务
        ="EMPTY"
    )
    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
    return RetrievalQA.from_chain_type(
        llm=llm,
        retriever=retriever,
        return_source_documents=True
    )

# ========== 4. 主流程 ==========
if __name__ == "__main__":
    DOC_DIR = "./docs"          # 存放待检索的 .txt 文档
    PERSIST_DIR = "./chroma_db" # 向量库持久化目录
    
    # 准备示例文档（若不存在则创建）
    os.makedirs(DOC_DIR, exist_ok=True)
    sample_file = os.path.join(DOC_DIR, "company.txt")
    if not os.path.exists(sample_file):
        with open(sample_file, "w", encoding="utf-8") as f:
            f.write("公司成立于2015年，总部位于深圳，主营AI客服机器人。\n"
                    "2023年营收达2.3亿元，员工总数约400人。\n"
                    "核心产品包括智能语音导航与在线文本机器人。")
    
    vectorstore = build_vectorstore(DOC_DIR, PERSIST_DIR)
    qa = create_qa_chain(vectorstore)
    
    # 交互式问答
    while True:
        query = input("\n请输入问题（输入 exit 退出）：").strip()
        if query.lower() == "exit":
            break
        result = qa.invoke({"query": query})
        print("\n回答：", result["result"])
        print("\n参考来源：")
        for i, doc in enumerate(result["source_documents"], 1):
            print(f"  [{i}] {doc.page_content[:80]}...")

运行方式

在代码同级目录创建 docs/ 文件夹，放入若干 .txt 文档（编码UTF-8）。
启动本地大模型推理服务（如 vLLM 或 Ollama），确保 base_url 指向正确端口。
执行 python rag_system.py，首次运行会自动下载嵌入模型并构建向量库。

验证方法与通过标准

检索验证：输入问题“公司总部在哪里？”，若回答包含“深圳”且参考来源指向 company.txt 对应片段，说明检索链路正常。
生成验证：输入“2023年营收是多少？”，回答应准确给出“2.3亿元”，且不出现与文档无关的编造内容。
边界测试：输入文档中不存在的问题（如“公司CEO是谁？”），系统应明确表示未找到相关信息，而非强行编造答案。

常见问题

首次运行下载慢：HuggingFace 嵌入模型约100MB，可提前设置 HF_ENDPOINT=https://hf-mirror.com 加速。
检索结果不相关：调小 chunk_size 至200或增大 k 值，同时确认文档编码为UTF-8。
LLM连接失败：检查 base_url 是否与推理服务实际端口一致，并确认模型名称正确。