app/main.py

"""Local Model Explorer: FastAPI backend + static frontend."""

from __future__ import annotations

import dataclasses
import json
import logging
import os
import threading
import time
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Literal

from fastapi import FastAPI, HTTPException, Request
from fastapi.exceptions import RequestValidationError
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field, ValidationError, model_validator

from .catalogue import MLX_EXPLORER, ONNX_EXPLORER, ROOT, Catalogue
from .detail import BUDGET, REPOS_ON_OPEN, REPOS_PER_MODEL, DetailCache
from .events import EventBatch, to_row
from .hardware import CONTEXTS, KV_TYPES, UNIFIED, Hardware, presets
from .parsing import SIZE_BUCKETS
from .ratelimit import RateLimiter
from .recommend import PRIORITIES, SORTS, WEIGHTS, Engine, Query, options_for, sort_results
from .sink import make_sink
from .stats import StatsCache

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
log = logging.getLogger("explorer")
logging.getLogger("uvicorn.access").disabled = True  # never log client addresses
logging.getLogger("httpx").setLevel(logging.WARNING)

APP_VERSION = os.environ.get("EXPLORER_VERSION", "0.1.0")
STATIC = ROOT / "static"
MAX_BODY = 64 * 1024
FRAME_ANCESTORS = os.environ.get("EXPLORER_FRAME_ANCESTORS", "'self' https://huggingface.co https://*.hf.space")
DETAIL_TOP_N = int(os.environ.get("EXPLORER_DETAIL_TOP_N", "12"))
PREWARM_N = int(os.environ.get("EXPLORER_PREWARM", "400"))
RELOAD_SECONDS = float(os.environ.get("EXPLORER_RELOAD_SECONDS", "900"))


class HardwareIn(BaseModel):
    kind: Literal["gpu", "unified", "cpu"]
    gpus: list[float] = Field(default_factory=list, max_length=16)
    ram_gb: float = Field(32, gt=0, le=4096)
    unified_system: str | None = None

    @model_validator(mode="after")
    def _check(self):
        if any(not (0 < v <= 256) for v in self.gpus):
            raise ValueError("each GPU's VRAM must be between 0 and 256 GB")
        if self.kind == "gpu" and not self.gpus:
            raise ValueError("a GPU setup needs at least one GPU")
        if self.unified_system is not None and self.unified_system not in UNIFIED:
            raise ValueError("unknown unified-memory system")
        return self

    def to_hardware(self) -> Hardware:
        share = UNIFIED[self.unified_system][1] if self.kind == "unified" and self.unified_system else None
        return Hardware(self.kind, tuple(self.gpus) if self.kind == "gpu" else (), self.ram_gb, share)


class RecommendRequest(BaseModel):
    hardware: HardwareIn | None = None
    context: int = 8192
    kv_type: Literal["f16", "q8_0", "q4_0"] = "f16"
    priority: Literal["balanced", "quality", "small"] = "balanced"
    min_bits: Literal[0, 2, 3, 4, 5, 6, 8] = 0
    vram_only: bool = False
    mmproj: Literal["gpu", "ram", "off"] = "gpu"
    sort: Literal["best", "popular", "recent", "community"] = "best"
    family: str | None = Field(None, max_length=40)
    size: str | None = Field(None, max_length=10)
    search: str | None = Field(None, max_length=80)
    moe_only: bool = False
    vision_only: bool = False
    fits_only: bool = False
    offset: int = Field(0, ge=0, le=100_000)
    limit: int = Field(24, ge=1, le=60)


class ModelRequest(BaseModel):
    key: str = Field(..., max_length=200)
    hardware: HardwareIn | None = None
    context: int = 8192
    kv_type: Literal["f16", "q8_0", "q4_0"] = "f16"
    priority: Literal["balanced", "quality", "small"] = "balanced"
    min_bits: Literal[0, 2, 3, 4, 5, 6, 8] = 0
    vram_only: bool = False
    mmproj: Literal["gpu", "ram", "off"] = "gpu"


class State:
    def __init__(self, catalogue=None, details=None, sink=None, fetch_live=True):
        base = Path(os.environ.get("EXPLORER_DATA_DIR", ROOT / ".runtime"))
        self.catalogue = catalogue or Catalogue()
        self.details = details or DetailCache()
        self.sink = sink or make_sink(base)
        self.engine = Engine()
        self.limiter = RateLimiter(rate_per_minute=float(os.environ.get("EXPLORER_RATE_PER_MIN", "120")), burst=60)
        self.existing_rows: list[dict] = []
        self.stats = StatsCache(lambda: self.existing_rows + self.sink.live_rows())
        self.fetch_live = fetch_live
        self.collection_enabled = os.environ.get("EXPLORER_SINK", "local") != "off"

    def prewarm(self):
        """File listings (not file downloads) for the most popular models, gently."""
        top = sorted(self.catalogue.groups.values(), key=lambda g: -g.downloads)[:PREWARM_N]
        for g in top:
            try:
                if any(self.details._stale(r.id) for r in g.repos[:REPOS_PER_MODEL]):
                    self.details.get(g, timeout=120, with_geometry=False)
                    time.sleep(2.0)  # leave most of the Hub budget for people using the page
            except Exception:
                pass

    def boot(self):
        self.catalogue.load_offline()
        if self.fetch_live:
            def initial():
                self.catalogue.refresh()
                self.prewarm()
            threading.Thread(target=initial, daemon=True, name="catalogue-initial").start()
            self.catalogue.start_background_refresh()
        self.sink.start()

        def reload_loop():
            for delay in (0, 90):
                time.sleep(delay)
                self.reload_existing()
            while True:
                time.sleep(RELOAD_SECONDS)
                self.reload_existing()
        threading.Thread(target=reload_loop, daemon=True, name="reload-existing").start()

    def reload_existing(self):
        started = time.time()
        rows = self.sink.read_existing()
        if rows is None:
            return
        self.existing_rows = rows
        self.sink.forget_flushed_before(started)
        self.stats.invalidate()


def _hidden(s, details) -> bool:
    return _only_auxiliary(s) or s.group.key in details.not_llama_cpp


def _only_auxiliary(s) -> bool:
    """Every file we could read is too small to be the model (a repo of draft models or
    speculative-decoding heads, which nobody runs on their own), or below the chosen bit level."""
    repos = (s.detail_repos or {}).values()
    return s.detail_loaded and not s.options and bool(repos) and all(rf.error is None for rf in repos)


def _query(hw: HardwareIn | None, context: int, kv_type: str, priority: str, min_bits: int = 0,
           vram_only: bool = False, mmproj: str = "gpu") -> Query:
    if context not in CONTEXTS:
        raise HTTPException(422, "context must be one of " + ",".join(map(str, CONTEXTS)))
    return Query(hw.to_hardware() if hw else None, context, kv_type, priority, min_bits,
                 vram_only and bool(hw) and hw.kind == "gpu", mmproj)


def create_app(state: State | None = None) -> FastAPI:
    st = state or State()

    @asynccontextmanager
    async def lifespan(app):
        if state is None:
            st.boot()
        yield
        st.sink.flush()

    app = FastAPI(title="Local Model Explorer", version=APP_VERSION, lifespan=lifespan,
                  docs_url=None, redoc_url=None, openapi_url=None)
    app.state.st = st

    @app.middleware("http")
    async def guard(request: Request, call_next):
        cl = request.headers.get("content-length")
        if cl and cl.isdigit() and int(cl) > MAX_BODY:
            return JSONResponse({"error": "payload too large"}, status_code=413)
        resp = await call_next(request)
        resp.headers["X-Content-Type-Options"] = "nosniff"
        resp.headers["Referrer-Policy"] = "no-referrer"
        resp.headers["Content-Security-Policy"] = (
            "default-src 'self'; img-src 'self' data:; style-src 'self'; script-src 'self'; "
            f"connect-src 'self'; frame-ancestors {FRAME_ANCESTORS}")
        return resp

    @app.exception_handler(RequestValidationError)
    async def _invalid(request, exc):
        errs = [{"loc": [str(x) for x in e.get("loc", [])][-3:], "msg": str(e.get("msg"))[:160]}
                for e in exc.errors()[:10]]
        return JSONResponse({"error": "invalid request", "details": errs}, status_code=422)

    def client_key(request: Request) -> str:
        fwd = request.headers.get("x-forwarded-for", "")
        return fwd.split(",")[0].strip() or (request.client.host if request.client else "unknown")

    @app.get("/", include_in_schema=False)
    def index():
        return FileResponse(STATIC / "index.html")

    @app.get("/stats", include_in_schema=False)
    def stats_page():
        return FileResponse(STATIC / "stats.html")

    @app.get("/healthz")
    def healthz():
        cat, sink = st.catalogue.status(), st.sink.status()
        ok = cat["groups"] > 0
        return JSONResponse({"ok": ok, "version": APP_VERSION, "catalogue": cat, "sink": sink,
                             "geometries": len(st.details.geometries)}, status_code=200 if ok else 503)

    @app.get("/api/meta")
    def meta():
        return {
            "version": APP_VERSION, "engine": st.engine.version, "families": st.catalogue.families(),
            "sizes": SIZE_BUCKETS, "priorities": PRIORITIES, "sorts": SORTS, "score_weights": WEIGHTS, "hardware": presets(),
            "catalogue": st.catalogue.status(), "collection_enabled": st.collection_enabled,
            "links": {"mlx": MLX_EXPLORER, "onnx": ONNX_EXPLORER},
        }

    @app.post("/api/recommend")
    def recommend(req: RecommendRequest):
        q = _query(req.hardware, req.context, req.kv_type, req.priority, req.min_bits, req.vram_only, req.mmproj)
        groups = st.catalogue.matching(req.family or None, req.size or None, req.search, req.moe_only, req.vision_only)
        community = st.stats.community()
        cached = {g.key: d for g in groups if (d := st.details.peek(g))}
        prelim = [s for s in st.engine.score(groups, q, cached, community) if not _hidden(s, st.details)]
        if (req.fits_only or q.vram_only) and q.hardware:
            prelim = [s for s in prelim if s.best and s.best.fit and s.best.fit.tier not in ("no", "no_vram")]
        prelim = sort_results(prelim, req.sort, community)
        # Read files for everything on the page (and a little beyond it on the first page, where
        # exact sizes can reorder the top results); reads that miss the wait keep going.
        window = prelim[: max(DETAIL_TOP_N, req.limit)] if req.offset == 0 else prelim[req.offset: req.offset + req.limit]
        # Exact file sizes and geometry for what is about to be shown.
        missing = [s.group for s in window if not s.detail_loaded or s.best is None or s.best.estimated]
        if missing and st.fetch_live:
            futs = [st.details.pool.submit(st.details.get, g, 4.0) for g in missing]
            deadline = time.time() + 4.0
            for f in futs:
                try:
                    f.result(timeout=max(deadline - time.time(), 0.05))
                except Exception:
                    pass
            fresh = {g.key: d for g in missing if (d := st.details.peek(g))}
            rescored = {s.group.key: s for s in st.engine.score([s.group for s in window], q, {**cached, **fresh}, community)}
            prelim = [rescored.get(s.group.key, s) for s in prelim]
            prelim = [s for s in prelim if not _hidden(s, st.details)]
            if (req.fits_only or q.vram_only) and q.hardware:
                prelim = [s for s in prelim if s.best and s.best.fit and s.best.fit.tier not in ("no", "no_vram")]
            prelim = sort_results(prelim, req.sort, community)
        page = prelim[req.offset: req.offset + req.limit]
        notices = []
        if st.catalogue.degraded:
            notices.append("Showing a cached copy of the model list.")
        # Only models whose files are still to be read; a model whose GGUFs can't be planned stays marked (est.)
        est = sum(1 for s in page if s.best and s.best.estimated and s.group.key not in st.details.geometry_failed)
        if est:
            notices.append(f"{est} result(s) are estimated from the parameter count until their files are read.")
        return {"total": len(prelim), "offset": req.offset, "results": [s.to_dict(q) for s in page],
                "notices": notices, "engine": st.engine.version}

    @app.post("/api/model")
    def model(req: ModelRequest):
        g = st.catalogue.groups.get(req.key)
        if g is None:
            raise HTTPException(404, "model not found")
        q = _query(req.hardware, req.context, req.kv_type, req.priority, req.min_bits, req.vram_only, req.mmproj)
        detail = st.details.get(g, timeout=25.0, repos=REPOS_ON_OPEN)
        community = st.stats.community()
        scored = st.engine.score([g], q, {g.key: detail}, community)[0]
        # The recommendation first, then what runs best, then higher quality, then more downloaded uploaders.
        best = scored.best
        options = sorted(scored.options, key=lambda o: (o is not best, -(o.fit.score if o.fit else 0),
                                                        -(o.quant.bpw or 0), -o.repo_downloads, o.repo))
        repos = [{**r.to_dict(), "files_error": (detail.repos.get(r.id).error if detail.repos.get(r.id) else None),
                  "loaded": r.id in detail.repos} for r in g.repos]
        other = {fmt: ids[:30] for fmt, ids in g.other_formats.items()}
        below = (len(options_for(g, detail, dataclasses.replace(q, min_bits=0))) - len(options)) if q.min_bits else 0
        return {
            **scored.to_dict(q),
            "options": [o.to_dict(q) for o in options],
            "hidden_below_min_bits": below,
            "repos": repos,
            "geometry": detail.geometry.to_dict() if detail.geometry else None,
            "geometry_source": detail.geometry_source,
            "kv_cache_gb": ({c: round(detail.geometry.kv_bytes(c, KV_TYPES[req.kv_type]) / 1e9, 2) for c in CONTEXTS}
                            if detail.geometry else None),
            "other_formats": other,
            "community": community.get(g.key),
        }

    @app.post("/api/events", status_code=202)
    async def events(request: Request):
        if not st.limiter.allow(client_key(request)):
            return JSONResponse({"error": "rate limited"}, status_code=429)
        raw = await request.body()
        if len(raw) > MAX_BODY:
            return JSONResponse({"error": "payload too large"}, status_code=413)
        try:
            body = json.loads(raw)
        except Exception:
            return JSONResponse({"error": "invalid json"}, status_code=400)
        try:
            batch = EventBatch.model_validate(body)
        except ValidationError as e:
            errs = [{"loc": [str(x) for x in err.get("loc", [])][-3:], "msg": str(err.get("msg"))[:160]}
                    for err in e.errors()[:10]]
            return JSONResponse({"error": "invalid events", "details": errs}, status_code=422)
        if not st.collection_enabled:
            return {"accepted": 0}
        rows = [to_row(ev, st.catalogue, APP_VERSION, st.engine.version) for ev in batch.events]
        st.sink.add(rows)
        st.stats.invalidate()
        summary = [{"flags": r["suspicious_flags"], "generation_tps": r.get("generation_tps"),
                    "prompt_tps": r.get("prompt_tps"), "backend": r.get("backend")} for r in rows]
        return {"accepted": len(rows), "rows": summary}

    @app.get("/api/stats")
    def stats():
        return st.stats.stats()

    app.mount("/static", StaticFiles(directory=STATIC), name="static")
    return app


app = create_app() if os.environ.get("EXPLORER_NO_AUTOAPP") != "1" else None