RAG 做得好不好，八成取决于检索而不是模型

"""RAG 检索管道：混合召回 + RRF 融合 + 重排 + 带引用生成。"""

from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Any

import boto3

BEDROCK = boto3.client("bedrock-runtime", region_name="ap-northeast-1")
OPENSEARCH_INDEX = "kb-chunks"

GEN_MODEL = "anthropic.claude-3-5-sonnet-20241022-v2:0"
EMBED_MODEL = "amazon.titan-embed-text-v2:0"


@dataclass
class Chunk:
    chunk_id: str
    text: str
    source: str
    page: int
    score: float = 0.0


def embed(text: str) -> list[float]:
    """生成查询向量。"""
    resp = BEDROCK.invoke_model(
        modelId=EMBED_MODEL,
        body=json.dumps({"inputText": text, "dimensions": 1024, "normalize": True}),
    )
    return json.loads(resp["body"].read())["embedding"]


def vector_search(os_client, query: str, k: int = 30) -> list[Chunk]:
    """向量召回：语义相近但用词不同的内容也能命中。"""
    body = {
        "size": k,
        "query": {"knn": {"embedding": {"vector": embed(query), "k": k}}},
        "_source": ["chunk_id", "text", "source", "page"],
    }
    hits = os_client.search(index=OPENSEARCH_INDEX, body=body)["hits"]["hits"]
    return [Chunk(**h["_source"], score=h["_score"]) for h in hits]


def keyword_search(os_client, query: str, k: int = 30) -> list[Chunk]:
    """BM25 召回：精确术语、编号、型号这类查询靠它兜底。"""
    body = {
        "size": k,
        "query": {
            "multi_match": {
                "query": query,
                "fields": ["title^3", "text"],
                "type": "best_fields",
            }
        },
        "_source": ["chunk_id", "text", "source", "page"],
    }
    hits = os_client.search(index=OPENSEARCH_INDEX, body=body)["hits"]["hits"]
    return [Chunk(**h["_source"], score=h["_score"]) for h in hits]


def rrf_fuse(lists: list[list[Chunk]], k: int = 60) -> list[Chunk]:
    """Reciprocal Rank Fusion：不同召回通道的排名融合，无需归一化分数。"""
    scores: dict[str, float] = {}
    pool: dict[str, Chunk] = {}
    for chunks in lists:
        for rank, c in enumerate(chunks, start=1):
            scores[c.chunk_id] = scores.get(c.chunk_id, 0.0) + 1.0 / (k + rank)
            pool[c.chunk_id] = c
    ordered = sorted(scores.items(), key=lambda kv: kv[1], reverse=True)
    out = []
    for cid, s in ordered:
        c = pool[cid]
        c.score = s
        out.append(c)
    return out


def rerank(query: str, chunks: list[Chunk], top_n: int = 5) -> list[Chunk]:
    """用 Rerank 模型精排。这一步对准确率的提升通常最明显。"""
    resp = BEDROCK.invoke_model(
        modelId="cohere.rerank-v3-5:0",
        body=json.dumps({
            "query": query,
            "documents": [c.text for c in chunks],
            "top_n": top_n,
        }),
    )
    results = json.loads(resp["body"].read())["results"]
    out = []
    for r in results:
        c = chunks[r["index"]]
        c.score = r["relevance_score"]
        out.append(c)
    return out


PROMPT = """你是企业知识库助手。严格依据 <参考资料> 回答问题。

规则：
1. 只使用参考资料中的信息，不要用你的先验知识补充
2. 每个结论后面用 [来源N] 标注依据
3. 如果参考资料不足以回答，直接说「现有资料无法回答这个问题」，并说明缺什么
4. 不要复述规则，直接给答案

<参考资料>
{context}
</参考资料>

问题：{question}"""


def answer(os_client, question: str) -> dict[str, Any]:
    """完整链路：召回 → 融合 → 重排 → 生成。"""
    candidates = rrf_fuse([
        vector_search(os_client, question),
        keyword_search(os_client, question),
    ])
    top = rerank(question, candidates[:50], top_n=5)

    context = "\n\n".join(
        f"[来源{i}] （出自 {c.source} 第 {c.page} 页）\n{c.text}"
        for i, c in enumerate(top, start=1)
    )

    resp = BEDROCK.converse(
        modelId=GEN_MODEL,
        messages=[{"role": "user", "content": [
            {"text": PROMPT.format(context=context, question=question)}
        ]}],
        inferenceConfig={"maxTokens": 1500, "temperature": 0.1, "topP": 0.9},
    )

    return {
        "answer": resp["output"]["message"]["content"][0]["text"],
        "citations": [
            {"source": c.source, "page": c.page, "score": round(c.score, 4)}
            for c in top
        ],
        "usage": resp["usage"],
    }