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 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]]

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 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 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)

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)

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回答：")