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

from __future__ import annotations

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

from tools.rag.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()