10 | RAG 入门——让 LLM 用你的私有知识回答问题

Indexing（离线）          Retrieval（在线）        Generation（在线）
文档 → 切分 → 向量化      用户提问 → 向量化         检索结果 + 问题
         → 存入向量库          → 在向量库中检索          → 交给 LLM 生成回答
                              → （可选）重排

from langchain.text_splitter import RecursiveCharacterTextSplitter

splitter = RecursiveCharacterTextSplitter(
    chunk_size=512,       # 每块最多 512 个字符
    chunk_overlap=64,     # 相邻两块重叠 64 个字符
    separators=["\n\n", "\n", "。", ".", " "]  # 按优先级尝试分割
)

chunks = splitter.split_text(document_text)
# 一篇 10,000 字的文档 → 大约 22 个 chunk（考虑 overlap）

# 按 Markdown 标题切分
from langchain.text_splitter import MarkdownHeaderTextSplitter

md_splitter = MarkdownHeaderTextSplitter(
    headers_to_split_on=[
        ("#", "h1"),
        ("##", "h2"),
        ("###", "h3"),
    ]
)

# 一篇技术文档会被按章节自然切分
# 每个 chunk 自带标题元数据，检索时更精准

"Python 异步编程"  →  [0.12, -0.34, 0.56, ..., 0.78]   (1536 维)
"asyncio 教程"     →  [0.11, -0.31, 0.54, ..., 0.76]   ← 距离很近
"红烧肉的做法"     →  [-0.89, 0.45, -0.12, ..., 0.03]  ← 距离很远

from langchain_openai import OpenAIEmbeddings

embeddings = OpenAIEmbeddings(
    model="text-embedding-3-small",  # 1536 维
    # 或 "text-embedding-3-large"  # 3072 维
)

# 将文本块批量转为向量
vectors = embeddings.embed_documents([chunk.text for chunk in chunks])
# 100 个 chunk → 100 个 1536 维向量

import chromadb

# 零配置，数据存在内存中（也可持久化到磁盘）
client = chromadb.Client()
collection = client.create_collection(
    name="my_docs",
    metadata={"hnsw:space": "cosine"}  # 使用余弦相似度
)

# 存入文档块及其向量
collection.add(
    documents=[chunk.text for chunk in chunks],
    embeddings=vectors,
    ids=[f"chunk_{i}" for i in range(len(chunks))],
    metadatas=[{"source": "api_doc.pdf", "page": 3} for _ in chunks]
)

results = collection.query(
    query_embeddings=[query_vector],
    n_results=5  # 返回最相似的 5 个
)
# results = [{"text": "...", "distance": 0.12}, ...]
# distance 越小越相似（cosine distance）

from langchain_community.vectorstores import Chroma

vectorstore = Chroma.from_documents(chunks, embeddings)

# 普通相似度搜索
docs_sim = vectorstore.similarity_search("Python 异步编程", k=5)
# 可能返回 5 个高度重复的 chunk

# MMR 搜索（lambda_mult 控制多样性，0=最多样，1=最相似）
docs_mmr = vectorstore.max_marginal_relevance_search(
    "Python 异步编程",
    k=5,
    lambda_mult=0.7  # 0.7 偏向相似，0.3 偏向多样
)
# 返回 5 个覆盖面更广的 chunk

from langchain.retrievers import ContextualCompressionRetriever
from langchain_cohere import CohereRerank

# 第一步：向量检索召回 top-20（宽松召回）
base_retriever = vectorstore.as_retriever(search_kwargs={"k": 20})

# 第二步：用 Cross-encoder 重排，保留 top-5
reranker = CohereRerank(model="rerank-multilingual-v3.0", top_n=5)

# 组合：先粗召回，再精排
compression_retriever = ContextualCompressionRetriever(
    base_compressor=reranker,
    base_retriever=base_retriever
)

# 使用
docs = compression_retriever.invoke("如何配置 Nginx 反向代理")
# 返回经过重排的 top-5，质量显著优于纯向量检索

RAG_PROMPT_TEMPLATE = """你是一个专业的技术文档助手。请基于以下参考文档回答用户的问题。

## 参考文档
{context}

## 回答要求
1. 只基于参考文档中的信息回答，不要编造
2. 如果参考文档中没有足够信息，明确说"根据现有文档无法确定"
3. 引用具体的文档来源（如 [文档1]）
4. 用中文回答，技术术语保留英文

## 用户问题
{question}
"""

def format_docs(docs):
    """把检索结果格式化为上下文"""
    formatted = []
    for i, doc in enumerate(docs):
        source = doc.metadata.get("source", "未知")
        formatted.append(f"[文档{i+1}] (来源: {source})\n{doc.page_content}")
    return "\n\n".join(formatted)

# 组装最终 Prompt
final_prompt = RAG_PROMPT_TEMPLATE.format(
    context=format_docs(docs),
    question=user_question
)

# pip install langchain langchain-openai langchain-chroma chromadb

from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_chroma import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.prompts import ChatPromptTemplate
from langchain.schema import StrOutputParser
from langchain.schema.runnable import RunnablePassthrough

# ============ 1. 准备文档 ============
# 假设你有一些技术文档
documents = [
    "LangChain 是一个用于构建 LLM 应用的框架。它提供了 Chain、Agent、Memory 等核心抽象。",
    "Chain 是 LangChain 中的核心概念，表示一系列按顺序执行的组件。可以用 LCEL（LangChain Expression Language）构建。",
    "Agent 使用 LLM 作为推理引擎，自主决定调用哪些工具。ReAct 是最常用的 Agent 架构。",
    "RAG（检索增强生成）通过检索外部知识库来增强 LLM 的回答能力，减少幻觉。",
    "向量数据库用于存储和检索文本的向量表示。常见选择包括 Chroma、FAISS、Pinecone。",
    # ... 更多文档
]

# ============ 2. 切分文档 ============
splitter = RecursiveCharacterTextSplitter(
    chunk_size=512,
    chunk_overlap=64,
    separators=["\n\n", "\n", "。", ".", " "]
)

# 这里文档已经很短，实际中你会 split_documents()
# chunks = splitter.split_documents(loaded_docs)

# ============ 3. 创建向量存储 ============
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

vectorstore = Chroma.from_texts(
    texts=documents,
    embedding=embeddings,
    collection_name="my_knowledge_base"
)

retriever = vectorstore.as_retriever(
    search_type="mmr",          # 使用 MMR 保证多样性
    search_kwargs={
        "k": 5,                 # 返回 top-5
        "lambda_mult": 0.7
    }
)

# ============ 4. 构建 Prompt ============
RAG_PROMPT = ChatPromptTemplate.from_template("""
你是一个专业的 AI 技术助手。请基于以下参考文档回答问题。

参考文档：
{context}

要求：
- 只基于参考文档回答，不确定时明确说明
- 引用文档编号 [1] [2] 等
- 用中文回答，技术术语保留英文

问题：{question}
""")

# ============ 5. 构建 RAG Chain ============
llm = ChatOpenAI(model="gpt-4o", temperature=0)

def format_docs(docs):
    return "\n\n".join(
        f"[{i+1}] {doc.page_content}" for i, doc in enumerate(docs)
    )

rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | RAG_PROMPT
    | llm
    | StrOutputParser()
)

# ============ 6. 使用 ============
answer = rag_chain.invoke("LangChain 中的 Agent 和 Chain 有什么区别？")
print(answer)
# 输出会引用 [1] [2] [3] 等文档来源，基于知识库回答