RAG 评估：把检索、生成、引用和运行成本分开测量

import re


cases = [
    {
        "question": "关键告警通知谁？",
        "relevant": {"duty:v1:002"},
        "required_terms": {"负责人"},
    },
    {
        "question": "普通告警需要停机吗？",
        "relevant": set(),
        "required_terms": {"无法确认"},
    },
]


def retrieval_scores(ranked_ids: list[str], relevant: set[str], k: int = 3) -> dict:
    if not relevant:
        return {"recall_at_k": None, "mrr": None}
    top = ranked_ids[:k]
    recall = len(set(top) & relevant) / len(relevant)
    ranks = [index for index, doc_id in enumerate(ranked_ids, 1) if doc_id in relevant]
    return {"recall_at_k": recall, "mrr": 1 / min(ranks) if ranks else 0.0}


def answer_scores(answer: str, required_terms: set[str], source_ids: set[str]) -> dict:
    cited = set(re.findall(r"\[(S\d+)\]", answer))
    return {
        "required_terms": int(all(term in answer for term in required_terms)),
        "citation_valid": int(bool(cited) and cited <= source_ids),
    }


print(retrieval_scores(["duty:v1:002", "other:v1:000"], cases[0]["relevant"]))
print(answer_scores("关键告警应通知负责人 [S1]。", {"负责人"}, {"S1", "S2"}))