手把手教你用Python构建RAG知识库系统（完整代码）

作者：未来情报所 | 发布时间：2026-04-01 | 阅读时长：约15分钟

---

前言

随着大语言模型（LLM）的普及，越来越多的开发者希望将自己的私有知识库与 AI 结合，打造专属的智能问答系统。RAG（Retrieval-Augmented Generation，检索增强生成） 正是解决这一问题的核心技术。

本文将从零开始，手把手带你用 Python 构建一个完整的 RAG 知识库系统，包含文档处理、向量化、检索和生成全流程，所有代码均可直接运行。

---

一、RAG 原理深度解析

1.1 为什么需要 RAG？

传统 LLM 存在两大核心痛点：

1. 知识截止问题：模型训练数据有时间限制，无法获取最新信息
2. 幻觉问题：模型可能生成看似合理但实际错误的内容

RAG 通过"先检索、后生成"的范式解决了这两个问题：

用户提问 → 向量检索相关文档 → 将文档作为上下文 → LLM 生成答案

1.2 RAG 核心组件

组件	功能	常用工具
文档加载器	解析各类文档格式	LangChain Loaders
文本分割器	将长文档切分为块	RecursiveCharacterTextSplitter
嵌入模型	将文本转为向量	OpenAI Embeddings / BGE
向量数据库	存储和检索向量	Chroma / Faiss / Milvus
检索器	根据查询找相关文档	VectorStoreRetriever
生成模型	基于上下文生成答案	GPT-4 / Qwen / Llama

1.3 RAG 工作流程

┌─────────────────────────────────────────────────────┐
│                   离线索引阶段                        │
│  文档 → 分块 → 向量化 → 存入向量数据库               │
└─────────────────────────────────────────────────────┘
                          ↓
┌─────────────────────────────────────────────────────┐
│                   在线检索阶段                        │
│  用户问题 → 向量化 → 相似度检索 → Top-K 文档块       │
└─────────────────────────────────────────────────────┘
                          ↓
┌─────────────────────────────────────────────────────┐
│                   生成阶段                            │
│  问题 + 检索文档 → Prompt 构建 → LLM → 最终答案      │
└─────────────────────────────────────────────────────┘

---

二、环境搭建

2.1 安装依赖

pip install langchain langchain-community langchain-openai
pip install chromadb
pip install sentence-transformers
pip install pypdf python-docx
pip install openai tiktoken
pip install fastapi uvicorn  # 可选，用于 API 服务

2.2 项目结构

rag_system/
├── config.py          # 配置文件
├── document_loader.py # 文档加载与处理
├── vector_store.py    # 向量数据库管理
├── retriever.py       # 检索器
├── rag_chain.py       # RAG 链
├── api.py             # FastAPI 接口
├── main.py            # 主程序入口
└── docs/              # 知识库文档目录
    ├── sample.pdf
    └── sample.txt

---

三、文档处理模块

3.1 配置文件

# config.py
import os
from dataclasses import dataclass

@dataclass
class RAGConfig:
    # OpenAI 配置（也可替换为本地模型）
    openai_api_key: str = os.getenv("OPENAI_API_KEY", "your-api-key")
    openai_base_url: str = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
    
    # 嵌入模型配置
    embedding_model: str = "text-embedding-3-small"  # 或使用本地 BGE 模型
    
    # LLM 配置
    llm_model: str = "gpt-4o-mini"
    llm_temperature: float = 0.1
    llm_max_tokens: int = 2048
    
    # 文档分割配置
    chunk_size: int = 500          # 每块字符数
    chunk_overlap: int = 50        # 块间重叠字符数
    
    # 检索配置
    top_k: int = 5                 # 检索返回的文档块数量
    score_threshold: float = 0.7   # 相似度阈值
    
    # 向量数据库配置
    chroma_persist_dir: str = "./chroma_db"
    collection_name: str = "knowledge_base"

config = RAGConfig()

3.2 文档加载器

# document_loader.py
import os
from typing import List
from pathlib import Path

from langchain.schema import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import (
    PyPDFLoader,
    TextLoader,
    Docx2txtLoader,
    DirectoryLoader,
    UnstructuredMarkdownLoader,
)

class DocumentProcessor:
    """文档处理器：支持 PDF、TXT、DOCX、MD 等格式"""
    
    def __init__(self, chunk_size: int = 500, chunk_overlap: int = 50):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            length_function=len,
            separators=["\n\n", "\n", "。", "！", "？", ".", "!", "?", " ", ""],
        )
    
    def load_file(self, file_path: str) -> List[Document]:
        """根据文件类型选择对应的加载器"""
        path = Path(file_path)
        suffix = path.suffix.lower()
        
        loader_map = {
            ".pdf": PyPDFLoader,
            ".txt": lambda p: TextLoader(p, encoding="utf-8"),
            ".docx": Docx2txtLoader,
            ".md": UnstructuredMarkdownLoader,
        }
        
        if suffix not in loader_map:
            raise ValueError(f"不支持的文件格式: {suffix}")
        
        loader_cls = loader_map[suffix]
        loader = loader_cls(file_path)
        documents = loader.load()
        
        # 添加元数据
        for doc in documents:
            doc.metadata["source"] = str(path.name)
            doc.metadata["file_type"] = suffix
        
        return documents
    
    def load_directory(self, dir_path: str) -> List[Document]:
        """加载目录下所有支持的文档"""
        all_docs = []
        supported_extensions = [".pdf", ".txt", ".docx", ".md"]
        
        for root, _, files in os.walk(dir_path):
            for file in files:
                file_path = os.path.join(root, file)
                ext = Path(file).suffix.lower()
                if ext in supported_extensions:
                    try:
                        docs = self.load_file(file_path)
                        all_docs.extend(docs)
                        print(f"✅ 已加载: {file} ({len(docs)} 页)")
                    except Exception as e:
                        print(f"❌ 加载失败: {file} - {e}")
        
        return all_docs
    
    def split_documents(self, documents: List[Document]) -> List[Document]:
        """将文档分割为小块"""
        chunks = self.text_splitter.split_documents(documents)
        print(f"📄 文档分割完成: {len(documents)} 个文档 → {len(chunks)} 个文本块")
        return chunks
    
    def process(self, source: str) -> List[Document]:
        """一键处理：加载 + 分割"""
        if os.path.isdir(source):
            docs = self.load_directory(source)
        else:
            docs = self.load_file(source)
        
        return self.split_documents(docs)

---

四、向量数据库模块

# vector_store.py
from typing import List, Optional
from langchain.schema import Document
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_community.embeddings import HuggingFaceEmbeddings

from config import config

class VectorStoreManager:
    """向量数据库管理器"""
    
    def __init__(self, use_local_embedding: bool = False):
        """
        Args:
            use_local_embedding: True 使用本地 BGE 模型，False 使用 OpenAI
        """
        if use_local_embedding:
            # 使用本地嵌入模型（无需 API Key，推荐国内用户）
            self.embeddings = HuggingFaceEmbeddings(
                model_name="BAAI/bge-small-zh-v1.5",
                model_kwargs={"device": "cpu"},
                encode_kwargs={"normalize_embeddings": True},
            )
        else:
            self.embeddings = OpenAIEmbeddings(
                model=config.embedding_model,
                openai_api_key=config.openai_api_key,
                openai_api_base=config.openai_base_url,
            )
        
        self.vector_store: Optional[Chroma] = None
    
    def create_from_documents(self, documents: List[Document]) -> Chroma:
        """从文档列表创建向量数据库"""
        print(f"🔄 正在向量化 {len(documents)} 个文本块...")
        
        self.vector_store = Chroma.from_documents(
            documents=documents,
            embedding=self.embeddings,
            persist_directory=config.chroma_persist_dir,
            collection_name=config.collection_name,
        )
        
        self.vector_store.persist()
        print(f"✅ 向量数据库创建完成，已持久化到: {config.chroma_persist_dir}")
        return self.vector_store
    
    def load_existing(self) -> Chroma:
        """加载已有的向量数据库"""
        self.vector_store = Chroma(
            persist_directory=config.chroma_persist_dir,
            embedding_function=self.embeddings,
            collection_name=config.collection_name,
        )
        count = self.vector_store._collection.count()
        print(f"✅ 已加载向量数据库，共 {count} 个向量")
        return self.vector_store
    
    def add_documents(self, documents: List[Document]):
        """向已有数据库追加文档"""
        if self.vector_store is None:
            raise RuntimeError("请先创建或加载向量数据库")
        self.vector_store.add_documents(documents)
        self.vector_store.persist()
        print(f"✅ 已追加 {len(documents)} 个文本块")
    
    def similarity_search(self, query: str, k: int = 5) -> List[Document]:
        """相似度搜索"""
        if self.vector_store is None:
            raise RuntimeError("请先创建或加载向量数据库")
        return self.vector_store.similarity_search(query, k=k)

---

五、RAG 检索链

# rag_chain.py
from typing import List, Dict, Any, Optional
from langchain.schema import Document
from langchain_openai import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.schema.runnable import RunnablePassthrough, RunnableLambda
from langchain.schema.output_parser import StrOutputParser
from langchain_community.vectorstores import Chroma

from config import config

# RAG 提示词模板
RAG_PROMPT_TEMPLATE = """你是一个专业的知识库问答助手。请根据以下检索到的相关文档内容，准确回答用户的问题。

【检索到的相关文档】
{context}

【用户问题】
{question}

【回答要求】
1. 仅基于提供的文档内容回答，不要编造信息
2. 如果文档中没有相关信息，请明确告知用户
3. 回答要简洁、准确、有条理
4. 如有必要，可以引用文档来源

【回答】"""

class RAGChain:
    """RAG 问答链"""
    
    def __init__(self, vector_store: Chroma):
        self.vector_store = vector_store
        self.retriever = vector_store.as_retriever(
            search_type="similarity_score_threshold",
            search_kwargs={
                "k": config.top_k,
                "score_threshold": config.score_threshold,
            },
        )
        
        self.llm = ChatOpenAI(
            model=config.llm_model,
            temperature=config.llm_temperature,
            max_tokens=config.llm_max_tokens,
            openai_api_key=config.openai_api_key,
            openai_api_base=config.openai_base_url,
        )
        
        self.prompt = ChatPromptTemplate.from_template(RAG_PROMPT_TEMPLATE)
        self.chain = self._build_chain()
    
    def _format_docs(self, docs: List[Document]) -> str:
        """格式化检索到的文档"""
        if not docs:
            return "未找到相关文档。"
        
        formatted = []
        for i, doc in enumerate(docs, 1):
            source = doc.metadata.get("source", "未知来源")
            formatted.append(f"[文档{i}] 来源: {source}\n{doc.page_content}")
        
        return "\n\n---\n\n".join(formatted)
    
    def _build_chain(self):
        """构建 RAG 链"""
        return (
            {
                "context": self.retriever | RunnableLambda(self._format_docs),
                "question": RunnablePassthrough(),
            }
            | self.prompt
            | self.llm
            | StrOutputParser()
        )
    
    def query(self, question: str) -> Dict[str, Any]:
        """执行 RAG 查询，返回答案和来源文档"""
        # 检索相关文档
        retrieved_docs = self.retriever.get_relevant_documents(question)
        
        # 生成答案
        answer = self.chain.invoke(question)
        
        # 整理来源信息
        sources = list(set([
            doc.metadata.get("source", "未知") 
            for doc in retrieved_docs
        ]))
        
        return {
            "question": question,
            "answer": answer,
            "sources": sources,
            "retrieved_docs_count": len(retrieved_docs),
        }
    
    def stream_query(self, question: str):
        """流式输出答案"""
        for chunk in self.chain.stream(question):
            yield chunk

---

六、完整 Demo 主程序

# main.py
import os
import sys
from pathlib import Path

from document_loader import DocumentProcessor
from vector_store import VectorStoreManager
from rag_chain import RAGChain
from config import config

def build_knowledge_base(docs_dir: str, use_local_embedding: bool = True):
    """构建知识库"""
    print("=" * 50)
    print("🚀 开始构建 RAG 知识库")
    print("=" * 50)
    
    # Step 1: 处理文档
    processor = DocumentProcessor(
        chunk_size=config.chunk_size,
        chunk_overlap=config.chunk_overlap,
    )
    chunks = processor.process(docs_dir)
    
    if not chunks:
        print("❌ 未找到任何文档，请检查文档目录")
        return None
    
    # Step 2: 创建向量数据库
    vs_manager = VectorStoreManager(use_local_embedding=use_local_embedding)
    vector_store = vs_manager.create_from_documents(chunks)
    
    print("\n✅ 知识库构建完成！")
    return vector_store

def load_knowledge_base(use_local_embedding: bool = True):
    """加载已有知识库"""
    vs_manager = VectorStoreManager(use_local_embedding=use_local_embedding)
    return vs_manager.load_existing()

def interactive_qa(rag_chain: RAGChain):
    """交互式问答"""
    print("\n" + "=" * 50)
    print("💬 RAG 知识库问答系统已启动")
    print("输入 'quit' 或 'exit' 退出")
    print("=" * 50 + "\n")
    
    while True:
        question = input("🙋 请输入您的问题: ").strip()
        
        if not question:
            continue
        if question.lower() in ["quit", "exit", "退出"]:
            print("👋 再见！")
            break
        
        print("\n🔍 正在检索相关文档...")
        result = rag_chain.query(question)
        
        print(f"\n📚 检索到 {result['retrieved_docs_count']} 个相关文档块")
        print(f"📎 来源: {', '.join(result['sources']) if result['sources'] else '无'}")
        print(f"\n🤖 AI 回答:\n{result['answer']}")
        print("\n" + "-" * 50 + "\n")

def main():
    docs_dir = "./docs"
    use_local = True  # True=本地BGE模型, False=OpenAI
    
    # 判断是否已有向量数据库
    if os.path.exists(config.chroma_persist_dir):
        print("📂 检测到已有知识库，直接加载...")
        vector_store = load_knowledge_base(use_local)
    else:
        print("📂 未检测到知识库，开始构建...")
        os.makedirs(docs_dir, exist_ok=True)
        
        # 创建示例文档
        sample_file = os.path.join(docs_dir, "sample.txt")
        if not os.path.exists(sample_file):
            with open(sample_file, "w", encoding="utf-8") as f:
                f.write("""RAG（检索增强生成）技术简介

RAG是一种结合信息检索和文本生成的AI技术框架。
它通过在生成答案前先检索相关文档，显著提升了大语言模型的准确性和可靠性。

RAG的主要优势：
1. 减少幻觉：基于真实文档生成答案
2. 知识更新：无需重新训练模型即可更新知识
3. 可追溯性：答案可以追溯到具体文档来源
4. 成本效益：比全量微调更经济高效

RAG的应用场景：
- 企业内部知识库问答
- 客服自动化系统
- 法律文档分析
- 医疗知识查询
- 代码文档助手
""")
            print(f"📝 已创建示例文档: {sample_file}")
        
        vector_store = build_knowledge_base(docs_dir, use_local)
        if vector_store is None:
            sys.exit(1)
    
    # 创建 RAG 链并启动问答
    rag_chain = RAGChain(vector_store)
    interactive_qa(rag_chain)

if __name__ == "__main__":
    main()

---

七、FastAPI 接口封装

# api.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional
import uvicorn

from vector_store import VectorStoreManager
from rag_chain import RAGChain

app = FastAPI(title="RAG 知识库 API", version="1.0.0")

# 全局 RAG 链实例
rag_chain: Optional[RAGChain] = None

class QueryRequest(BaseModel):
    question: str
    top_k: Optional[int] = 5

class QueryResponse(BaseModel):
    question: str
    answer: str
    sources: List[str]
    retrieved_docs_count: int

@app.on_event("startup")
async def startup_event():
    global rag_chain
    vs_manager = VectorStoreManager(use_local_embedding=True)
    vector_store = vs_manager.load_existing()
    rag_chain = RAGChain(vector_store)
    print("✅ RAG 服务启动成功")

@app.post("/query", response_model=QueryResponse)
async def query(request: QueryRequest):
    if rag_chain is None:
        raise HTTPException(status_code=503, detail="RAG 服务未初始化")
    
    result = rag_chain.query(request.question)
    return QueryResponse(**result)

@app.get("/health")
async def health():
    return {"status": "ok", "service": "RAG Knowledge Base"}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

---

八、性能优化技巧

8.1 混合检索（Hybrid Search）

from langchain.retrievers import BM25Retriever, EnsembleRetriever

# 结合 BM25 关键词检索 + 向量语义检索
bm25_retriever = BM25Retriever.from_documents(chunks)
bm25_retriever.k = 5

vector_retriever = vector_store.as_retriever(search_kwargs={"k": 5})

# 集成检索器（权重可调）
ensemble_retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, vector_retriever],
    weights=[0.4, 0.6],  # BM25 40% + 向量 60%
)

8.2 重排序（Reranking）

from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import CrossEncoderReranker
from langchain_community.cross_encoders import HuggingFaceCrossEncoder

# 使用 Cross-Encoder 对检索结果重排序
model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-base")
compressor = CrossEncoderReranker(model=model, top_n=3)

compression_retriever = ContextualCompressionRetriever(
    base_compressor=compressor,
    base_retriever=vector_retriever,
)

8.3 HyDE（假设文档嵌入）

from langchain.chains import HypotheticalDocumentEmbedder

# 先让 LLM 生成假设答案，再用假设答案检索
hyde_embeddings = HypotheticalDocumentEmbedder.from_llm(
    llm=llm,
    base_embeddings=embeddings,
    custom_prompt=hyde_prompt,
)

---

总结

本文完整实现了一个生产可用的 RAG 知识库系统，核心要点：

1. 文档处理：支持多格式，合理设置 chunk_size 和 overlap
2. 向量化：国内推荐使用 BGE 本地模型，避免 API 依赖
3. 检索优化：混合检索 + 重排序可显著提升召回质量
4. 生成控制：低 temperature + 严格 Prompt 减少幻觉

完整代码已上传 GitHub，欢迎 Star ⭐

未来情报所 | 专注 AI 工程化落地，关注获取更多实战干货

---

标签：#Python #RAG #LangChain #向量数据库 #大模型 #AI工程化