dp-zp-agent/evaluation/rag_answer_data.py

from __future__ import annotations

import hashlib
import json
from pathlib import Path
from typing import Any

from evaluation.retrieval_runner import filter_questions_by_split


DEFAULT_DELAY = 0.5

ALLOWED_OVERRIDE_FIELDS = {
    "question",
    "expected_answer_contains",
    "expected_source_urls",
    "should_answer",
    "note",
}


def apply_question_overrides(
    questions: list[dict[str, Any]],
    path: Path,
) -> tuple[list[dict[str, Any]], list[str]]:
    if not path.exists():
        raise FileNotFoundError(
            f"RAG answer override súbor neexistuje: {path}"
        )

    with path.open("r", encoding="utf-8") as file:
        overrides = json.load(file)

    if not isinstance(overrides, dict):
        raise ValueError(
            "rag_answer_overrides.json musí obsahovať JSON objekt."
        )

    known_ids = {str(item["id"]) for item in questions}
    unknown_ids = set(overrides.keys()) - known_ids

    if unknown_ids:
        raise ValueError(
            "Override súbor obsahuje neznáme question ID: "
            + ", ".join(sorted(unknown_ids))
        )

    result: list[dict[str, Any]] = []
    applied_ids: list[str] = []

    for question in questions:
        item = dict(question)
        question_id = str(item["id"])
        override = overrides.get(question_id)

        if override is None:
            result.append(item)
            continue

        if not isinstance(override, dict):
            raise ValueError(
                f"{question_id}: override musí byť JSON objekt."
            )

        unsupported_fields = (
            set(override.keys())
            - ALLOWED_OVERRIDE_FIELDS
        )

        if unsupported_fields:
            raise ValueError(
                f"{question_id}: nepovolené override polia: "
                + ", ".join(
                    sorted(unsupported_fields)
                )
            )

        if "question" in override:
            overridden_question = override["question"]

            if (
                not isinstance(
                    overridden_question,
                    str,
                )
                or not overridden_question.strip()
            ):
                raise ValueError(
                    f"{question_id}: "
                    "override question musí byť "
                    "neprázdny reťazec."
                )

        item.update(override)
        result.append(item)
        applied_ids.append(question_id)

    return result, applied_ids


def validate_string_list(
    question_id: str,
    field_name: str,
    value: Any,
) -> None:
    if not isinstance(value, list):
        raise ValueError(
            f"{question_id}: "
            f"{field_name} musí byť JSON pole."
        )

    for index, item in enumerate(value):
        if not isinstance(item, str):
            raise ValueError(
                f"{question_id}: "
                f"{field_name}[{index}] "
                "musí byť reťazec."
            )


def validate_answer_questions(
    questions: list[dict[str, Any]],
) -> None:
    seen_ids: set[str] = set()

    for index, question in enumerate(
        questions,
        start=1,
    ):
        if not isinstance(question, dict):
            raise ValueError(
                f"Otázka #{index} "
                "musí byť JSON objekt."
            )

        question_id = question.get("id")

        if (
            not isinstance(question_id, str)
            or not question_id.strip()
        ):
            raise ValueError(
                f"Otázka #{index}: "
                "id musí byť neprázdny reťazec."
            )

        if question_id in seen_ids:
            raise ValueError(
                f"Duplicitné question ID: {question_id}"
            )

        seen_ids.add(question_id)

        split = question.get("split")

        if split not in {
            "dev",
            "test",
        }:
            raise ValueError(
                f"{question_id}: "
                "split musí byť "
                "'dev' alebo 'test'."
            )

        question_text = question.get(
            "question"
        )

        if (
            not isinstance(
                question_text,
                str,
            )
            or not question_text.strip()
        ):
            raise ValueError(
                f"{question_id}: "
                "question musí byť "
                "neprázdny reťazec."
            )

        validate_string_list(
            question_id,
            "expected_documents",
            question.get(
                "expected_documents"
            ),
        )

        validate_string_list(
            question_id,
            "expected_source_urls",
            question.get(
                "expected_source_urls"
            ),
        )

        validate_string_list(
            question_id,
            "expected_answer_contains",
            question.get(
                "expected_answer_contains"
            ),
        )

        if not isinstance(
            question.get(
                "should_answer"
            ),
            bool,
        ):
            raise ValueError(
                f"{question_id}: "
                "should_answer musí byť boolean."
            )

        for field_name in (
            "category",
            "difficulty",
        ):
            value = question.get(
                field_name
            )

            if (
                value is not None
                and not isinstance(
                    value,
                    str,
                )
            ):
                raise ValueError(
                    f"{question_id}: "
                    f"{field_name} musí byť "
                    "reťazec."
                )

        note = question.get("note")

        if (
            note is not None
            and not isinstance(
                note,
                str,
            )
        ):
            raise ValueError(
                f"{question_id}: "
                "note musí byť reťazec "
                "alebo null."
            )


def select_questions(
    questions: list[dict[str, Any]],
    *,
    split: str,
    question_ids: list[str],
    limit: int | None,
) -> list[dict[str, Any]]:
    selected = filter_questions_by_split(
        questions,
        split,
    )

    if question_ids:
        wanted = set(question_ids)

        selected = [
            item
            for item in selected
            if item.get("id") in wanted
        ]

    if limit is not None:
        selected = selected[:limit]

    return selected


def result_prefix(
    *,
    split: str,
    limit: int | None,
    question_ids: list[str],
) -> str:
    prefix = f"rag_answers_{split}"

    if question_ids:
        return prefix + "_selected"

    if limit is not None:
        return prefix + f"_limit{limit}"

    return prefix


def questions_fingerprint(
    questions: list[dict[str, Any]],
) -> str:
    canonical = json.dumps(
        questions,
        ensure_ascii=False,
        sort_keys=True,
        separators=(
            ",",
            ":",
        ),
        allow_nan=False,
    ).encode("utf-8")

    return hashlib.sha256(
        canonical
    ).hexdigest()