RAG 高级优化：混合检索、重排序、查询改写与评测闭环

Naive RAG                          Advanced RAG
┌──────────────┐                   ┌──────────────────────────────┐
│ 用户查询      │                   │ 用户查询                      │
└──────┬───────┘                   └──────────┬───────────────────┘
       ▼                                      ▼
┌──────────────┐                   ┌──────────────────────────────┐
│  Embedding   │                   │  查询改写 (Query Rewrite)     │
└──────┬───────┘                   │  HyDE / Multi-Query / Step   │
       ▼                           └──────────┬───────────────────┘
┌──────────────┐                              ▼
│  向量检索     │                   ┌──────────────────────────────┐
│  Top-K       │                   │  混合检索 (Hybrid Search)     │
└──────┬───────┘                   │  向量 + BM25 + RRF 融合       │
       ▼                           └──────────┬───────────────────┘
┌──────────────┐                              ▼
│  拼接 Prompt  │                   ┌──────────────────────────────┐
│  LLM 生成    │                   │  重排序 (Reranking)           │
└──────────────┘                   │  Cross-Encoder 精排            │
                                   └──────────┬───────────────────┘
                                              ▼
                                   ┌──────────────────────────────┐
                                   │  上下文优化                    │
                                   │  压缩 / 去重 / 相关性过滤       │
                                   └──────────┬───────────────────┘
                                              ▼
                                   ┌──────────────────────────────┐
                                   │  生成优化 + 引用溯源           │
                                   └──────────┬───────────────────┘
                                              ▼
                                   ┌──────────────────────────────┐
                                   │  评测闭环 (RAGAS / TruLens)   │
                                   └──────────────────────────────┘

import numpy as np
from dataclasses import dataclass, field


@dataclass
class SearchResult:
    doc_id: str
    content: str
    score: float
    rank: int = 0


def reciprocal_rank_fusion(
    result_lists: list[list[SearchResult]],
    k: int = 60,
    top_n: int = 10
) -> list[SearchResult]:
    doc_scores: dict[str, float] = {}
    doc_map: dict[str, SearchResult] = {}

    for results in result_lists:
        for result in results:
            rrf_score = 1.0 / (k + result.rank)
            doc_scores[result.doc_id] = doc_scores.get(result.doc_id, 0.0) + rrf_score
            if result.doc_id not in doc_map:
                doc_map[result.doc_id] = result

    sorted_docs = sorted(doc_scores.items(), key=lambda x: x[1], reverse=True)[:top_n]

    fused_results = []
    for doc_id, score in sorted_docs:
        result = doc_map[doc_id]
        result.score = score
        fused_results.append(result)

    return fused_results

from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np


class HybridRetriever:
    def __init__(self, documents: list[dict]):
        self.documents = documents
        self.contents = [doc["content"] for doc in documents]

        tokenized_corpus = [list(content) for content in self.contents]
        self.bm25 = BM25Okapi(tokenized_corpus)

        self.embedding_model = SentenceTransformer("BAAI/bge-large-zh-v1.5")
        embeddings = self.embedding_model.encode(self.contents, normalize_embeddings=True)
        dimension = embeddings.shape[1]

        self.index = faiss.IndexFlatIP(dimension)
        self.index.add(embeddings.astype("float32"))

    def bm25_search(self, query: str, top_k: int = 20) -> list[SearchResult]:
        tokenized_query = list(query)
        scores = self.bm25.get_scores(tokenized_query)
        top_indices = np.argsort(scores)[::-1][:top_k]

        results = []
        for rank, idx in enumerate(top_indices):
            results.append(SearchResult(
                doc_id=str(idx),
                content=self.contents[idx],
                score=float(scores[idx]),
                rank=rank + 1
            ))
        return results

    def vector_search(self, query: str, top_k: int = 20) -> list[SearchResult]:
        query_embedding = self.embedding_model.encode([query], normalize_embeddings=True)
        distances, indices = self.index.search(query_embedding.astype("float32"), top_k)

        results = []
        for rank, (idx, score) in enumerate(zip(indices[0], distances[0])):
            if idx < 0:
                continue
            results.append(SearchResult(
                doc_id=str(idx),
                content=self.contents[idx],
                score=float(score),
                rank=rank + 1
            ))
        return results

    def hybrid_search(self, query: str, top_k: int = 10) -> list[SearchResult]:
        bm25_results = self.bm25_search(query, top_k=20)
        vector_results = self.vector_search(query, top_k=20)
        return reciprocal_rank_fusion([bm25_results, vector_results], top_n=top_k)

def weighted_rrf(
    result_lists: list[tuple[list[SearchResult], float]],
    k: int = 60,
    top_n: int = 10
) -> list[SearchResult]:
    doc_scores: dict[str, float] = {}
    doc_map: dict[str, SearchResult] = {}

    for results, weight in result_lists:
        for result in results:
            rrf_score = weight / (k + result.rank)
            doc_scores[result.doc_id] = doc_scores.get(result.doc_id, 0.0) + rrf_score
            if result.doc_id not in doc_map:
                doc_map[result.doc_id] = result

    sorted_docs = sorted(doc_scores.items(), key=lambda x: x[1], reverse=True)[:top_n]

    return [
        SearchResult(doc_id=did, content=doc_map[did].content, score=score)
        for did, score in sorted_docs
    ]

bm25_results = retriever.bm25_search(query, top_k=20)
vector_results = retriever.vector_search(query, top_k=20)
fused = weighted_rrf([
    (bm25_results, 0.3),
    (vector_results, 0.7),
], top_n=10)

from sentence_transformers import CrossEncoder

cross_encoder = CrossEncoder("BAAI/bge-reranker-v2-m3", max_length=512)

query = "如何配置 Nginx 反向代理？"
candidates = [
    "Nginx 是一个高性能的 HTTP 和反向代理服务器...",
    "反向代理配置示例：server { listen 80; location / { proxy_pass http://backend; } }",
    "Docker 容器编排最佳实践...",
]

pairs = [(query, doc) for doc in candidates]
scores = cross_encoder.predict(pairs)

ranked = sorted(zip(candidates, scores), key=lambda x: x[1], reverse=True)
for doc, score in ranked:
    print(f"[{score:.4f}] {doc[:60]}...")

import cohere

co = cohere.ClientV2(api_key="your-cohere-api-key")

query = "什么是 RAG？"
documents = [
    "RAG（Retrieval-Augmented Generation）是一种结合检索和生成的架构...",
    "Transformer 是一种基于自注意力机制的神经网络架构...",
    "向量数据库用于存储和检索高维向量表示...",
]

response = co.rerank(
    model="rerank-v3.5",
    query=query,
    documents=documents,
    top_n=3,
    return_documents=True
)

for result in response.results:
    print(f"[{result.relevance_score:.4f}] {result.document.text[:80]}")

class RetrievalRerankingPipeline:
    def __init__(self, hybrid_retriever: HybridRetriever, reranker=None):
        self.retriever = hybrid_retriever
        self.reranker = reranker or CrossEncoder("BAAI/bge-reranker-v2-m3")

    def retrieve(self, query: str, initial_k: int = 20, final_k: int = 5) -> list[dict]:
        candidates = self.retriever.hybrid_search(query, top_k=initial_k)

        if not candidates:
            return []

        pairs = [(query, c.content) for c in candidates]
        scores = self.reranker.predict(pairs)

        for candidate, score in zip(candidates, scores):
            candidate.score = float(score)

        candidates.sort(key=lambda x: x.score, reverse=True)
        return candidates[:final_k]

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)


def hyde_retrieval(query: str, retriever, embedding_model, llm=llm):
    prompt = ChatPromptTemplate.from_template(
        "请写一段详细的段落来回答以下问题。\n"
        "问题：{query}\n"
        "回答："
    )
    chain = prompt | llm
    hypothetical_answer = chain.invoke({"query": query}).content

    results = retriever.hybrid_search(hypothetical_answer, top_k=10)
    return results

from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.3)

QUERY_VARIANT_PROMPT = ChatPromptTemplate.from_template(
    "你是一位搜索查询优化专家。请基于以下用户查询，生成 {n} 个不同角度的搜索查询变体，"
    "每个变体应该从不同的方面或使用不同的术语来表达相同的搜索意图。\n\n"
    "原始查询：{query}\n\n"
    "请以 JSON 数组格式输出，例如：[\"变体1\", \"变体2\", \"变体3\"]"
)


def generate_query_variants(query: str, n: int = 3) -> list[str]:
    chain = QUERY_VARIANT_PROMPT | llm
    response = chain.invoke({"query": query, "n": n})

    import json
    try:
        variants = json.loads(response.content)
    except json.JSONDecodeError:
        variants = [line.strip() for line in response.content.split("\n") if line.strip()]

    return [query] + variants[:n]


def multi_query_retrieval(query: str, retriever, n_variants: int = 3) -> list[SearchResult]:
    queries = generate_query_variants(query, n=n_variants)

    all_results: list[list[SearchResult]] = []
    for q in queries:
        results = retriever.hybrid_search(q, top_k=10)
        all_results.append(results)

    fused = reciprocal_rank_fusion(all_results, top_n=10)
    return fused

STEP_BACK_PROMPT = ChatPromptTemplate.from_template(
    "你是一位搜索策略专家。给定以下具体查询，请生成一个更宽泛、更抽象的查询，"
    "以获取解决该问题所需的背景知识。\n\n"
    "原始查询：{query}\n\n"
    "只输出抽象后的查询，不要其他内容："
)


def step_back_retrieval(query: str, retriever, llm=llm):
    chain = STEP_BACK_PROMPT | llm
    abstract_query = chain.invoke({"query": query}).content.strip()

    specific_results = retriever.hybrid_search(query, top_k=10)
    broad_results = retriever.hybrid_search(abstract_query, top_k=10)

    fused = reciprocal_rank_fusion([specific_results, broad_results], top_n=10)
    return fused

DECOMPOSE_PROMPT = ChatPromptTemplate.from_template(
    "请将以下复合查询分解为多个独立的子查询。每个子查询应该只关注一个具体的信息点。\n\n"
    "复合查询：{query}\n\n"
    "以 JSON 数组格式输出子查询列表：[\"子查询1\", \"子查询2\", ...]"
)


def decompose_and_retrieve(query: str, retriever, llm=llm) -> dict:
    chain = DECOMPOSE_PROMPT | llm
    response = chain.invoke({"query": query})

    import json
    try:
        sub_queries = json.loads(response.content)
    except json.JSONDecodeError:
        sub_queries = [query]

    results_by_query: dict[str, list[SearchResult]] = {}
    for sq in sub_queries:
        results = retriever.hybrid_search(sq, top_k=5)
        results_by_query[sq] = results

    all_results = list(results_by_query.values())
    fused = reciprocal_rank_fusion(all_results, top_n=10)

    return {
        "sub_queries": sub_queries,
        "results_by_query": results_by_query,
        "fused_results": fused
    }

from llmlingua import PromptCompressor

prompt_compressor = PromptCompressor(
    model_name="microsoft/llmlingua-2-xlm-roberta-large-meetingbank",
    device_map="cpu"
)


def compress_context(query: str, contexts: list[str], rate: float = 0.5) -> str:
    full_context = "\n\n".join(contexts)

    compressed = prompt_compressor.compress_prompt(
        full_context,
        rate=rate,
        question=query,
        force_tokens=["\n", "?", "!", "."],
        drop_consecutive=True
    )

    return compressed["compressed_prompt"]

import hashlib
from difflib import SequenceMatcher


def deduplicate_contexts(
    contexts: list[str],
    similarity_threshold: float = 0.85
) -> list[str]:
    seen_hashes = set()
    unique_contexts = []

    for ctx in contexts:
        content_hash = hashlib.md5(ctx.strip().encode()).hexdigest()
        if content_hash in seen_hashes:
            continue

        is_duplicate = False
        for existing in unique_contexts:
            similarity = SequenceMatcher(None, ctx, existing).ratio()
            if similarity > similarity_threshold:
                is_duplicate = True
                break

        if not is_duplicate:
            seen_hashes.add(content_hash)
            unique_contexts.append(ctx)

    return unique_contexts


def relevance_filter(
    query: str,
    contexts: list[str],
    scores: list[float],
    threshold: float = 0.3
) -> list[str]:
    return [ctx for ctx, score in zip(contexts, scores) if score >= threshold]

def mitigate_lost_in_middle(contexts: list[str], query: str) -> list[str]:
    if len(contexts) <= 2:
        return contexts

    sorted_by_relevance = sorted(
        enumerate(contexts),
        key=lambda x: len(x[1]),
        reverse=True
    )

    reordered = []
    left, right = 0, len(sorted_by_relevance) - 1
    pick_left = True

    for _, ctx in sorted_by_relevance:
        if pick_left:
            reordered.insert(0, ctx)
        else:
            reordered.append(ctx)
        pick_left = not pick_left

    return reordered

以下是一组参考文档，请仔细阅读所有文档后回答问题。特别注意文档 2 和文档 3 的内容。

[文档 1] ...
[文档 2] ...（可能被忽略的关键信息在这里）
[文档 3] ...（同样重要但容易被跳过的信息）
[文档 4] ...

CITATION_PROMPT = """你是一位严谨的技术文档回答者。请根据以下参考文档回答用户问题。

要求：
1. 每个论断后必须标注来源，格式为 [文档X]
2. 如果参考文档不足以完整回答问题，明确说明哪些部分是基于参考文档的，哪些部分信息不足
3. 不要编造参考文档中不存在的信息

参考文档：
{contexts}

用户问题：{question}

回答："""


def generate_with_citations(
    query: str,
    retrieved_docs: list[dict],
    llm=llm
) -> str:
    contexts = "\n\n".join(
        f"[文档{i+1}] {doc['content']}"
        for i, doc in enumerate(retrieved_docs)
    )

    prompt = CITATION_PROMPT.format(contexts=contexts, question=query)
    response = llm.invoke([("user", prompt)])
    return response.content

VERIFICATION_PROMPT = """你是一位事实核查专家。请判断以下回答是否忠实于给定的参考文档。

参考文档：
{contexts}

待验证回答：
{answer}

请逐条检查回答中的每个论断，判断是否有参考文档支撑。输出 JSON 格式：
{{
    "is_faithful": true/false,
    "unsupported_claims": ["不被支撑的论断1", ...],
    "confidence": 0.0-1.0
}}"""


def verify_answer(answer: str, contexts: str, llm=llm) -> dict:
    prompt = VERIFICATION_PROMPT.format(contexts=contexts, answer=answer)
    response = llm.invoke([("user", prompt)])

    import json
    try:
        return json.loads(response.content)
    except json.JSONDecodeError:
        return {"is_faithful": False, "unsupported_claims": [], "confidence": 0.0}

def generate_with_voting(
    query: str,
    retrieved_docs: list[dict],
    n_paths: int = 3,
    llm=llm
) -> str:
    candidates = []
    for _ in range(n_paths):
        answer = generate_with_citations(query, retrieved_docs, llm=llm)
        candidates.append(answer)

    judge_prompt = (
        f"以下是针对同一问题的三个回答，请选出最准确、最完整的一个。\n\n"
        f"问题：{query}\n\n"
    )
    for i, c in enumerate(candidates):
        judge_prompt += f"回答 {i+1}：\n{c}\n\n"
    judge_prompt += "请输出最佳回答的编号（1、2 或 3）以及选择理由。"

    response = llm.invoke([("user", judge_prompt)])

    import re
    match = re.search(r"[123]", response.content)
    best_idx = int(match.group()) - 1 if match else 0

    return candidates[min(best_idx, len(candidates) - 1)]

from ragas import evaluate
from ragas.metrics import (
    faithfulness,
    answer_relevancy,
    context_precision,
    context_recall,
)
from datasets import Dataset

eval_data = {
    "question": [
        "什么是 RAG？",
        "向量数据库有哪些？",
        "如何优化检索质量？",
    ],
    "answer": [
        "RAG 是检索增强生成...",
        "常见的向量数据库包括 Pinecone、Weaviate、Milvus...",
        "可以通过混合检索、重排序和查询改写来优化检索质量...",
    ],
    "contexts": [
        ["RAG（Retrieval-Augmented Generation）是一种将检索与生成结合的架构..."],
        ["向量数据库是专门用于存储和检索高维向量的数据库系统..."],
        ["检索质量优化包括：混合检索、Cross-Encoder 重排序、查询改写..."],
    ],
    "ground_truth": [
        "RAG 是检索增强生成，结合外部知识库与 LLM 生成能力",
        "Pinecone、Weaviate、Milvus、Qdrant、Chroma",
        "混合检索、重排序、查询改写",
    ],
}

dataset = Dataset.from_dict(eval_data)

result = evaluate(
    dataset=dataset,
    metrics=[faithfulness, answer_relevancy, context_precision, context_recall],
)

print(result)
# {'faithfulness': 0.85, 'answer_relevancy': 0.92,
#  'context_precision': 0.78, 'context_recall': 0.88}

from trulens.core import Feedback, TruSession
from trulens.providers.openai import OpenAI

session = TruSession()
provider = OpenAI(model_engine="gpt-4o-mini")

f_relevance = Feedback(
    provider.relevance_with_cot_reasons,
    name="Relevance"
).on_input_output()

f_groundedness = Feedback(
    provider.groundedness_measure_with_cot_reasons,
    name="Groundedness"
).on_context()

f_answer_relevance = Feedback(
    provider.relevance_with_cot_reasons,
    name="Answer Relevance"
).on_input_output()

import json
import random


class RAGEvalDatasetBuilder:
    def __init__(self):
        self.evaluation_pairs = []

    def add_pair(
        self,
        question: str,
        ground_truth_answer: str,
        relevant_doc_ids: list[str],
        difficulty: str = "medium",
        category: str = "general"
    ):
        self.evaluation_pairs.append({
            "question": question,
            "ground_truth": ground_truth_answer,
            "relevant_doc_ids": relevant_doc_ids,
            "difficulty": difficulty,
            "category": category,
        })

    def generate_negative_samples(
        self,
        all_doc_ids: list[str],
        n_per_positive: int = 2
    ) -> list[dict]:
        negatives = []
        for pair in self.evaluation_pairs:
            irrelevant_ids = [
                did for did in all_doc_ids
                if did not in pair["relevant_doc_ids"]
            ]
            sampled = random.sample(
                irrelevant_ids,
                min(n_per_positive, len(irrelevant_ids))
            )
            negatives.append({
                "question": pair["question"],
                "type": "negative_retrieval",
                "expected_irrelevant": sampled,
            })
        return negatives

    def export(self, path: str):
        with open(path, "w", encoding="utf-8") as f:
            json.dump(self.evaluation_pairs, f, ensure_ascii=False, indent=2)

    def compute_difficulty_distribution(self) -> dict[str, int]:
        dist = {}
        for pair in self.evaluation_pairs:
            d = pair["difficulty"]
            dist[d] = dist.get(d, 0) + 1
        return dist

┌─────────────────────────────────────────────────────────────┐
│                    RAG 优化闭环                               │
│                                                             │
│  ┌──────────┐     ┌──────────┐     ┌──────────┐            │
│  │  评测     │────▶│  分析     │────▶│  调优     │            │
│  │ Eval      │     │ Analysis │     │ Tuning   │            │
│  └──────────┘     └──────────┘     └──────────┘            │
│       ▲                                    │                │
│       │              ┌──────────┐          │                │
│       └──────────────│  重评测   │◀─────────┘                │
│                      │ Re-eval  │                           │
│                      └──────────┘                           │
└─────────────────────────────────────────────────────────────┘

class RAGOptimizationLoop:
    def __init__(self, pipeline, eval_dataset, llm=llm):
        self.pipeline = pipeline
        self.eval_dataset = eval_dataset
        self.llm = llm
        self.history = []

    def evaluate(self) -> dict:
        results = {
            "faithfulness": [],
            "answer_relevancy": [],
            "context_precision": [],
            "context_recall": [],
        }

        for item in self.eval_dataset:
            retrieved = self.pipeline.retrieve(item["question"])
            contexts = [r.content for r in retrieved]

            generated = generate_with_citations(item["question"], retrieved, llm=self.llm)

            eval_result = evaluate(
                Dataset.from_dict({
                    "question": [item["question"]],
                    "answer": [generated],
                    "contexts": [contexts],
                    "ground_truth": [item["ground_truth"]],
                }),
                metrics=[faithfulness, answer_relevancy, context_precision, context_recall],
            )

            for key in results:
                results[key].append(eval_result[key])

        avg_results = {k: sum(v) / len(v) for k, v in results.items()}
        self.history.append(avg_results)
        return avg_results

    def diagnose(self, results: dict) -> list[str]:
        recommendations = []

        if results["context_recall"] < 0.7:
            recommendations.append(
                "上下文召回率偏低：建议启用混合检索或增加检索 Top-K 数量"
            )
        if results["context_precision"] < 0.6:
            recommendations.append(
                "上下文精确度偏低：建议加入重排序环节或收紧相关性阈值"
            )
        if results["faithfulness"] < 0.8:
            recommendations.append(
                "忠实度偏低：建议优化 Prompt 中的引用约束或切换更强的生成模型"
            )
        if results["answer_relevancy"] < 0.7:
            recommendations.append(
                "答案相关性偏低：建议在 Prompt 中强调紧扣问题回答"
            )

        return recommendations

    def run_optimization_cycle(self) -> dict:
        print("=" * 50)
        print(f"开始第 {len(self.history) + 1} 轮评测...")
        results = self.evaluate()

        print(f"\n评测结果：")
        for metric, score in results.items():
            print(f"  {metric}: {score:.4f}")

        recommendations = self.diagnose(results)
        if recommendations:
            print(f"\n优化建议：")
            for i, rec in enumerate(recommendations, 1):
                print(f"  {i}. {rec}")
        else:
            print("\n所有指标达标，系统状态良好。")

        return results