RAG 让 AI 能"查资料"：检索增强生成完全指南

from langchain_community.document_loaders import PyPDFLoader, TextLoader

# 加载 PDF
loader = PyPDFLoader("employee_handbook.pdf")
docs = loader.load()

# 加载文本
loader = TextLoader("faq.txt", encoding="utf-8")
docs = loader.load()

from langchain.text_splitter import RecursiveCharacterTextSplitter

splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,       # 每块 500 字符
    chunk_overlap=50,     # 相邻块重叠 50 字符
    separators=["\n\n", "\n", "。", "！", "？", " "],
)

chunks = splitter.split_documents(docs)

from langchain_openai import OpenAIEmbeddings

embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

# 将文本转为向量
vector = embeddings.embed_query("公司的年假政策")
# 返回：[0.023, -0.156, 0.089, ...]  （1536 维的浮点数数组）

from langchain_community.vectorstores import Chroma

vectorstore = Chroma.from_documents(
    documents=chunks,
    embedding=embeddings,
    persist_directory="./chroma_db",  # 持久化到磁盘
)

# 检索
retriever = vectorstore.as_retriever(search_kwargs={"k": 4})
relevant_docs = retriever.invoke("年假政策是什么")

# 生成
from langchain_openai import ChatOpenAI
from langchain.chains import RetrievalQA

llm = ChatOpenAI(model="gpt-4o")
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=retriever,
)

result = qa_chain.invoke("年假政策是什么")
print(result["result"])

# 余弦相似度（最常用）
import numpy as np

def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

# 欧氏距离
def euclidean_distance(a, b):
    return np.linalg.norm(np.array(a) - np.array(b))

# 余弦相似度范围：[-1, 1]
# 1 = 完全相同
# 0 = 无关
# -1 = 完全相反

# 方案 1：按语义切分（而非固定长度）
from langchain.text_splitter import MarkdownTextSplitter

splitter = MarkdownTextSplitter(chunk_size=500, chunk_overlap=50)
# 按 Markdown 标题切分，保持语义完整性

# 方案 2：递归切分（推荐）
splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=50,
    separators=["\n\n", "\n", "。", "！", "？", ".", " "],
    # 先尝试按段落切，不行再按句子切，最后按词切

# 方案 3：父子块策略
# 小块用于检索（精确匹配），大块用于生成（完整上下文）
small_splitter = RecursiveCharacterTextSplitter(chunk_size=200)
large_splitter = RecursiveCharacterTextSplitter(chunk_size=1000)

# 方案 1：混合检索（关键词 + 语义）
from langchain.retrievers import EnsembleRetriever
from langchain_community.retrievers import BM25Retriever

bm25 = BM25Retriever.from_documents(chunks)
bm25.k = 4

semantic = vectorstore.as_retriever(search_kwargs={"k": 4})

hybrid = EnsembleRetriever(
    retrievers=[bm25, semantic],
    weights=[0.3, 0.7],  # 语义检索权重更高
)

# 方案 2：重排序（Reranking）
# 先粗检索 top 20，再用模型精排取 top 4
from langchain.retrievers import ContextualCompressionRetriever
from langchain_cohere import CohereRerank

reranker = CohereRerank(model="rerank-v3.5", top_n=4)
compression_retriever = ContextualCompressionRetriever(
    base_compressor=reranker,
    base_retriever=vectorstore.as_retriever(search_kwargs={"k": 20}),
)

# 方案 3：查询改写
# 用 LLM 改写用户问题，提高检索效果
rewrite_prompt = """
请将以下用户问题改写为更适合搜索的形式，
保留关键信息，去除口语化表达：
用户问题：{question}
改写后的搜索词：
"""

# 方案 1：引用来源
prompt = """
请根据以下参考资料回答问题。
回答时必须引用信息来源（文档名和段落编号）。

参考资料：
{context}

问题：{question}

回答格式：
回答内容...
来源：《文档名》第 X 段
"""

# 方案 2：限制范围
prompt = """
你是一个知识库助手。请严格遵守以下规则：
1. 只根据提供的参考资料回答
2. 如果参考资料中没有相关信息，回答"根据现有资料，无法回答"
3. 不要编造任何信息
4. 回答要简洁准确

参考资料：{context}
问题：{question}
"""

# 方案 3：多轮追问
# 第一轮：检索相关文档
# 第二轮：让 AI 评估文档是否足够回答问题
# 第三轮：如果不够，让 AI 提出需要补充的信息