src/store/hf_sync.py

"""Best-effort Hugging Face dataset sync for the local store."""
from __future__ import annotations

import json
import logging
import os
import sqlite3
import threading
from contextlib import closing
from io import BytesIO

LOG = logging.getLogger(__name__)

_LOCK = threading.Lock()
_TIMER: threading.Timer | None = None
_DEBOUNCE_SECONDS = 3.0


def _hf_config() -> tuple[str | None, str | None]:
    repo_id = os.environ.get("HF_STORE_DATASET")
    token = os.environ.get("HF_STORE_TOKEN") or os.environ.get("HF_TOKEN")
    return repo_id, token


def pull_into_sqlite(db_path: str) -> None:
    """Pull `store.jsonl` from a private HF dataset into SQLite if configured.

    This is intentionally guarded: no dataset/token/file/package should prevent
    app startup. Rows are merged with local data through SQLite upserts/inserts.
    """
    repo_id, token = _hf_config()
    if not repo_id or not token:
        return
    try:
        from huggingface_hub import hf_hub_download

        path = hf_hub_download(
            repo_id=repo_id,
            filename="store.jsonl",
            repo_type="dataset",
            token=token,
        )
        n_sample = n_review = n_bad = 0
        with closing(sqlite3.connect(db_path)) as conn, conn, open(path, encoding="utf-8") as fh:
            for line in fh:
                line = line.strip()
                if not line:
                    continue
                # Per-line guard: one malformed/oversized row must NOT abort the
                # whole pull (a single bad line used to roll back every prior
                # insert, leaving the board showing only a stale subset).
                try:
                    row = json.loads(line)
                    kind = row.pop("kind", "")
                    if kind == "sample":
                        conn.execute(
                            """
                            INSERT OR REPLACE INTO samples (
                                ts, host, inchikey, guest_name, model, prompt_version,
                                prompt_label, batch, prompt, results_json, true_logka
                            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                            """,
                            (
                                row.get("ts", ""),
                                row.get("host", ""),
                                row.get("inchikey", ""),
                                row.get("guest_name", ""),
                                row.get("model", ""),
                                row.get("prompt_version", ""),
                                row.get("prompt_label", ""),
                                row.get("batch", ""),
                                row.get("prompt", ""),
                                row.get("results_json", "{}"),
                                row.get("true_logka"),
                            ),
                        )
                        n_sample += 1
                    elif kind == "review":
                        conn.execute(
                            """
                            INSERT OR IGNORE INTO reviews (
                                id, ts, inchikey, guest_name, model, prompt_version,
                                rating, comment, reviewer
                            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                            """,
                            (
                                row.get("id"),
                                row.get("ts", ""),
                                row.get("inchikey", ""),
                                row.get("guest_name", ""),
                                row.get("model", ""),
                                row.get("prompt_version", ""),
                                row.get("rating", ""),
                                row.get("comment", ""),
                                row.get("reviewer", ""),
                            ),
                        )
                        n_review += 1
                except Exception as line_exc:  # noqa: BLE001 - skip the bad row, keep going
                    n_bad += 1
                    if n_bad <= 3:
                        LOG.warning("HF store pull: skipped a bad row: %s", line_exc)
        LOG.warning("HF store pull: %d samples, %d reviews loaded (%d skipped)",
                    n_sample, n_review, n_bad)
    except Exception as exc:  # noqa: BLE001 - sync must never block app startup
        LOG.warning("HF store pull skipped: %s", exc)


def schedule_push(db_path: str) -> None:
    """Debounce an async best-effort upload of all samples/reviews to HF."""
    repo_id, token = _hf_config()
    if not repo_id or not token:
        return

    def _push() -> None:
        try:
            from huggingface_hub import HfApi

            buf = BytesIO()
            with closing(sqlite3.connect(db_path)) as conn:
                conn.row_factory = sqlite3.Row
                for row in conn.execute("SELECT * FROM samples ORDER BY id"):
                    obj = {"kind": "sample", **dict(row)}
                    buf.write((json.dumps(obj, ensure_ascii=False) + "\n").encode("utf-8"))
                for row in conn.execute("SELECT * FROM reviews ORDER BY id"):
                    obj = {"kind": "review", **dict(row)}
                    buf.write((json.dumps(obj, ensure_ascii=False) + "\n").encode("utf-8"))
            HfApi(token=token).upload_file(
                path_or_fileobj=buf.getvalue(),
                path_in_repo="store.jsonl",
                repo_id=repo_id,
                repo_type="dataset",
            )
        except Exception as exc:  # noqa: BLE001 - never raise into callers
            LOG.warning("HF store push failed: %s", exc)

    global _TIMER
    with _LOCK:
        if _TIMER is not None:
            _TIMER.cancel()
        _TIMER = threading.Timer(_DEBOUNCE_SECONDS, _push)
        _TIMER.daemon = True
        _TIMER.start()