dp-zp-agent/tools/rag/evaluation/metrics.py

from __future__ import annotations

from collections import defaultdict
from typing import Any

EVALUATION_K_VALUES = (
    1,
    3,
    5,
)


def unique_document_ranking(
    results: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    selected: list[dict[str, Any]] = []

    seen: set[str] = set()

    for item in results:
        document_path = str(item["document_path"])

        if document_path in seen:
            continue

        seen.add(document_path)

        selected.append(item)

    return selected


def first_relevant_rank(
    ranked_documents: list[dict[str, Any]],
    expected_documents: set[str],
) -> int | None:
    for rank, item in enumerate(
        ranked_documents,
        start=1,
    ):
        if item["document_path"] in expected_documents:
            return rank

    return None


def recall_at_k(
    ranked_documents: list[dict[str, Any]],
    expected_documents: set[str],
    k: int,
) -> float:
    if not expected_documents:
        return 0.0

    retrieved = {str(item["document_path"]) for item in ranked_documents[:k]}

    relevant_retrieved = retrieved & expected_documents

    return len(relevant_retrieved) / len(expected_documents)


def evaluate_question_mode(
    question: dict[str, Any],
    mode: str,
    results: list[dict[str, Any]],
) -> dict[str, Any]:
    ranked_documents = unique_document_ranking(results)

    expected_documents = {str(value) for value in question["expected_documents"]}

    rank = first_relevant_rank(
        ranked_documents,
        expected_documents,
    )

    reciprocal_rank = 0.0 if rank is None else 1.0 / rank

    row: dict[str, Any] = {
        "id": question["id"],
        "split": question.get("split"),
        "category": question.get(
            "category",
            "unknown",
        ),
        "difficulty": question.get(
            "difficulty",
            "unknown",
        ),
        "question": question["question"],
        "mode": mode,
        "expected_documents": sorted(expected_documents),
        "first_relevant_rank": rank,
        "reciprocal_rank": round(
            reciprocal_rank,
            6,
        ),
        "top_documents": [item["document_path"] for item in ranked_documents],
        "top_source_urls": [item.get("source_url") for item in ranked_documents],
    }

    for k in EVALUATION_K_VALUES:
        row[f"hit_at_{k}"] = 1 if (rank is not None and rank <= k) else 0

        row[f"recall_at_{k}"] = round(
            recall_at_k(
                ranked_documents,
                expected_documents,
                k,
            ),
            6,
        )

    return row


def average(
    values: list[float],
) -> float:
    if not values:
        return 0.0

    return sum(values) / len(values)


def aggregate_metrics(
    rows: list[dict[str, Any]],
) -> dict[str, Any]:
    if not rows:
        return {
            "questions": 0,
            "hit_at_1": 0.0,
            "hit_at_3": 0.0,
            "hit_at_5": 0.0,
            "mrr": 0.0,
            "recall_at_5": 0.0,
        }

    return {
        "questions": len(rows),
        "hit_at_1": round(
            average([float(row["hit_at_1"]) for row in rows]),
            6,
        ),
        "hit_at_3": round(
            average([float(row["hit_at_3"]) for row in rows]),
            6,
        ),
        "hit_at_5": round(
            average([float(row["hit_at_5"]) for row in rows]),
            6,
        ),
        "mrr": round(
            average([float(row["reciprocal_rank"]) for row in rows]),
            6,
        ),
        "recall_at_5": round(
            average([float(row["recall_at_5"]) for row in rows]),
            6,
        ),
    }


def aggregate_by_category(
    rows: list[dict[str, Any]],
) -> dict[
    str,
    dict[str, Any],
]:
    grouped: dict[
        str,
        list[dict[str, Any]],
    ] = defaultdict(list)

    for row in rows:
        category = str(
            row.get(
                "category",
                "unknown",
            )
        )

        grouped[category].append(row)

    return {
        category: aggregate_metrics(category_rows)
        for (
            category,
            category_rows,
        ) in sorted(grouped.items())
    }


def aggregate_by_difficulty(
    rows: list[dict[str, Any]],
) -> dict[
    str,
    dict[str, Any],
]:
    grouped: dict[
        str,
        list[dict[str, Any]],
    ] = defaultdict(list)

    for row in rows:
        difficulty = str(
            row.get(
                "difficulty",
                "unknown",
            )
        )

        grouped[difficulty].append(row)

    return {
        difficulty: aggregate_metrics(difficulty_rows)
        for (
            difficulty,
            difficulty_rows,
        ) in sorted(grouped.items())
    }