RAG 添加重排序 (Reranking)

在你的分层架构中集成重排序（Reranking）同样非常干净——只需在 RAGService 的 ask 方法中，检索合并后与生成前插入一个“重排序”步骤。

我会为你提供两种主流方案的集成代码，你可以根据预算和隐私需求选择：

🔧 方案一：免费、本地部署（推荐个人开发者）

使用 BGE-reranker（由智源开源），效果媲美 Cohere，数据不出网，完全免费。

1. 安装依赖

pip install FlagEmbedding

2. 创建重排序模块

在 src/infrastructure/rerankers/ 目录下新建 bge_reranker.py：

# src/infrastructure/rerankers/bge_reranker.py
from typing import List, Tuple
from langchain_core.documents import Document
from FlagEmbedding import FlagReranker

class BGEReranker:
    def __init__(self, model_name: str = "BAAI/bge-reranker-v2-m3", use_fp16: bool = True):
        """
        BGE 重排序模型
        - model_name: 可选 "BAAI/bge-reranker-v2-m3"（多语言，推荐）或 "BAAI/bge-reranker-large"
        - use_fp16: 节省显存和内存
        """
        self.reranker = FlagReranker(model_name, use_fp16=use_fp16)

    def rerank(self, query: str, documents: List[Document], top_n: int = 3) -> List[Document]:
        """
        根据相关性对文档重新排序，返回 top_n 个最相关的文档
        """
        if not documents:
            return []

        # 构造 (query, document) 对
        pairs = [[query, doc.page_content] for doc in documents]
        # 计算相关性分数
        scores = self.reranker.compute_score(pairs, normalize=True)  # 归一化到 0~1

        # 按分数降序排序，取前 top_n
        sorted_pairs = sorted(zip(documents, scores), key=lambda x: x[1], reverse=True)
        reranked_docs = [doc for doc, _ in sorted_pairs[:top_n]]

        print(f"🔄 重排序完成：从 {len(documents)} 个文档中选出 Top-{len(reranked_docs)}")
        return reranked_docs

3. 集成到 RAGService

修改 src/application/services/rag_service.py，在 __init__ 中初始化 Reranker，并在 ask 中调用：

# src/application/services/rag_service.py
from typing import List, Tuple
from langchain_core.documents import Document
from langchain_core.prompts import PromptTemplate
from src.application.services.query_rewriter import QueryRewriter
from src.infrastructure.rerankers.bge_reranker import BGEReranker  # 新增

class RAGService:
    def __init__(self, llm_client, retriever, top_k_rerank: int = 3):
        self.llm = llm_client
        self.retriever = retriever
        self.rewriter = QueryRewriter(llm_client)
        # ✅ 初始化重排序器
        self.reranker = BGEReranker()
        self.top_k_rerank = top_k_rerank

        self.prompt_template = PromptTemplate.from_template(
            "你是一个专业的AI知识助手。请基于以下上下文回答用户的问题。\n"
            "如果上下文不足以回答，请如实说明，不要编造。\n\n"
            "上下文：\n{context}\n\n"
            "用户问题：{question}\n\n"
            "回答："
        )

    def format_docs(self, docs: List[Document]) -> str:
        return "\n\n".join([doc.page_content for doc in docs])

    def ask(self, query: str) -> Tuple[str, List[Document]]:
        # 1. 查询改写
        rewritten_queries = self.rewriter.rewrite(query)

        # 2. 多路检索 + 合并去重
        all_docs = []
        seen_contents = set()
        for q in rewritten_queries:
            docs = self.retriever.invoke(q)
            for doc in docs:
                content = doc.page_content.strip()
                if content and content not in seen_contents:
                    seen_contents.add(content)
                    all_docs.append(doc)
        print(f"🔍 多查询检索后，去重合并共召回 {len(all_docs)} 个片段")

        # ---------- 3. 新增：重排序 ----------
        if all_docs:
            all_docs = self.reranker.rerank(query, all_docs, top_n=self.top_k_rerank)

        # 4. 构建上下文并生成
        context = self.format_docs(all_docs)
        final_prompt = self.prompt_template.format(context=context, question=query)
        try:
            response = self.llm.invoke(final_prompt)
            answer = response.content
        except Exception as e:
            answer = f"生成失败: {e}"

        return answer, all_docs

4. 在 main.py 中微调（可选）

你可以在 RAGService 的构造中增加参数控制 top_k_rerank（最终送入 LLM 的文档数），例如：

rag_service = RAGService(
    llm_client=llm,
    retriever=ensemble_retriever,
    top_k_rerank=3   # 只取最相关的3个片段给大模型
)

💰 方案二：Cohere Rerank API（企业级，付费，但效果顶级）

如果你预算充足，可以换成 Cohere 的 API（每月有免费额度）。

安装

pip install cohere

创建 cohere_reranker.py

# src/infrastructure/rerankers/cohere_reranker.py
import cohere
from typing import List
from langchain_core.documents import Document

class CohereReranker:
    def __init__(self, api_key: str, model: str = "rerank-multilingual-v3.0"):
        self.client = cohere.Client(api_key)
        self.model = model

    def rerank(self, query: str, documents: List[Document], top_n: int = 3) -> List[Document]:
        if not documents:
            return []
        # 准备文档列表（只传文本内容）
        doc_texts = [doc.page_content for doc in documents]
        response = self.client.rerank(
            query=query,
            documents=doc_texts,
            model=self.model,
            top_n=top_n
        )
        # 按返回的索引重排文档
        reranked_docs = [documents[res.index] for res in response.results]
        print(f"🔄 Cohere 重排序完成，保留 Top-{len(reranked_docs)}")
        return reranked_docs

然后在 RAGService 中替换 BGEReranker() 为 CohereReranker(api_key=settings.COHERE_API_KEY)（记得在 settings.py 中添加该配置）。

🧪 验证效果

运行你的 Gradio 界面，输入一个复杂问题，观察终端输出：

🔍 多查询检索后，去重合并共召回 12 个片段
🔄 重排序完成：从 12 个文档中选出 Top-3

你可以在 respond 中打印重排序后的文档来源，会发现相关性明显提升。

💡 进阶调优

• 调整检索器召回数量：在你的 ensemble_retriever 中，k 值可以调大（如 k=10），让重排序器有更多候选，然后重排序只取前 top_n=3。这样能在最终答案质量与延迟间取得平衡。
• 缓存模型：BGE-reranker 首次加载会下载约 2GB 模型，请耐心等待。之后会自动缓存到本地。

现在你的 RAG 机器人已经拥有 “改写 → 检索 → 重排序 → 生成” 的完整工业级链路，距离“企业级智能 Agent”又近了一大步！🚀 如果运行中遇到模型加载慢或内存不足，随时告诉我，我教你用更小的模型。

本作品采用《CC 协议》，转载必须注明作者和本文链接