dp-zp-agent/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()
        )
    }