training_coach/storage.py

import json
import logging
import os
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Protocol

from training_coach.models import CompletedSession


DEFAULT_HISTORY_PATH = Path("data/completed_sessions.json")
DEFAULT_HF_HISTORY_FILE = "completed_sessions.json"
logger = logging.getLogger(__name__)


class HistoryStore(Protocol):
    def load_completed_sessions(self) -> list[CompletedSession]:
        pass

    def save_completed_sessions(self, sessions: list[CompletedSession]) -> None:
        pass

    def append_completed_session(self, session: CompletedSession) -> None:
        pass


def _history_payload(sessions: list[CompletedSession]) -> dict:
    return {
        "completed_sessions": [
            session.model_dump(mode="json") for session in sessions
        ]
    }


def _sessions_from_payload(data: dict) -> list[CompletedSession]:
    sessions = data.get("completed_sessions", [])
    return [CompletedSession.model_validate(session) for session in sessions]


def _is_missing_hf_history_error(error: Exception) -> bool:
    if error.__class__.__name__ in {"EntryNotFoundError", "RepositoryNotFoundError"}:
        return True
    if error.__class__.__name__ != "HfHubHTTPError":
        return False
    response = getattr(error, "response", None)
    return getattr(response, "status_code", None) == 404


class JsonHistoryStore:
    def __init__(self, path: str | Path = DEFAULT_HISTORY_PATH):
        self.path = Path(path)

    def load_completed_sessions(self) -> list[CompletedSession]:
        if not self.path.exists():
            logger.info(
                "event=history_load_complete backend=json path=%s sessions=0 missing_file=true",
                self.path,
            )
            return []

        logger.info("event=history_load_start backend=json path=%s", self.path)
        data = json.loads(self.path.read_text(encoding="utf-8"))
        sessions = _sessions_from_payload(data)
        logger.info(
            "event=history_load_complete backend=json path=%s sessions=%s",
            self.path,
            len(sessions),
        )
        return sessions

    def save_completed_sessions(self, sessions: list[CompletedSession]) -> None:
        logger.info(
            "event=history_save_start backend=json path=%s sessions=%s",
            self.path,
            len(sessions),
        )
        self.path.parent.mkdir(parents=True, exist_ok=True)
        data = _history_payload(sessions)
        self.path.write_text(
            json.dumps(data, indent=2) + "\n",
            encoding="utf-8",
        )
        logger.info(
            "event=history_save_complete backend=json path=%s sessions=%s",
            self.path,
            len(sessions),
        )

    def append_completed_session(self, session: CompletedSession) -> None:
        logger.info(
            "event=history_append_start backend=json day_number=%s sets=%s",
            session.day_number,
            len(session.completed_sets),
        )
        sessions = self.load_completed_sessions()
        sessions.append(session)
        self.save_completed_sessions(sessions)
        logger.info(
            "event=history_append_complete backend=json sessions=%s",
            len(sessions),
        )


class HuggingFaceDatasetHistoryStore:
    def __init__(
        self,
        repo_id: str,
        token: str | None = None,
        path_in_repo: str = DEFAULT_HF_HISTORY_FILE,
        private: bool = True,
        api=None,
        downloader=None,
    ):
        self.repo_id = repo_id
        self.token = token
        self.path_in_repo = path_in_repo
        self.private = private
        self.api = api
        self.downloader = downloader

    def _api(self):
        if self.api is not None:
            return self.api
        try:
            from huggingface_hub import HfApi
        except ImportError as error:
            raise RuntimeError(
                "Install huggingface_hub to use HuggingFaceDatasetHistoryStore."
            ) from error
        self.api = HfApi()
        return self.api

    def _download(self) -> str:
        if self.downloader is not None:
            return self.downloader(
                repo_id=self.repo_id,
                filename=self.path_in_repo,
                repo_type="dataset",
                token=self.token,
            )
        try:
            from huggingface_hub import hf_hub_download
        except ImportError as error:
            raise RuntimeError(
                "Install huggingface_hub to use HuggingFaceDatasetHistoryStore."
            ) from error
        return hf_hub_download(
            repo_id=self.repo_id,
            filename=self.path_in_repo,
            repo_type="dataset",
            token=self.token,
            force_download=True,
        )

    def load_completed_sessions(self) -> list[CompletedSession]:
        logger.info(
            "event=history_load_start backend=hf_dataset repo_id=%s path_in_repo=%s",
            self.repo_id,
            self.path_in_repo,
        )
        try:
            downloaded_path = self._download()
        except Exception as error:
            if _is_missing_hf_history_error(error):
                logger.info(
                    "event=history_load_complete backend=hf_dataset repo_id=%s "
                    "path_in_repo=%s sessions=0 missing_file=true",
                    self.repo_id,
                    self.path_in_repo,
                )
                return []
            logger.exception(
                "event=history_load_failed backend=hf_dataset repo_id=%s error_type=%s",
                self.repo_id,
                type(error).__name__,
            )
            raise

        data = json.loads(Path(downloaded_path).read_text(encoding="utf-8"))
        sessions = _sessions_from_payload(data)
        logger.info(
            "event=history_load_complete backend=hf_dataset repo_id=%s "
            "path_in_repo=%s sessions=%s",
            self.repo_id,
            self.path_in_repo,
            len(sessions),
        )
        return sessions

    def save_completed_sessions(self, sessions: list[CompletedSession]) -> None:
        logger.info(
            "event=history_save_start backend=hf_dataset repo_id=%s path_in_repo=%s sessions=%s",
            self.repo_id,
            self.path_in_repo,
            len(sessions),
        )
        api = self._api()
        api.create_repo(
            repo_id=self.repo_id,
            repo_type="dataset",
            private=self.private,
            exist_ok=True,
            token=self.token,
        )
        with TemporaryDirectory() as directory:
            path = Path(directory) / Path(self.path_in_repo).name
            path.write_text(
                json.dumps(_history_payload(sessions), indent=2) + "\n",
                encoding="utf-8",
            )
            api.upload_file(
                path_or_fileobj=str(path),
                path_in_repo=self.path_in_repo,
                repo_id=self.repo_id,
                repo_type="dataset",
                token=self.token,
                commit_message="Update completed training sessions",
            )
        logger.info(
            "event=history_save_complete backend=hf_dataset repo_id=%s "
            "path_in_repo=%s sessions=%s",
            self.repo_id,
            self.path_in_repo,
            len(sessions),
        )

    def append_completed_session(self, session: CompletedSession) -> None:
        logger.info(
            "event=history_append_start backend=hf_dataset repo_id=%s day_number=%s sets=%s",
            self.repo_id,
            session.day_number,
            len(session.completed_sets),
        )
        sessions = self.load_completed_sessions()
        sessions.append(session)
        self.save_completed_sessions(sessions)
        logger.info(
            "event=history_append_complete backend=hf_dataset repo_id=%s sessions=%s",
            self.repo_id,
            len(sessions),
        )


def create_history_store() -> HistoryStore:
    backend = os.getenv("HISTORY_STORE", "").strip().lower()
    repo_id = os.getenv("HF_HISTORY_DATASET_REPO", "").strip()
    if backend == "hf_dataset" or repo_id:
        if not repo_id:
            logger.error(
                "event=history_store_config_failed backend=hf_dataset "
                "reason=missing_repo_id"
            )
            raise RuntimeError(
                "Set HF_HISTORY_DATASET_REPO when HISTORY_STORE=hf_dataset."
            )
        logger.info(
            "event=history_store_created backend=hf_dataset repo_id=%s path_in_repo=%s",
            repo_id,
            os.getenv("HF_HISTORY_FILE", DEFAULT_HF_HISTORY_FILE),
        )
        return HuggingFaceDatasetHistoryStore(
            repo_id=repo_id,
            token=os.getenv("HF_TOKEN"),
            path_in_repo=os.getenv("HF_HISTORY_FILE", DEFAULT_HF_HISTORY_FILE),
            private=os.getenv("HF_HISTORY_PRIVATE", "1") not in {"0", "false", "False"},
        )
    path = os.getenv("LOCAL_HISTORY_PATH", DEFAULT_HISTORY_PATH)
    logger.info("event=history_store_created backend=json path=%s", path)
    return JsonHistoryStore(path)