backend/hf_sync.py

"""Hugging Face Dataset acts as the persistent store for teacher rubric files.

Layout inside the dataset repo:
    rubrics/<filename>           <- raw rubric files
    vector_cache/manifest.json   <- vector cache metadata
    vector_cache/index/**        <- persisted LanceDB files

Local disk on HF Spaces is ephemeral, so we try to restore the vector cache
from the dataset first and only rebuild from rubric files when the cache is stale.
"""
from __future__ import annotations

import json
import logging
from pathlib import Path
import shutil
from tempfile import TemporaryDirectory
from typing import List

from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError

from config import DATASET_ID, HF_TOKEN

logger = logging.getLogger(__name__)

RUBRIC_PREFIX = "rubrics/"
VECTOR_PREFIX = "vector_cache/"
VECTOR_INDEX_PREFIX = f"{VECTOR_PREFIX}index/"
VECTOR_MANIFEST_PATH = f"{VECTOR_PREFIX}manifest.json"
VECTOR_CACHE_SCHEMA_VERSION = 1


def is_configured() -> bool:
    return bool(HF_TOKEN and DATASET_ID)


def _vector_manifest(rubric_filenames: List[str], embed_model: str) -> dict:
    return {
        "schema_version": VECTOR_CACHE_SCHEMA_VERSION,
        "embed_model": embed_model,
        "rubric_files": sorted(rubric_filenames),
    }


def _api() -> HfApi:
    if not HF_TOKEN or not DATASET_ID:
        raise RuntimeError(
            "HF_TOKEN and DATASET_ID must be set as environment variables / Space secrets."
        )
    return HfApi(token=HF_TOKEN)


def _ensure_dataset_exists(api: HfApi) -> None:
    try:
        api.repo_info(repo_id=DATASET_ID, repo_type="dataset")
    except RepositoryNotFoundError:
        try:
            api.create_repo(
                repo_id=DATASET_ID,
                repo_type="dataset",
                private=True,
                exist_ok=True,
            )
            logger.info("Created dataset repo %s", DATASET_ID)
        except Exception as exc:  # noqa: BLE001
            raise RuntimeError(
                f"Dataset {DATASET_ID!r} does not exist and could not be created automatically. "
                "Create it on Hugging Face or update DATASET_ID / HF_TOKEN permissions."
            ) from exc


def list_remote_rubrics() -> List[str]:
    api = _api()
    try:
        files = api.list_repo_files(repo_id=DATASET_ID, repo_type="dataset")
    except RepositoryNotFoundError:
        return []
    return [
        f[len(RUBRIC_PREFIX):]
        for f in files
        if f.startswith(RUBRIC_PREFIX) and not f.endswith("/")
    ]


def download_vector_store(
    target_dir: Path,
    expected_rubric_filenames: List[str],
    embed_model: str,
) -> bool:
    if not is_configured():
        return False

    api = _api()
    try:
        manifest_path = hf_hub_download(
            repo_id=DATASET_ID,
            repo_type="dataset",
            filename=VECTOR_MANIFEST_PATH,
            token=HF_TOKEN,
        )
    except (EntryNotFoundError, RepositoryNotFoundError):
        logger.info("No vector cache manifest in dataset yet")
        return False
    except Exception as exc:  # noqa: BLE001
        logger.warning("Could not download vector cache manifest: %s", exc)
        return False

    try:
        manifest = json.loads(Path(manifest_path).read_text("utf-8"))
    except Exception as exc:  # noqa: BLE001
        logger.warning("Could not parse vector cache manifest: %s", exc)
        return False

    expected_manifest = _vector_manifest(expected_rubric_filenames, embed_model)
    if manifest != expected_manifest:
        logger.info("Vector cache manifest is stale; rebuilding from rubric files")
        return False

    try:
        files = api.list_repo_files(repo_id=DATASET_ID, repo_type="dataset")
    except RepositoryNotFoundError:
        return False

    vector_files = [f for f in files if f.startswith(VECTOR_INDEX_PREFIX)]
    if not vector_files:
        logger.info("Vector cache manifest exists but no cache files were found")
        return False

    shutil.rmtree(target_dir, ignore_errors=True)
    target_dir.mkdir(parents=True, exist_ok=True)

    for repo_path in vector_files:
        relative_path = repo_path[len(VECTOR_INDEX_PREFIX):]
        if not relative_path:
            continue
        downloaded = hf_hub_download(
            repo_id=DATASET_ID,
            repo_type="dataset",
            filename=repo_path,
            token=HF_TOKEN,
        )
        destination = target_dir / relative_path
        destination.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(downloaded, destination)

    logger.info("Vector cache restored from dataset (%d files)", len(vector_files))
    return True


def delete_vector_store() -> bool:
    """Remove the persisted vector cache folder from the dataset."""
    if not is_configured():
        return False

    api = _api()
    try:
        api.delete_folder(
            path_in_repo=VECTOR_PREFIX.rstrip("/"),
            repo_id=DATASET_ID,
            repo_type="dataset",
            commit_message="Delete vector cache",
        )
        logger.info("Vector cache deleted from HF Dataset")
        return True
    except (EntryNotFoundError, RepositoryNotFoundError):
        return False


def upload_vector_store(local_dir: Path, rubric_filenames: List[str], embed_model: str) -> bool:
    if not is_configured():
        return False

    api = _api()
    _ensure_dataset_exists(api)

    local_files = [p for p in local_dir.rglob("*") if p.is_file()]
    if not rubric_filenames or not local_files:
        delete_vector_store()
        return bool(not rubric_filenames)

    with TemporaryDirectory(prefix="feedback-vector-cache-") as tmp_dir:
        stage_dir = Path(tmp_dir)
        index_dir = stage_dir / "index"
        shutil.copytree(local_dir, index_dir, dirs_exist_ok=True)
        (stage_dir / "manifest.json").write_text(
            json.dumps(_vector_manifest(rubric_filenames, embed_model), indent=2, ensure_ascii=False),
            "utf-8",
        )

        api.upload_folder(
            repo_id=DATASET_ID,
            repo_type="dataset",
            folder_path=stage_dir,
            path_in_repo=VECTOR_PREFIX.rstrip("/"),
            commit_message="Sync vector cache",
            delete_patterns="**",
            ignore_patterns=["**/.DS_Store"],
        )

    logger.info("Vector cache synced to HF Dataset")
    return True


def download_all_rubrics(target_dir: Path) -> List[Path]:
    target_dir.mkdir(parents=True, exist_ok=True)
    paths: List[Path] = []
    for name in list_remote_rubrics():
        try:
            local = hf_hub_download(
                repo_id=DATASET_ID,
                repo_type="dataset",
                filename=f"{RUBRIC_PREFIX}{name}",
                local_dir=str(target_dir),
                token=HF_TOKEN,
            )
            paths.append(Path(local))
        except EntryNotFoundError:
            continue
    return paths


def upload_rubric(local_path: Path, filename: str) -> None:
    api = _api()
    try:
        _ensure_dataset_exists(api)
        api.upload_file(
            path_or_fileobj=str(local_path),
            path_in_repo=f"{RUBRIC_PREFIX}{filename}",
            repo_id=DATASET_ID,
            repo_type="dataset",
            commit_message=f"Add {filename}",
        )
    except RuntimeError:
        raise
    except Exception as exc:  # noqa: BLE001
        raise RuntimeError(
            f"Failed to upload {filename} to dataset {DATASET_ID!r}: {exc}"
        ) from exc


def delete_rubric(filename: str) -> bool:
    api = _api()
    try:
        api.delete_file(
            path_in_repo=f"{RUBRIC_PREFIX}{filename}",
            repo_id=DATASET_ID,
            repo_type="dataset",
            commit_message=f"Delete {filename}",
        )
        return True
    except EntryNotFoundError:
        return False