从零搭建RAG系统：检索增强生成完整实战

from langchain_community.document_loaders import (
    PyPDFLoader,
    DirectoryLoader,
    TextLoader,
    UnstructuredMarkdownLoader
)
import os

def load_documents(doc_dir: str):
    """递归加载目录下所有支持格式的文档"""
    all_docs = []
    
    # PDF文件
    pdf_loader = DirectoryLoader(
        doc_dir, glob="**/*.pdf",
        loader_cls=PyPDFLoader, show_progress=True
    )
    all_docs.extend(pdf_loader.load())
    
    # Markdown文件
    md_loader = DirectoryLoader(
        doc_dir, glob="**/*.md",
        loader_cls=UnstructuredMarkdownLoader, show_progress=True
    )
    all_docs.extend(md_loader.load())
    
    # 纯文本文件
    txt_loader = DirectoryLoader(
        doc_dir, glob="**/*.txt",
        loader_cls=TextLoader,
        loader_kwargs={"encoding": "utf-8"}, show_progress=True
    )
    all_docs.extend(txt_loader.load())
    
    print(f"共加载 {len(all_docs)} 个文档片段")
    return all_docs

from langchain.text_splitter import RecursiveCharacterTextSplitter

def split_documents(documents, chunk_size=500, chunk_overlap=80):
    """将文档分割成适合检索的小块"""
    
    # 中文文本的分隔符优先级
    separators = [
        "\n\n",    # 段落分隔
        "\n",       # 换行
        "。",        # 句号
        "！",        # 感叹号
        "？",        # 问号
        "；",        # 分号
        "，",        # 逗号
        " ",         # 空格
        ""           # 字符级（兜底）
    ]
    
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        length_function=len,
        separators=separators,
        is_separator_regex=False
    )
    
    chunks = splitter.split_documents(documents)
    print(f"分割成 {len(chunks)} 个文本块")
    print(f"平均块大小: {sum(len(c.page_content) for c in chunks) / len(chunks):.0f} 字符")
    
    return chunks

from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma

def create_vectorstore(chunks, persist_dir="./chroma_db"):
    """创建向量数据库"""
    
    # 选用中文优化的嵌入模型
    embeddings = HuggingFaceEmbeddings(
        model_name="BAAI/bge-small-zh-v1.5",
        model_kwargs={"device": "cpu"},
        encode_kwargs={
            "normalize_embeddings": True,
            "batch_size": 64
        }
    )
    
    # 批量创建向量并存储
    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=persist_dir,
        collection_metadata={"hnsw:space": "cosine"}
    )
    
    print(f"向量数据库已持久化到 {persist_dir}")
    print(f"共 {vectorstore._collection.count()} 条向量")
    
    return vectorstore

from langchain_community.chat_models import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate

def build_qa_chain(vectorstore):
    """构建RAG问答链"""
    
    llm = ChatOpenAI(
        base_url="http://localhost:11434/v1",
        api_key="ollama",
        model="deepseek-r1:7b",
        temperature=0.1,     # 低温度减少随机性
        max_tokens=2048
    )
    
    # 自定义提示词——这是效果的关键
    prompt = PromptTemplate(
        template="""你是一个严谨的技术文档问答助手。请严格基于以下参考文档回答用户的问题。

## 参考文档
{context}

## 用户问题
{question}

## 回答规则
1. 只根据参考文档的内容回答，绝对不要编造信息
2. 如果参考文档中没有相关信息，直接说"根据现有文档，我无法回答这个问题"
3. 引用文档中的具体段落来支撑你的回答
4. 如果文档信息不足以完整回答，说明哪些部分有依据、哪些部分文档未提及""",
        input_variables=["context", "question"]
    )
    
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vectorstore.as_retriever(
            search_type="similarity",
            search_kwargs={"k": 4}     # 检索最相关的4个片段
        ),
        chain_type_kwargs={"prompt": prompt},
        return_source_documents=True,
        verbose=False
    )
    
    return qa_chain

# 摄取阶段（一次性）
documents = load_documents("./docs")
chunks = split_documents(documents, chunk_size=500, chunk_overlap=80)
vectorstore = create_vectorstore(chunks)

# 查询阶段（日常使用）
qa_chain = build_qa_chain(vectorstore)
result = qa_chain({"query": "我们的退款政策是什么？"})

print(result["result"])
print("\n--- 引用来源 ---")
for doc in result["source_documents"]:
    print(f"文件: {doc.metadata.get('source', '未知')}")
    print(f"内容: {doc.page_content[:150]}...")