RAG 检索增强生成实战：从文档切分到向量检索完整方案

1
2
3
4

用户提问 → 查询理解 → 向量检索 → 重排序 → 上下文构建 → Prompt 组装 → 大模型生成 → 回答
                ↑
           文档处理（离线）：
           文档加载 → 文档切分 → 向量化 → 存入向量数据库

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

from langchain_community.document_loaders import (
    PyPDFLoader, Docx2txtLoader, TextLoader, UnstructuredHTMLLoader
)

def load_document(file_path: str):
    ext = file_path.split('.')[-1].lower()
    loaders = {
        'pdf': PyPDFLoader,
        'docx': Docx2txtLoader,
        'doc': Docx2txtLoader,
        'txt': TextLoader,
        'md': TextLoader,
        'html': UnstructuredHTMLLoader,
    }
    loader_class = loaders.get(ext, TextLoader)
    loader = loader_class(file_path)
    return loader.load()

1
2
3
4
5
6
7
8
9

from langchain.text_splitter import RecursiveCharacterTextSplitter

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,        # 每个 chunk 500 字符
    chunk_overlap=50,      # 相邻 chunk 重叠 50 字符，防止上下文断裂
    separators=["\n\n", "\n", "。", "！", "？", ".", " ", ""],  # 优先在段落、句子边界切分
)

chunks = text_splitter.split_documents(documents)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

from langchain_community.embeddings import HuggingFaceBgeEmbeddings

# 用 BGE 中文 Embedding 模型，效果比 OpenAI 的在中文上好
embeddings = HuggingFaceBgeEmbeddings(
    model_name="BAAI/bge-large-zh-v1.5",
    model_kwargs={'device': 'cuda'},
    encode_kwargs={'normalize_embeddings': True},  # 归一化，方便余弦相似度
)

# 批量向量化
texts = [chunk.page_content for chunk in chunks]
vectors = embeddings.embed_documents(texts)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

from langchain_community.vectorstores import Milvus

# 存入 Milvus
vector_store = Milvus.from_documents(
    documents=chunks,
    embedding=embeddings,
    connection_args={"host": "milvus", "port": "19530"},
    collection_name="knowledge_base",
)

# 也可以用 pgvector（PostgreSQL 插件，不用单独维护向量数据库）
# from langchain_community.vectorstores import PGVector
# vector_store = PGVector.from_documents(...)

1
2
3
4

def retrieve(query: str, top_k: int = 5):
    # 向量相似度检索
    docs = vector_store.similarity_search(query, k=top_k)
    return docs

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

def query_understanding(query: str) -> str:
    """查询优化：提取关键词、扩展同义词、纠正错别字"""
    prompt = f"""
    你是一个搜索查询优化专家。请对以下用户问题进行优化，使其更适合向量检索。
    要求：
    1. 提取核心关键词
    2. 扩展相关同义词
    3. 去除无关词汇
    4. 输出优化后的查询语句，不要解释

    用户问题：{query}
    优化后查询：
    """
    optimized = llm.invoke(prompt).content
    return optimized

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

from langchain.retrievers import EnsembleRetriever
from langchain_community.retrievers import BM25Retriever

# BM25 关键词检索
bm25_retriever = BM25Retriever.from_documents(chunks)
bm25_retriever.k = 5

# 向量检索
vector_retriever = vector_store.as_retriever(search_kwargs={"k": 5})

# 混合检索（RRF 融合）
ensemble_retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, vector_retriever],
    weights=[0.5, 0.5],
)

docs = ensemble_retriever.get_relevant_documents(query)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

from sentence_transformers import CrossEncoder

reranker = CrossEncoder('BAAI/bge-reranker-large')

def rerank(query: str, docs, top_n: int = 3):
    # 计算 query 和每个 doc 的相关性分数
    pairs = [(query, doc.page_content) for doc in docs]
    scores = reranker.predict(pairs)

    # 按分数排序，取 top_n
    scored_docs = list(zip(docs, scores))
    scored_docs.sort(key=lambda x: x[1], reverse=True)
    return [doc for doc, score in scored_docs[:top_n]]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

def build_prompt(query: str, docs) -> str:
    # 把检索到的文档格式化成上下文
    context = "\n\n".join([
        f"【文档{i+1}】{doc.page_content}\n（来源：{doc.metadata.get('source', '未知')}）"
        for i, doc in enumerate(docs)
    ])

    prompt = f"""
    你是一个专业的招聘平台知识库助手。请根据以下参考文档回答用户的问题。

    【参考文档】
    {context}

    【用户问题】
    {query}

    【回答要求】
    1. 只根据参考文档回答，不要编造信息
    2. 如果参考文档里没有相关信息，回答"根据现有知识库无法回答该问题"
    3. 回答要简洁准确，引用文档时标注来源
    4. 用中文回答
    """
    return prompt

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

def rag_answer(query: str) -> str:
    # 1. 查询优化
    optimized_query = query_understanding(query)

    # 2. 混合检索
    docs = ensemble_retriever.get_relevant_documents(optimized_query)

    # 3. 重排序
    relevant_docs = rerank(query, docs, top_n=3)

    # 4. 构建 Prompt
    prompt = build_prompt(query, relevant_docs)

    # 5. 大模型生成
    answer = llm.invoke(prompt).content

    return answer