用 RAG Fusion 多查询召回融合提升检索

# 依赖：pip install sentence-transformers==2.7.0 openai==1.30.0 numpy==1.26.4
# 单一查询召回易只命中一种表述，漏掉同义或角度不同的相关片段。
# RAG Fusion 先让 LLM 生成多个改写查询，各查询分别向量召回后，
# 用互逆排名融合(RRF, Reciprocal Rank Fusion)合并排序，比单查询召回更稳更全。
import os
import numpy as np
from sentence_transformers import SentenceTransformer
from openai import OpenAI

client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
embedder = SentenceTransformer("BAAI/bge-small-zh")

docs = [
    "降低幻觉的关键：在 prompt 中限制模型只基于给定上下文回答，禁止补充外部知识。",
    "RAG 是检索增强生成，先检索知识库再让大模型生成答案。",
    "减少大模型虚构回答，可要求模型对无依据内容明确标注'未知'。",
    "向量检索常用 FAISS，支持百万级向量秒级搜索。",
]
doc_vecs = embedder.encode(docs, normalize_embeddings=True)

def rewrite(query: str, n: int = 3) -> list[str]:
    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user",
                   "content": f"把下面问题改写成{n}个不同表述的问法，每行一个：{query}"}],
    )
    return [line.strip("- ").strip()
            for line in resp.choices[0].message.content.splitlines()
            if line.strip()]

def retrieve(query: str, k: int = 4):
    qv = embedder.encode([query], normalize_embeddings=True)[0]
    scores = doc_vecs @ qv
    order = np.argsort(-scores)[:k]
    return [(int(i), float(scores[i])) for i in order]

def rrf(all_rankings: list, k: int = 60):
    # 互逆排名融合：每个文档得分 = sum(1/(rank + k))，rank 从 1 开始
    fused = {}
    for ranking in all_rankings:
        for rank, (doc_id, _) in enumerate(ranking, start=1):
            fused[doc_id] = fused.get(doc_id, 0.0) + 1.0 / (rank + k)
    return sorted(fused.items(), key=lambda x: x[1], reverse=True)

if __name__ == "__main__":
    queries = rewrite("RAG 怎么降低幻觉？")
    rankings = [retrieve(q) for q in queries]
    fused = rrf(rankings)
    for doc_id, score in fused[:3]:
        print(round(score, 4), docs[doc_id])

0.0298 降低幻觉的关键：在 prompt 中限制模型只基于给定上下文回答，禁止补充外部知识。
0.0185 减少大模型虚构回答，可要求模型对无依据内容明确标注'未知'。
0.0149 RAG 是检索增强生成，先检索知识库再让大模型生成答案