RAG 完整指南：从基础检索到进阶架构

# 简化版 RAG 流程：建库 + 问答
class SimpleRAG:
    def __init__(self, embedding_model, vector_store, llm):
        self.embedding_model = embedding_model
        self.vector_store = vector_store
        self.llm = llm

    def build_knowledge_base(self, documents):
        """离线阶段：切块 -> 向量化 -> 入库"""
        for doc in documents:
            chunks = self.split_into_chunks(doc, chunk_size=500, overlap=50)
            vectors = self.embedding_model.encode(chunks)
            self.vector_store.add(chunks, vectors)

    def split_into_chunks(self, text, chunk_size, overlap):
        """最朴素的定长切块：每 chunk_size 个字符切一段，相邻块重叠 overlap"""
        chunks = []
        start = 0
        while start < len(text):
            end = start + chunk_size
            chunks.append(text[start:end])
            start = end - overlap
        return chunks

    def answer(self, question, top_k=5):
        """在线阶段：检索 -> 注入上下文 -> 生成"""
        query_vector = self.embedding_model.encode([question])[0]
        retrieved_chunks = self.vector_store.search(query_vector, top_k=top_k)

        context = "\n\n".join(retrieved_chunks)
        prompt = f"""请仅根据下面的资料回答问题，如果资料中没有相关信息，请明确说不知道。

资料：
{context}

问题：{question}

回答："""
        return self.llm.generate(prompt)

def hybrid_search(query, vector_store, bm25_index, top_k=10, alpha=0.5):
    """
    alpha 控制向量检索与关键词检索的权重，alpha 越大越偏向语义匹配
    """
    # 向量检索：返回 (chunk_id, 相似度分数)
    dense_results = vector_store.search(query, top_k=top_k * 2)

    # 关键词检索：返回 (chunk_id, BM25 分数)
    sparse_results = bm25_index.search(query, top_k=top_k * 2)

    # 分数归一化后加权融合（Reciprocal Rank Fusion 是更常用的融合方式）
    scores = {}
    for rank, (chunk_id, _) in enumerate(dense_results):
        scores[chunk_id] = scores.get(chunk_id, 0) + alpha * (1 / (rank + 1))
    for rank, (chunk_id, _) in enumerate(sparse_results):
        scores[chunk_id] = scores.get(chunk_id, 0) + (1 - alpha) * (1 / (rank + 1))

    ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
    return [chunk_id for chunk_id, _ in ranked[:top_k]]

def reciprocal_rank_fusion(result_lists, k=60, top_k=10):
    """
    通用 RRF 融合：输入任意多路检索结果，输出融合后的排序。
    result_lists: List[List[chunk_id]]，每一路已按相关性从高到低排好序
    k: 平滑常数，经验值 60，越大则不同排名之间的分差越平缓
    """
    fused_scores = {}
    for single_list in result_lists:
        for rank, chunk_id in enumerate(single_list):
            # 注意 rank 从 0 开始，所以用 rank + 1 表示真实名次
            fused_scores[chunk_id] = fused_scores.get(chunk_id, 0.0) + 1.0 / (k + rank + 1)

    ranked = sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)
    return [chunk_id for chunk_id, _ in ranked[:top_k]]


# 使用示例：融合向量、BM25、以及一路"标题精确匹配"的结果
def hybrid_search_rrf(query, vector_store, bm25_index, title_index, top_k=10):
    dense = [cid for cid, _ in vector_store.search(query, top_k=30)]
    sparse = [cid for cid, _ in bm25_index.search(query, top_k=30)]
    title = [cid for cid, _ in title_index.search(query, top_k=30)]
    return reciprocal_rank_fusion([dense, sparse, title], k=60, top_k=top_k)

def retrieve_with_rerank(query, vector_store, reranker, coarse_k=50, final_k=5):
    # 第一阶段：粗召回
    candidates = vector_store.search(query, top_k=coarse_k)

    # 第二阶段：精排——reranker 直接对 (query, chunk) 打相关性分
    pairs = [(query, chunk) for chunk in candidates]
    scores = reranker.predict(pairs)

    reranked = sorted(zip(candidates, scores), key=lambda x: x[1], reverse=True)
    return [chunk for chunk, _ in reranked[:final_k]]

from sentence_transformers import CrossEncoder

# 加载一个多语言重排模型（首次运行会自动下载权重）
reranker = CrossEncoder("BAAI/bge-reranker-v2-m3", max_length=512)

def rerank(query, candidates, final_k=5):
    """
    query: 用户问题（字符串）
    candidates: 粗召回回来的候选片段列表 List[str]
    返回：按相关性排序后的 (片段, 分数) 列表，取前 final_k 个
    """
    # 构造 (query, passage) 句对，模型对每一对输出一个相关性 logit
    pairs = [[query, passage] for passage in candidates]
    scores = reranker.predict(pairs)  # 返回一维分数数组

    scored = sorted(zip(candidates, scores), key=lambda x: x[1], reverse=True)
    return scored[:final_k]


# 端到端串联：混合检索粗召回 -> 交叉编码器精排
def search_pipeline(query, vector_store, bm25_index, reranker_fn):
    coarse_ids = hybrid_search(query, vector_store, bm25_index, top_k=50)
    coarse_texts = [vector_store.get_text(cid) for cid in coarse_ids]
    top = reranker_fn(query, coarse_texts, final_k=5)
    return [text for text, score in top if score > 0.0]  # 顺手用分数阈值过滤明显不相关的

import requests

def cohere_rerank(query, documents, top_n=5, api_key="YOUR_KEY"):
    """调用商用重排 API 的通用写法，返回排序后的文档下标与分数"""
    resp = requests.post(
        "https://api.cohere.com/v1/rerank",
        headers={"Authorization": f"Bearer {api_key}"},
        json={
            "model": "rerank-multilingual-v3.0",
            "query": query,
            "documents": documents,
            "top_n": top_n,
        },
        timeout=10,
    )
    resp.raise_for_status()
    results = resp.json()["results"]
    # 每个 result 含 index（对应原 documents 下标）和 relevance_score
    return [(documents[r["index"]], r["relevance_score"]) for r in results]

def hyde_retrieve(query, llm, vector_store, top_k=5):
    # 第一步：让模型生成一段"假设"能回答该问题的文档片段
    hypo_prompt = f"请写一段简短的文字来回答这个问题（即使你并不确定答案是否准确）：{query}"
    hypothetical_doc = llm.generate(hypo_prompt)

    # 第二步：用这段假想文档（而非原始问题）去做向量检索
    return vector_store.search_by_text(hypothetical_doc, top_k=top_k)

def multi_query_retrieve(query, llm, vector_store, n_variants=4, top_k=10):
    # 第一步：让模型把原问题改写成多个不同角度的等价查询
    rewrite_prompt = f"""请把下面的问题改写成 {n_variants} 个措辞不同但意思相同的检索查询，
每行一个，不要编号，不要多余解释。

原问题：{query}"""
    variants_text = llm.generate(rewrite_prompt)
    queries = [q.strip() for q in variants_text.split("\n") if q.strip()]
    queries.append(query)  # 别忘了把原始问题也加进去

    # 第二步：每个查询各检索一路，收集为多个有序列表
    result_lists = []
    for q in queries:
        ids = [cid for cid, _ in vector_store.search(q, top_k=top_k)]
        result_lists.append(ids)

    # 第三步：用 RRF 融合所有路的结果
    return reciprocal_rank_fusion(result_lists, k=60, top_k=top_k)

def decompose_and_retrieve(question, llm, retriever, top_k=4):
    # 第一步：把复合问题拆成若干可独立检索的子问题
    decompose_prompt = f"""把下面这个问题拆解成 2-4 个可以各自独立检索的子问题，
每行一个子问题，不要编号。如果问题本身很简单无需拆解，就只输出原问题。

问题：{question}"""
    sub_questions = [
        s.strip() for s in llm.generate(decompose_prompt).split("\n") if s.strip()
    ]

    # 第二步：逐个子问题检索，带上子问题本身作为材料的标注
    gathered = []
    for sub_q in sub_questions:
        chunks = retriever.search(sub_q, top_k=top_k)
        gathered.append(f"【子问题：{sub_q}】\n" + "\n".join(chunks))

    # 第三步：把所有材料 + 原问题交给模型综合作答
    context = "\n\n".join(gathered)
    final_prompt = f"""请综合下面按子问题分组的资料，完整回答用户的原始问题。

{context}

原始问题：{question}

回答："""
    return llm.generate(final_prompt)

def build_indexed_chunks(documents):
    """建库时给每个 chunk 附上结构化元数据"""
    indexed = []
    for doc in documents:
        for chunk in split_into_chunks(doc.text, chunk_size=500, overlap=50):
            indexed.append({
                "text": chunk,
                "metadata": {
                    "doc_id": doc.id,
                    "department": doc.department,   # 部门
                    "year": doc.year,               # 年份
                    "doc_type": doc.doc_type,        # 文档类型：政策/手册/财报...
                    "updated_at": doc.updated_at,    # 更新时间戳
                    "access_level": doc.access_level # 权限等级
                }
            })
    return indexed


def filtered_search(query, vector_store, filters, top_k=5):
    """带元数据过滤的检索：先按 filters 精确筛，再在子集内比相似度"""
    # 大多数向量库（如 Qdrant / Weaviate / Milvus）都支持 filter 参数
    return vector_store.search(
        query=query,
        top_k=top_k,
        filter={
            "must": [
                {"key": "year", "match": {"value": filters["year"]}},
                {"key": "department", "match": {"value": filters["department"]}},
                {"key": "access_level", "range": {"lte": filters["max_access_level"]}},
            ]
        }
    )

def contextualize_chunk(llm, full_document, chunk):
    """为每个 chunk 生成简短的文档级上下文，再拼接到 chunk 前面"""
    prompt = f"""这是完整文档：
{full_document}

这是文档中的一个片段：
{chunk}

请用1-2句话说明这个片段在整篇文档中的上下文背景（比如属于哪个章节、涉及哪个主体），
不要复述片段内容本身，只给出定位信息。"""
    context_note = llm.generate(prompt)
    return f"{context_note}\n\n{chunk}"

def build_contextual_index(llm, embedding_model, vector_store, documents):
    for doc in documents:
        chunks = split_into_chunks(doc, chunk_size=500, overlap=50)
        contextual_chunks = [contextualize_chunk(llm, doc, c) for c in chunks]
        vectors = embedding_model.encode(contextual_chunks)
        vector_store.add(contextual_chunks, vectors)

def recursive_split(text, chunk_size=500, overlap=50,
                    separators=None):
    """
    递归字符切分：按分隔符优先级从粗到细尝试，尽量在自然边界切分。
    separators 默认按"段落 -> 行 -> 句 -> 词 -> 字符"逐级降级。
    """
    if separators is None:
        separators = ["\n\n", "\n", "。", "！", "？", "；", " ", ""]

    def _split(text, seps):
        # 已经足够短，直接作为一个块返回
        if len(text) <= chunk_size or not seps:
            return [text]

        sep = seps[0]
        # 用当前分隔符切；空字符串分隔符表示按字符逐个切
        pieces = list(text) if sep == "" else text.split(sep)

        chunks, buffer = [], ""
        for piece in pieces:
            candidate = buffer + (sep if buffer else "") + piece
            if len(candidate) <= chunk_size:
                buffer = candidate
            else:
                if buffer:
                    chunks.append(buffer)
                # 单个 piece 仍然超长，用更细的分隔符递归处理
                if len(piece) > chunk_size:
                    chunks.extend(_split(piece, seps[1:]))
                    buffer = ""
                else:
                    buffer = piece
        if buffer:
            chunks.append(buffer)
        return chunks

    raw_chunks = _split(text, separators)

    # 加入相邻块重叠：把上一块尾部的 overlap 个字符接到下一块开头
    result = []
    for i, ch in enumerate(raw_chunks):
        if i > 0 and overlap > 0:
            tail = raw_chunks[i - 1][-overlap:]
            ch = tail + ch
        result.append(ch)
    return result

import tiktoken

def token_based_split(text, max_tokens=256, overlap_tokens=32,
                      encoding_name="cl100k_base"):
    """按 Token 数切分，保证每个 chunk 不超过 embedding/模型的 Token 上限"""
    enc = tiktoken.get_encoding(encoding_name)
    token_ids = enc.encode(text)

    chunks = []
    start = 0
    while start < len(token_ids):
        end = start + max_tokens
        piece_ids = token_ids[start:end]
        chunks.append(enc.decode(piece_ids))
        if end >= len(token_ids):
            break
        start = end - overlap_tokens  # 重叠 overlap_tokens 个 token，避免边界信息割裂
    return chunks

def build_sentence_window_index(document, window_size=2):
    """把文档切成句子，每个句子记录它前后 window_size 句作为返回窗口"""
    sentences = split_into_sentences(document)  # 按标点切句
    records = []
    for i, sent in enumerate(sentences):
        left = max(0, i - window_size)
        right = min(len(sentences), i + window_size + 1)
        records.append({
            "text": sent,                             # 用这句话做向量化和检索
            "window": "".join(sentences[left:right]),  # 命中后实际返回这个窗口
            "position": i,
        })
    return records


def sentence_window_retrieve(query, vector_store, top_k=5):
    """检索命中的是单句，但取出它预存的上下文窗口交给模型"""
    hits = vector_store.search(query, top_k=top_k)
    return [hit["window"] for hit in hits]  # 返回窗口而非单句

from sentence_transformers import SentenceTransformer

# 自托管一个多语言 embedding 模型
model = SentenceTransformer("BAAI/bge-m3")  # 支持中英等多语言、最长 8192 token

def embed_documents(texts):
    # normalize_embeddings=True 让向量归一化，之后可直接用点积当余弦相似度
    return model.encode(texts, normalize_embeddings=True, batch_size=32)

def embed_query(query):
    # 有些模型建议给"查询"加特定指令前缀以提升检索效果，务必看模型卡说明
    return model.encode(
        [f"为这个句子生成用于检索的表示：{query}"],
        normalize_embeddings=True,
    )[0]

def mrr(retrieved_lists, relevant_ids_list):
    """retrieved_lists: 每次查询返回的排序后 chunk_id 列表
       relevant_ids_list: 每次查询对应的真正相关 chunk_id 集合"""
    reciprocal_ranks = []
    for retrieved, relevant_ids in zip(retrieved_lists, relevant_ids_list):
        rank = next(
            (i + 1 for i, cid in enumerate(retrieved) if cid in relevant_ids),
            None
        )
        reciprocal_ranks.append(1 / rank if rank else 0)
    return sum(reciprocal_ranks) / len(reciprocal_ranks)

import math

def recall_at_k(retrieved, relevant_ids, k):
    """Top-k 里命中的相关片段数 / 全部相关片段数"""
    top_k = retrieved[:k]
    hit = sum(1 for cid in top_k if cid in relevant_ids)
    return hit / len(relevant_ids) if relevant_ids else 0.0


def ndcg_at_k(retrieved, relevance_map, k):
    """
    relevance_map: {chunk_id: 相关性等级}，如 2=高度相关, 1=部分相关, 0=不相关
    DCG 对靠前位置的命中给更高权重；IDCG 是理想排序下的 DCG，用于归一化。
    """
    def dcg(scores):
        return sum(rel / math.log2(i + 2) for i, rel in enumerate(scores))

    gains = [relevance_map.get(cid, 0) for cid in retrieved[:k]]
    ideal = sorted(relevance_map.values(), reverse=True)[:k]

    idcg = dcg(ideal)
    return dcg(gains) / idcg if idcg > 0 else 0.0

def faithfulness_score(answer, contexts, judge_llm):
    """
    忠实度 = 回答中能被检索资料支持的论断比例。
    answer: 模型的最终回答
    contexts: 检索到的资料片段列表
    judge_llm: 充当裁判的强模型
    """
    # 第一步：把回答拆解成一条条可独立验证的原子论断
    decompose_prompt = f"""把下面这段回答拆解成一系列独立的、可单独验证的简单陈述句，
每行一句，不要编号。

回答：{answer}"""
    claims = [c.strip() for c in judge_llm.generate(decompose_prompt).split("\n") if c.strip()]
    if not claims:
        return 1.0

    # 第二步：逐条判断该论断能否被资料支持
    context_text = "\n\n".join(contexts)
    supported = 0
    for claim in claims:
        verify_prompt = f"""资料：
{context_text}

陈述：{claim}

这条陈述能否由上面的资料直接推出或支持？只回答 YES 或 NO。"""
        verdict = judge_llm.generate(verify_prompt).strip().upper()
        if verdict.startswith("YES"):
            supported += 1

    return supported / len(claims)

def answer_relevance_score(question, answer, judge_llm, embed_fn, n=3):
    """从答案反推问题，再和原问题比相似度，衡量答案是否切题"""
    gen_prompt = f"""根据下面这个回答，反推出它最可能在回答的 {n} 个问题，每行一个：

回答：{answer}"""
    reverse_questions = [
        q.strip() for q in judge_llm.generate(gen_prompt).split("\n") if q.strip()
    ]

    q_vec = embed_fn(question)
    sims = []
    for rq in reverse_questions:
        rq_vec = embed_fn(rq)
        # 向量已归一化时，点积即余弦相似度
        sims.append(sum(a * b for a, b in zip(q_vec, rq_vec)))
    return sum(sims) / len(sims) if sims else 0.0

def generate_eval_dataset(chunks, llm, questions_per_chunk=1):
    """从知识库 chunk 自动生成 (问题, 黄金片段) 评估对"""
    dataset = []
    for chunk in chunks:
        prompt = f"""根据下面这段资料，提出 {questions_per_chunk} 个用户可能会问、
且答案完全包含在这段资料里的问题，每行一个，不要编号：

资料：{chunk['text']}"""
        questions = [q.strip() for q in llm.generate(prompt).split("\n") if q.strip()]
        for q in questions:
            dataset.append({"question": q, "gold_chunk_id": chunk["id"]})
    return dataset

def agentic_rag(question, llm, retriever, max_iterations=4):
    """让模型自主决定检索动作，直到它认为信息足够为止"""
    gathered_context = []
    for _ in range(max_iterations):
        decision_prompt = f"""问题：{question}

已收集到的资料：
{gathered_context if gathered_context else "（暂无）"}

请判断：现有资料是否足以回答这个问题？
如果足够，回复：ANSWER
如果不够，回复：SEARCH: <你想检索的具体查询词>"""
        decision = llm.generate(decision_prompt)

        if decision.startswith("SEARCH:"):
            sub_query = decision.replace("SEARCH:", "").strip()
            new_chunks = retriever.search(sub_query, top_k=5)
            gathered_context.extend(new_chunks)
        else:
            break

    final_prompt = f"资料：\n{gathered_context}\n\n问题：{question}\n请给出最终回答："
    return llm.generate(final_prompt)

# 以 Claude 的工具调用为例声明一个检索工具
tools = [{
    "name": "search_knowledge_base",
    "description": "在企业知识库中检索与查询相关的文档片段，返回最相关的若干段落。",
    "input_schema": {
        "type": "object",
        "properties": {
            "query": {"type": "string", "description": "检索查询词，应聚焦、具体"},
            "top_k": {"type": "integer", "description": "返回片段数量，默认 5"},
        },
        "required": ["query"],
    },
}]

def agentic_rag_with_tools(question, client, retriever, max_turns=5):
    messages = [{"role": "user", "content": question}]
    for _ in range(max_turns):
        resp = client.messages.create(
            model="claude-sonnet-4-5",
            max_tokens=1024,
            tools=tools,
            messages=messages,
        )
        messages.append({"role": "assistant", "content": resp.content})

        # 模型没有再调用工具，说明它认为信息够了，直接产出答案
        tool_uses = [b for b in resp.content if getattr(b, "type", None) == "tool_use"]
        if not tool_uses:
            return resp

        # 执行模型请求的每一次检索，把结果作为 tool_result 回灌
        tool_results = []
        for tu in tool_uses:
            chunks = retriever.search(tu.input["query"], top_k=tu.input.get("top_k", 5))
            tool_results.append({
                "type": "tool_result",
                "tool_use_id": tu.id,
                "content": "\n\n".join(chunks),
            })
        messages.append({"role": "user", "content": tool_results})
    return resp

def multi_hop_retrieve(question, llm, retriever):
    # 第一跳：检索能回答"子问题1"的资料
    sub_q1 = llm.generate(f"要回答「{question}」，第一步需要先查清楚什么？只给出子问题。")
    hop1_chunks = retriever.search(sub_q1, top_k=3)

    # 从第一跳结果里提取出关键实体/线索，构造第二跳的查询
    bridge_entity = llm.generate(f"根据资料：{hop1_chunks}\n\n提取出回答「{sub_q1}」的关键实体。")
    sub_q2 = llm.generate(f"结合「{bridge_entity}」，构造用于回答「{question}」的下一步检索查询。")
    hop2_chunks = retriever.search(sub_q2, top_k=3)

    all_context = hop1_chunks + hop2_chunks
    return llm.generate(f"资料：{all_context}\n\n问题：{question}\n回答：")

def extract_graph_from_document(llm, document):
    """用 LLM 从文档中抽取 (实体, 关系, 实体) 三元组，构建知识图谱的原料"""
    prompt = f"""从下面的文本中抽取实体和它们之间的关系，
以三元组列表输出，每行格式为：实体1 | 关系 | 实体2

文本：{document}"""
    triples = []
    for line in llm.generate(prompt).split("\n"):
        parts = [p.strip() for p in line.split("|")]
        if len(parts) == 3:
            triples.append(tuple(parts))
    return triples


def graph_rag_query(question, graph_db, vector_store, llm, hops=2):
    """结合图谱关系遍历与向量检索的混合查询"""
    # 第一步：从问题中识别出起始实体
    entities = llm.generate(f"从问题中提取出涉及的关键实体，逗号分隔：{question}")
    seed_entities = [e.strip() for e in entities.split(",") if e.strip()]

    # 第二步：沿图谱的边做多跳遍历，收集关联子图
    subgraph = graph_db.traverse(seed_entities, max_hops=hops)

    # 第三步：把结构化的关系路径 + 向量检索到的文本证据一起交给模型
    text_evidence = vector_store.search(question, top_k=5)
    context = f"实体关系：\n{subgraph}\n\n文本证据：\n{text_evidence}"
    return llm.generate(f"资料：{context}\n\n问题：{question}\n回答：")

def answer_with_citations(question, retrieved_chunks, llm):
    """让模型在回答里标注引用编号，并把编号映射回来源元数据"""
    # 给每个片段编号，同时保留其来源信息
    numbered = []
    source_map = {}
    for i, chunk in enumerate(retrieved_chunks, start=1):
        numbered.append(f"[{i}] {chunk['text']}")
        source_map[i] = {
            "doc_title": chunk["metadata"]["doc_title"],
            "url": chunk["metadata"].get("url"),
            "page": chunk["metadata"].get("page"),
        }

    context = "\n\n".join(numbered)
    prompt = f"""请仅根据下面带编号的资料回答问题。
在每个论断后面用方括号标注它依据的资料编号，如 [1] 或 [2][3]。
如果资料中没有相关信息，请明确说不知道，不要编造。

资料：
{context}

问题：{question}

回答（记得标注引用编号）："""

    answer = llm.generate(prompt)
    return {"answer": answer, "sources": source_map}

class SemanticCache:
    def __init__(self, embed_fn, threshold=0.95):
        self.embed_fn = embed_fn
        self.threshold = threshold          # 相似度阈值，越高越"严格"，越不容易误命中
        self.cache = []                     # [(向量, 问题, 答案)]

    def get(self, question):
        q_vec = self.embed_fn(question)
        best_sim, best_answer = 0.0, None
        for vec, _, answer in self.cache:
            sim = sum(a * b for a, b in zip(q_vec, vec))  # 归一化向量点积=余弦相似度
            if sim > best_sim:
                best_sim, best_answer = sim, answer
        # 只有相似度足够高才算命中缓存
        return best_answer if best_sim >= self.threshold else None

    def set(self, question, answer):
        self.cache.append((self.embed_fn(question), question, answer))


def cached_rag(question, cache, rag_pipeline):
    hit = cache.get(question)
    if hit is not None:
        return hit          # 缓存命中，直接返回，省下检索 + 生成的全部开销
    answer = rag_pipeline(question)
    cache.set(question, answer)
    return answer

def index_multimodal_document(page, vision_llm, embedding_model, vector_store):
    """图转文路线：为页面中的每张图生成描述后一并入库"""
    # 文本部分照常切块入库
    for chunk in split_into_chunks(page.text, chunk_size=500, overlap=50):
        vec = embedding_model.encode([chunk])[0]
        vector_store.add(chunk, vec, metadata={"type": "text", "page": page.number})

    # 图片部分：用多模态模型生成详细描述，把描述当文本索引
    for img in page.images:
        caption = vision_llm.describe(
            img,
            prompt="详细描述这张图/图表的内容，包括其中的数据、趋势、标签和关键结论。"
        )
        vec = embedding_model.encode([caption])[0]
        vector_store.add(
            caption, vec,
            metadata={"type": "image", "page": page.number, "image_ref": img.id}
        )


def multimodal_answer(question, vector_store, vision_llm, top_k=5):
    """检索命中图片时，把原图而非仅描述交给多模态模型作答"""
    hits = vector_store.search(question, top_k=top_k)
    text_context, image_refs = [], []
    for hit in hits:
        if hit["metadata"]["type"] == "image":
            image_refs.append(hit["metadata"]["image_ref"])
        text_context.append(hit["text"])
    return vision_llm.generate(
        prompt=f"资料：\n{text_context}\n\n问题：{question}",
        images=[load_image(ref) for ref in image_refs],
    )

import time

def instrumented_rag(question, retriever, reranker, llm, logger):
    """带全链路埋点的 RAG，用于生产监控"""
    trace = {"question": question, "ts": time.time()}

    t0 = time.time()
    candidates = retriever.search(question, top_k=50)
    trace["retrieval_ms"] = (time.time() - t0) * 1000
    trace["candidate_count"] = len(candidates)

    t1 = time.time()
    reranked = reranker(question, candidates, final_k=5)
    trace["rerank_ms"] = (time.time() - t1) * 1000
    trace["top_rerank_score"] = reranked[0][1] if reranked else 0.0
    # 重排最高分过低，标记为疑似"知识库无相关内容"
    trace["low_confidence"] = trace["top_rerank_score"] < 0.3

    t2 = time.time()
    context = [text for text, _ in reranked]
    answer = llm.generate(f"资料：\n{context}\n\n问题：{question}\n回答：")
    trace["generation_ms"] = (time.time() - t2) * 1000
    trace["total_ms"] = (time.time() - trace["ts"] * 1000)

    logger.emit(trace)   # 打到监控系统（Prometheus / 日志平台 / 数据仓库）
    return answer, trace

import hashlib

def incremental_update(new_docs, vector_store, embedding_model, doc_hashes):
    """基于内容哈希的增量更新：只处理真正变化的文档"""
    for doc in new_docs:
        content_hash = hashlib.sha256(doc.text.encode("utf-8")).hexdigest()
        if doc_hashes.get(doc.id) == content_hash:
            continue  # 内容没变，跳过

        # 内容变了：先删掉该文档的旧 chunk，再重新切块入库
        vector_store.delete(filter={"doc_id": doc.id})
        chunks = recursive_split(doc.text, chunk_size=500, overlap=50)
        vectors = embedding_model.encode(chunks)
        vector_store.add(
            chunks, vectors,
            metadata=[{"doc_id": doc.id, "updated_at": doc.updated_at} for _ in chunks]
        )
        doc_hashes[doc.id] = content_hash

from rank_bm25 import BM25Okapi

def build_index(docs):
    # 1) 递归切分：512 token 分块，重叠 64
    chunks = []
    for d in docs:
        chunks += recursive_split(d, size=512, overlap=64)
    # 2) 向量化 + BM25 双索引
    vectors = embed([c.text for c in chunks])
    vector_store.add(chunks, vectors)
    bm25 = BM25Okapi([tokenize(c.text) for c in chunks])
    return chunks, bm25

def retrieve(query, chunks, bm25, k=20):
    # 3) 混合检索：向量 topK + BM25 topK
    dense = vector_store.search(embed([query])[0], k=k)
    sparse_scores = bm25.get_scores(tokenize(query))
    sparse = top_k_by_score(chunks, sparse_scores, k=k)
    # 4) RRF 融合
    fused = reciprocal_rank_fusion([dense, sparse], k=60)
    return fused[:k]

def answer(query, chunks, bm25):
    candidates = retrieve(query, chunks, bm25)
    # 5) Cross-Encoder 重排，取前 5
    reranked = cross_encoder_rerank(query, candidates)[:5]
    # 6) 置信度阈值：都太低就坦白不知道
    if reranked[0].score < 0.3:
        return "抱歉，我在现有资料里没找到相关信息。"
    # 7) 带引用生成
    context = "\n\n".join(
        f"[{i+1}] {c.text}" for i, c in enumerate(reranked)
    )
    prompt = (
        "只依据下面资料回答，并在句末用 [编号] 标注来源；"
        "资料没提到就说不知道。\n\n" + context + "\n\n问题：" + query
    )
    return call_llm(prompt)

def recall_at_k(retrieved_ids, gold_ids, k=10):
    hit = len(set(retrieved_ids[:k]) & set(gold_ids))
    return hit / max(len(gold_ids), 1)

def faithfulness(answer, context, judge_llm):
    # 用 LLM 裁判：答案里每句是否都能从 context 找到支撑
    prompt = (
        "判断答案是否完全由上下文支撑，只回 0-1 的小数。\n"
        f"上下文：{context}\n答案：{answer}"
    )
    return float(judge_llm(prompt))

# 汇总一个评估集
scores = {"recall": [], "faith": []}
for case in eval_set:
    r = retrieve(case["q"])
    a = answer(case["q"])
    scores["recall"].append(recall_at_k([d.id for d in r], case["gold"]))
    scores["faith"].append(faithfulness(a, r, judge_llm))

print("Recall@10:", sum(scores["recall"]) / len(scores["recall"]))
print("Faithfulness:", sum(scores["faith"]) / len(scores["faith"]))