NotesAgentic/backend/app/benchmarks/datasets.py

"""Benchmark Dataset 注册：从受控目录加载 JSON 数据集并校验。

Dataset 只能来自配置目录（settings.benchmark_datasets_path），API 不接受调用方提交
任意文件路径。目录不存在或为空时按「无数据集」处理，不报错。
"""

from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass, field
from pathlib import Path

from pydantic import BaseModel, Field, ValidationError

from app.config import get_settings
from app.contracts import (
    BenchmarkDatasetInfo,
    BenchmarkKind,
    RAGDatasetCase,
)
from app.errors import ApiError


@dataclass
class RAGDataset:
    """内存中的 RAG 数据集：元信息 + 已校验的 Case 列表 + 内容哈希。"""

    dataset_id: str
    kind: BenchmarkKind
    version: str
    description: str
    cases: list[RAGDatasetCase] = field(default_factory=list)
    content_hash: str = ""


class _DatasetMeta(BaseModel):
    """Dataset 元数据的最小校验模型。

    list_datasets 用它逐文件校验元信息字段结构，把「合法 JSON 但字段类型错误」
    （如 cases: 42）这类损坏文件隔离掉，而不是让 len() 抛 TypeError 拖垮整个列表。
    """

    dataset_id: str = Field(min_length=1)
    kind: str = ""
    version: str = ""
    description: str = ""
    cases: list = Field(default_factory=list)


def _datasets_dir() -> Path:
    return get_settings().benchmark_datasets_path


def _dataset_files() -> list[Path]:
    directory = _datasets_dir()
    if not directory.is_dir():
        return []
    return sorted(directory.glob("*.json"))


def _content_hash(raw: bytes) -> str:
    return "sha256:" + hashlib.sha256(raw).hexdigest()


def _read_json(path: Path) -> tuple[dict, bytes]:
    """读取并解析 JSON 文件，返回 (dict, 原始字节)；非法 JSON 抛 BENCHMARK_DATASET_INVALID。"""
    try:
        raw_bytes = path.read_bytes()
        return json.loads(raw_bytes.decode("utf-8")), raw_bytes
    except (json.JSONDecodeError, OSError, UnicodeDecodeError) as exc:
        raise ApiError(
            422,
            "BENCHMARK_DATASET_INVALID",
            f"Dataset file is not valid JSON: {path.name}",
            {"path": str(path)},
        ) from exc


def _dataset_from_raw(raw: dict, raw_bytes: bytes, kind: BenchmarkKind) -> RAGDataset:
    """把单个数据集 JSON 解析为 RAGDataset，非法结构抛 BENCHMARK_DATASET_INVALID。"""
    dataset_id = raw.get("dataset_id")
    if not isinstance(dataset_id, str) or not dataset_id:
        raise ApiError(
            422,
            "BENCHMARK_DATASET_INVALID",
            "Dataset must declare a non-empty string 'dataset_id'.",
            {},
        )
    file_kind = raw.get("kind", kind.value)
    if file_kind != kind.value:
        raise ApiError(
            422,
            "BENCHMARK_DATASET_INVALID",
            f"Dataset kind mismatch: expected '{kind.value}', got '{file_kind}'.",
            {"dataset_id": dataset_id},
        )
    raw_cases = raw.get("cases")
    if not isinstance(raw_cases, list) or not raw_cases:
        raise ApiError(
            422,
            "BENCHMARK_DATASET_INVALID",
            "Dataset 'cases' must be a non-empty list.",
            {"dataset_id": dataset_id},
        )

    cases: list[RAGDatasetCase] = []
    for index, case in enumerate(raw_cases):
        try:
            parsed = RAGDatasetCase.model_validate(case)
        except ValidationError as exc:
            raise ApiError(
                422,
                "BENCHMARK_DATASET_INVALID",
                f"Dataset case #{index} is invalid.",
                {"dataset_id": dataset_id, "case_index": index, "errors": exc.errors()},
            ) from exc
        # 每个 Case 至少要声明一个期望 id，否则无法计算命中/召回
        if not parsed.expected_note_ids and not parsed.expected_block_ids:
            raise ApiError(
                422,
                "BENCHMARK_DATASET_INVALID",
                f"Dataset case '{parsed.case_id}' must declare expected_note_ids or expected_block_ids.",
                {"dataset_id": dataset_id, "case_id": parsed.case_id},
            )
        # citation_required=true 时必须声明 expected_block_ids，否则无法计算 Citation Hit Rate
        if parsed.citation_required and not parsed.expected_block_ids:
            raise ApiError(
                422,
                "BENCHMARK_DATASET_INVALID",
                f"Dataset case '{parsed.case_id}' requires expected_block_ids when citation_required is true.",
                {"dataset_id": dataset_id, "case_id": parsed.case_id},
            )
        cases.append(parsed)

    return RAGDataset(
        dataset_id=dataset_id,
        kind=kind,
        version=str(raw.get("version", "")),
        description=str(raw.get("description", "")),
        cases=cases,
        content_hash=_content_hash(raw_bytes),
    )


def list_datasets(kind: BenchmarkKind) -> list[BenchmarkDatasetInfo]:
    """枚举受控目录下指定 kind 的数据集元信息（不含 Case 内容）。

    逐文件用 _DatasetMeta 校验元信息字段结构，单个损坏文件隔离跳过而非整体失败，
    保证列表接口健壮；损坏细节由 load_dataset 抛出。
    """
    infos: list[BenchmarkDatasetInfo] = []
    for path in _dataset_files():
        try:
            raw, raw_bytes = _read_json(path)
            meta = _DatasetMeta.model_validate(raw)
        except (ApiError, ValidationError):
            continue
        if meta.kind not in ("", kind.value):
            continue
        infos.append(
            BenchmarkDatasetInfo(
                dataset_id=meta.dataset_id,
                kind=kind,
                version=meta.version,
                description=meta.description,
                case_count=len(meta.cases),
                content_hash=_content_hash(raw_bytes),
            )
        )
    return infos


def load_dataset(dataset_id: str, kind: BenchmarkKind) -> RAGDataset:
    """按 id 加载并校验数据集；找不到抛 BENCHMARK_DATASET_NOT_FOUND。"""
    for path in _dataset_files():
        raw, raw_bytes = _read_json(path)
        if raw.get("dataset_id") != dataset_id:
            continue
        return _dataset_from_raw(raw, raw_bytes, kind)
    raise ApiError(
        404,
        "BENCHMARK_DATASET_NOT_FOUND",
        f"Benchmark dataset does not exist: {dataset_id}",
        {"dataset_id": dataset_id, "kind": kind.value},
    )