app/main.py

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

from __future__ import annotations

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

from .catalogue import ROOT, Catalogue
from .events import EventBatch, to_row
from .memory import CONTEXTS, RAM_CLASSES
from .model_detail import DetailCache
from .parsing import PARAM_BUCKETS, QUANT_BUCKETS
from .ratelimit import RateLimiter
from .recommend import PRIORITIES, HeuristicEngine, Query, perplexity_refs, 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")
# Never log client addresses: silence uvicorn's per-request access log.
logging.getLogger("uvicorn.access").disabled = True
logging.getLogger("httpx").setLevel(logging.WARNING)  # per-request Hub fetch logs are noise

APP_VERSION = os.environ.get("EXPLORER_VERSION", "0.1.0")
STATIC = ROOT / "static"
MAX_BODY = 64 * 1024
# Who may embed the app: the Hub page, other Spaces (e.g. a static wrapper Space), and itself.
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", "24"))
PREWARM_N = int(os.environ.get("EXPLORER_PREWARM", "300"))
RELOAD_SECONDS = float(os.environ.get("EXPLORER_RELOAD_SECONDS", "900"))


class RecommendRequest(BaseModel):
    family: str | None = Field(None, max_length=40)
    size: str | None = Field(None, max_length=10)
    quant: str | None = Field(None, max_length=10)
    context: int = 8192
    ram_gb: float | None = Field(None, gt=0, le=1024)
    ram_source: Literal["confirmed", "estimated"] | None = None
    priority: Literal["balanced", "quality", "speed", "memory", "long_context"] = "balanced"
    sort: Literal["recommended", "popular", "recent", "community", "all"] = "recommended"
    llm_only: bool = True
    search: str | None = Field(None, max_length=80)
    offset: int = Field(0, ge=0, le=100_000)
    limit: int = Field(24, ge=1, le=100)


class State:
    """Everything a request needs; swapped for fakes in tests."""

    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 = HeuristicEngine()
        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 ppl_refs(self, community):
        by_id = self.catalogue.by_id
        return perplexity_refs(community, {mid: r.base_model for mid, r in by_id.items()},
                               {mid: r.quant_bucket for mid, r in by_id.items()})

    def prewarm(self):
        """Fetch config + file sizes for the most downloaded LLMs, gently, so first results are exact."""
        top = sorted((r for r in self.catalogue.records if r.is_llm), key=lambda r: -r.downloads)[:PREWARM_N]
        for r in top:
            try:
                if self.details.peek(r.id) is None:
                    self.details.get(r.id)
                    time.sleep(0.25)
            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():
            # Load at boot, again shortly after (a restart can race the previous container's
            # final flush), then regularly so shards written elsewhere or by compaction appear.
            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 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()  # don't lose buffered events on shutdown

    app = FastAPI(title="MLX 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")

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

    @app.get("/bench/mlx_explorer_bench.py", include_in_schema=False)
    def bench_script():
        return FileResponse(ROOT / "bench" / "mlx_explorer_bench.py", media_type="text/x-python")

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

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

    # ------------------------------------------------------------------ api
    @app.get("/api/meta")
    def meta(llm_only: bool = True):
        return {
            "version": APP_VERSION,
            "engine": st.engine.version,
            "families": st.catalogue.families(llm_only),
            "sizes": PARAM_BUCKETS + ["MoE"],
            "quantizations": [q for q in QUANT_BUCKETS if q != "unknown"] + ["unknown"],
            "contexts": CONTEXTS,
            "ram_classes": RAM_CLASSES,
            "priorities": PRIORITIES,
            "catalogue": st.catalogue.status(),
            "collection_enabled": st.collection_enabled,
        }

    @app.post("/api/recommend")
    def recommend(req: RecommendRequest):
        if req.context not in CONTEXTS:
            raise HTTPException(422, "context must be one of " + ",".join(map(str, CONTEXTS)))
        q = Query(req.family or None, req.size or None, req.quant or None, req.context,
                  req.ram_gb, req.ram_source, req.priority)
        matches = st.catalogue.matching(q.family, q.size, q.quant, req.llm_only, req.search)
        community = st.stats.community()
        cached = {m.id: d for m in matches if (d := st.details.peek(m.id))}
        norms = st.engine.norms(matches, q, cached)
        norms["ppl_refs"] = st.ppl_refs(community)
        prelim = sort_results(st.engine.score(matches, q, cached, community, norms), req.sort)
        window = prelim[: max(DETAIL_TOP_N, req.offset + req.limit)]
        details = st.details.get_many([s.model.id for s in window], timeout=6.0)
        rescored = st.engine.score([s.model for s in window], q, details, community, norms)
        final = sort_results(rescored, req.sort) + prelim[len(window):]
        page = final[req.offset: req.offset + req.limit]
        notices = []
        if st.catalogue.degraded:
            notices.append("Showing a cached copy of the model list; some model information isn't available.")
        missing = sum(1 for s in page if not s.detail_loaded)
        if missing:
            notices.append(f"Some model information isn't available for {missing} result(s); "
                           "their memory estimates use parameter counts only.")
        return {
            "total": len(matches),
            "offset": req.offset,
            "results": [s.to_dict() for s in page],
            "notices": notices,
            "engine": st.engine.version,
        }

    @app.get("/api/model/{model_id:path}")
    def model(model_id: str, context: int = 8192, ram_gb: float | None = None,
              priority: str = "balanced", ram_source: str | None = None):
        rec = st.catalogue.by_id.get(model_id)
        if rec is None:
            raise HTTPException(404, "model not found in the mlx-community catalogue")
        if context not in CONTEXTS:
            context = 8192
        if priority not in PRIORITIES:
            priority = "balanced"
        detail = st.details.get(model_id)
        community = st.stats.community()
        if ram_gb is not None and not (0 < ram_gb <= 1024):
            ram_gb = None
        q = Query(context=context, ram_gb=ram_gb, priority=priority,
                  ram_source="confirmed" if ram_source == "confirmed" else ("estimated" if ram_gb else None))
        sibs = st.catalogue.siblings(rec)
        sib_details = st.details.get_many([s.id for s in sibs], timeout=4.0)
        norms = st.engine.norms([rec] + sibs, q, {model_id: detail, **sib_details})
        norms["ppl_refs"] = st.ppl_refs(community)
        scored = st.engine.score([rec], q, {model_id: detail}, community, norms)[0]
        by_context = []
        for c in CONTEXTS:
            mem, _ = st.engine.memory_for(rec, Query(context=c, ram_gb=ram_gb), detail)
            by_context.append(mem.to_dict())
        sib_scored = st.engine.score(sibs, q, sib_details, community, norms)
        weight_files = [f for f in detail.files if f["path"].endswith((".safetensors", ".npz", ".gguf"))]
        return {
            **scored.to_dict(),
            "files": detail.files[:60],
            "file_count": len(detail.files),
            "weight_files": len(weight_files),
            "detail_error": detail.error,
            "memory_by_context": by_context,
            "siblings": [s.to_dict() for s in sib_scored],
        }

    @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:  # also covers chunked uploads with no content-length
            return JSONResponse({"error": "payload too large"}, status_code=413)
        try:
            import json as _json
            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]
        kept = [r for r in rows if r is not None]  # legacy clickstream events from cached pages are dropped
        if kept:
            st.sink.add(kept)
            st.stats.invalidate()
        return {"accepted": len(kept), "flags": [r["suspicious_flags"] if r else ["dropped_legacy_event"] for r in rows]}

    @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