__init__.py

import os

from benchmarks.adapters.longmemeval import LongMemEvalAdapter
from benchmarks.adapters.locomo import LoCoMoAdapter

# Shared dataset registry — single source of truth for the ingest and judge stages.
_DATA_ROOT = os.path.join(
    os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
    "data",
)

DATASETS = {
    "longmemeval_s": os.path.join(_DATA_ROOT, "LongMemEval", "longmemeval_s_cleaned.json"),
    "longmemeval_m": os.path.join(_DATA_ROOT, "LongMemEval", "longmemeval_m_cleaned.json"),
    "longmemeval_oracle": os.path.join(_DATA_ROOT, "LongMemEval", "longmemeval_oracle.json"),
    "locomo": os.path.join(_DATA_ROOT, "LoCoMo", "locomo10.json"),
}


def build_adapter(dataset: str):
    """Return the adapter instance for a dataset key in DATASETS."""
    path = DATASETS[dataset]
    if dataset.startswith("longmemeval"):
        return LongMemEvalAdapter(path)
    if dataset == "locomo":
        return LoCoMoAdapter(path)
    raise ValueError(f"Unknown dataset: {dataset}")


__all__ = [
    "LongMemEvalAdapter",
    "LoCoMoAdapter",
    "DATASETS",
    "build_adapter",
]