从零带你使用Python实现RAG(检索增强生成)系统

pip install langchain langchain-community \
             chromadb sentence-transformers \
             pypdf unstructured \
             openai tiktoken

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: 元数据（页码、来源等）

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

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

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

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

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"
    )

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

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

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

# ===== 三种检索方式对比 =====

# 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
)

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)

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 系统 —— 从文档到问答
"""
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}")

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

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]]