RAG 检索增强生成实战：Embedding + Vector DB + ReRank

大模型闭卷考试会产生幻觉，RAG（检索增强生成）在生成前先从私有知识库检索最相关 Top-K 段落，再带着上下文答题，是目前让大模型接入企业私有数据最主流、最落地的生产方案。

文档分块与向量入库

分块是 RAG 效果的第一道闸门：太小丢上下文，太大污染检索。实际生产通常 200-800 tokens 加 15-20% 重叠窗口；按自然段/标题边界切开比硬切字符效果强很多。用 BGE-M3 / E5-large-v2 这类开源中英双语 Embedding 效果远超 openai-ada-002。

from __future__ import annotations
import json, re
from dataclasses import dataclass
from typing import Iterable
import numpy as np
import psycopg2.extras
from sentence_transformers import SentenceTransformer, CrossEncoder

CHUNK_SIZE = 512
CHUNK_OVERLAP = 80
EMBED_MODEL = "BAAI/bge-m3"
RERANK_MODEL = "BAAI/bge-reranker-v2-m3"
TOP_K_RETRIEVAL = 20
TOP_K_FINAL = 4

embedder = SentenceTransformer(EMBED_MODEL)
reranker = CrossEncoder(RERANK_MODEL)

@dataclass
class Chunk:
    doc_id: str
    chunk_idx: int
    text: str
    metadata: dict

def smart_chunk(doc_id: str, text: str, meta: dict | None = None) -> list[Chunk]:
    paragraphs = re.split(r"\n\s*\n", text.strip())
    chunks: list[Chunk] = []
    buf, buf_len, idx = [], 0, 0
    for p in paragraphs:
        p_len = len(p.split())
        if buf and buf_len + p_len > CHUNK_SIZE:
            chunks.append(Chunk(doc_id, idx, "\n\n".join(buf), meta or {}))
            idx += 1
            drop = max(1, len(buf) - (CHUNK_OVERLAP * len(buf) // max(1, buf_len)))
            buf, buf_len = buf[drop:], sum(len(b.split()) for b in buf[drop:])
        buf.append(p); buf_len += p_len
    if buf: chunks.append(Chunk(doc_id, idx, "\n\n".join(buf), meta or {}))
    return chunks

def pgvector_store(conn, chunks: Iterable[Chunk]) -> None:
    cur = conn.cursor()
    cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
    cur.execute(f"CREATE TABLE IF NOT EXISTS docs (id BIGSERIAL PRIMARY KEY, doc_id TEXT NOT NULL, chunk_idx INT NOT NULL, content TEXT NOT NULL, metadata JSONB NOT NULL DEFAULT '{{}}'::jsonb, embedding vector(1024) NOT NULL, UNIQUE(doc_id, chunk_idx))")
    rows = []
    for c in chunks:
        emb = embedder.encode(c.text, normalize_embeddings=True)
        rows.append((c.doc_id, c.chunk_idx, c.text, json.dumps(c.metadata or {}), psycopg2.extras.Json(emb.tolist())))
    psycopg2.extras.execute_batch(cur, "INSERT INTO docs(doc_id, chunk_idx, content, metadata, embedding) VALUES (%s,%s,%s,%s,%s) ON CONFLICT DO NOTHING", rows)
    conn.commit()

def hybrid_search(conn, query: str, top_k: int = TOP_K_RETRIEVAL, category: str | None = None):
    q_emb = embedder.encode(query, normalize_embeddings=True).tolist()
    cur = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
    where_sql = "AND metadata->>'category' = %s" if category else ""
    params = [q_emb, top_k] + ([category] if category else [])
    cur.execute(f"SELECT id, content, metadata, 1 - (embedding <=> %s::vector) AS score FROM docs WHERE true {where_sql} ORDER BY embedding <=> %s::vector LIMIT %s", params)
    return [(h["content"], {"id": h["id"], "meta": dict(h["metadata"]), "vec_score": float(h["score"])}) for h in cur.fetchall()]

def rag_answer(conn, query: str, **filters) -> dict:
    raw_hits = hybrid_search(conn, query, **filters)
    if not raw_hits: return {"answer": "抱歉，未找到相关资料。", "citations": []}
    texts = [h[0] for h in raw_hits]
    scores = reranker.predict([(query, t) for t in texts], apply_softmax=True)
    ranked = sorted(zip(scores, raw_hits), reverse=True)[:TOP_K_FINAL]
    contexts = [f"[{i+1}] {t[1][0]}" for i, t in enumerate(ranked)]
    citations = [{"id": t[1][1]["id"], "rerank_score": float(t[0]), **t[1][1]["meta"]} for t in ranked]
    prompt = f"基于以下参考资料回答问题，每句结论必须标注引用编号[1]-[{TOP_K_FINAL}]；资料外的内容用'根据所提供资料无法确定'回答。\n参考资料：\n{chr(10).join(contexts)}\n\n问题：{query}\n带引用的答案："
    return {"answer_prompt": prompt, "citations": citations}

检索+重排两阶段流水线

向量库 ANN 召回 Top 20-50 保证召回率（不要一开始就取 Top3），再用 CrossEncoder（Reranker）精排到 Top 4 喂给大模型，综合相关性质量提升 20-40%。最终给 Prompt 注入参考资料编号，要求回答强制标注引用编号，可追溯防幻觉。

模块	朴素方案（Demo）	生产方案（推荐）	质量提升
分块	硬切 1000 字符	自然段边界+512词+15%重叠	+20-35%
Embedding	openai text-ada-002	BAAI/bge-m3 多语言	+15-25%
检索	向量余弦 Top3	ANN Top30 + Metadata过滤	+10-20%
精修排序	不做	CrossEncoder Reranker	+20-40%
回答生成	直接拼接context	引用编号强制溯源+拒答兜底	-40-60%幻觉率