runtime.py

from __future__ import annotations

from pathlib import Path
from threading import Lock
from typing import Optional
import time

import requests
from huggingface_hub import hf_hub_download

from config import (
    CACHE_DIR,
    FAISS_LOCAL,
    FAISS_URL,
    IDLE_UNLOAD_SECONDS,
    MANIFEST_LOCAL,
    MANIFEST_URL,
    RERANKER_FILENAME,
    RERANKER_REPO_ID,
    VECTORS_LOCAL,
    VECTORS_URL,
)
from index.store import IndexStore


def _download_to(url: str, dst: Path) -> Path:
    dst.parent.mkdir(parents=True, exist_ok=True)
    if dst.exists() and dst.stat().st_size > 0:
        return dst

    with requests.get(url, stream=True, timeout=600) as r:
        r.raise_for_status()
        with open(dst, "wb") as f:
            for chunk in r.iter_content(chunk_size=1024 * 1024):
                if chunk:
                    f.write(chunk)
    return dst


class RuntimeState:
    def __init__(self) -> None:
        self.store: Optional[IndexStore] = None
        self.reranker_model_path: Optional[Path] = None
        self.loading = False
        self.last_used_ts = 0.0
        self._lock = Lock()

    def ensure_assets(self) -> None:
        _download_to(FAISS_URL, FAISS_LOCAL)
        _download_to(MANIFEST_URL, MANIFEST_LOCAL)
        _download_to(VECTORS_URL, VECTORS_LOCAL)

        local_path = hf_hub_download(
            repo_id=RERANKER_REPO_ID,
            filename=RERANKER_FILENAME,
        )
        self.reranker_model_path = Path(local_path)

    def maybe_unload(self) -> None:
        if IDLE_UNLOAD_SECONDS <= 0:
            return

        now = time.time()
        if self.store is not None and (now - self.last_used_ts) > IDLE_UNLOAD_SECONDS:
            self.store = None

    def ensure_loaded(self) -> None:
        self.maybe_unload()

        if self.store is not None:
            self.last_used_ts = time.time()
            return

        with self._lock:
            if self.store is not None:
                self.last_used_ts = time.time()
                return

            self.loading = True
            try:
                self.ensure_assets()
                self.store = IndexStore(root=CACHE_DIR).load_all()
                self.last_used_ts = time.time()
            finally:
                self.loading = False

    def status_text(self) -> str:
        if self.loading:
            return "Loading assets..."
        if self.store is None:
            return "Cold state"
        return "Ready"


RUNTIME = RuntimeState()