app/main.py

"""DecisionLab: every LightDec (FalconDec) and Arthur model found in the models folder, plus Laya, side by side. Serves the web UI and a small JSON API on port 9910."""
from __future__ import annotations

import os
from contextlib import asynccontextmanager
import platform
import threading
import time
from pathlib import Path
from typing import Any

import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

from .assets import stamp_assets
from .demos import DEMOS, GROUPS
from .models import BACKENDS, gpu_label, load_all
from .security import Gate, SecurityMiddleware
from .validation import validate_models, validate_questions

STATIC = Path(__file__).parent / "static"
MAX_OPTIONS = int(os.getenv("MAX_OPTIONS", "40"))
MAX_BODY_BYTES = int(os.getenv("MAX_BODY_BYTES", str(256 * 1024)))
DECIDE_SLOTS = Gate(int(os.getenv("MAX_PENDING_DECIDES", "4")))      # DL-SA-004

# DL-SA-007: no /docs, /redoc or /openapi.json
VERSION = "2.3.1"
PAGE = stamp_assets((STATIC / "index.html").read_text(encoding="utf-8"), VERSION)


@asynccontextmanager
async def lifespan(_app):
    """Start loading every model in the background when the server starts (replaces the deprecated on_event)."""
    threading.Thread(target=lambda: load_all(), name="model-loader", daemon=True).start()
    yield


app = FastAPI(title="DecisionLab", version=VERSION, docs_url=None, redoc_url=None, openapi_url=None, lifespan=lifespan)
app.mount("/static", StaticFiles(directory=STATIC), name="static")
# Body limit and security headers on the lab's own routes (no token auth: operator ruling 2026-09-28)
app.add_middleware(SecurityMiddleware, max_body=MAX_BODY_BYTES)


@app.get("/", include_in_schema=False)
def index() -> HTMLResponse:
    # Asset URLs carry the release, so browsers never reuse app.js / app.css from an older release.
    return HTMLResponse(PAGE)


@app.get("/api/health")
def health() -> dict:
    return {"ok": True}


@app.get("/api/status")
def status() -> dict:
    gpu = gpu_label()
    return {
        "models": {k: b.describe() for k, b in BACKENDS.items()},
        "order": list(BACKENDS),
        "env": {"torch": torch.__version__, "python": platform.python_version(), "gpu": gpu,
                "threads": torch.get_num_threads(),
                "models_dir": os.getenv("MODELS_DIR") or "/models"},
    }


@app.get("/api/demos")
def demos(group: str | None = None) -> list[dict]:
    return [d for d in DEMOS if d["group"] == group] if group else DEMOS


@app.get("/api/groups")
def groups() -> list[dict]:
    return [dict(g, demos=sum(d["group"] == g["id"] for d in DEMOS),
                 questions=sum(len(d["questions"]) for d in DEMOS if d["group"] == g["id"])) for g in GROUPS]


@app.post("/api/reload/{key}")
def reload(key: str) -> dict:
    b = BACKENDS.get(key)
    if not b:
        raise HTTPException(404, f"unknown model '{key}'")
    b.start_load()                    # DL-SA-005: a second request while loading starts nothing
    return {"status": "loading"}


class DecideRequest(BaseModel):
    state: Any
    questions: dict[str, dict]
    models: list[str] | None = None   # every model when omitted


class Busy(Exception):
    """Too many decisions are already running."""


def run_decision(state, questions: dict, models: list[str] | None = None) -> dict:
    """One decision, run on the requested models (every model if None). Shared by /api/decide and the Gradio app.
    Raises ValueError for an invalid request and Busy when the server is at its limit."""
    validate_questions(questions, MAX_OPTIONS)
    keys = validate_models(models, list(BACKENDS))                        # DL-SA-008
    if not DECIDE_SLOTS.try_enter():
        raise Busy("DecisionLab is busy with other requests. Try again in a moment.")
    try:
        return RUN_MODELS(state, questions, keys)
    finally:
        DECIDE_SLOTS.leave()


def run_models(state, questions: dict, keys: list[str]) -> dict:
    """Run the models one after another (they never compete for the device); each model times itself."""
    out, t0 = {}, time.perf_counter()
    for key in keys:
        try:
            out[key] = BACKENDS[key].decide(state, questions)
        except Exception as exc:
            out[key] = {"error": f"{type(exc).__name__}: {exc}"[:500]}
    return {"results": out, "server_ms": round((time.perf_counter() - t0) * 1000, 1)}


# The model-running step. The Hugging Face Space replaces it with a @spaces.GPU version (ZeroGPU); the checks and the
# concurrency limit above stay in the main process.
RUN_MODELS = run_models


@app.post("/api/decide")
def decide(req: DecideRequest) -> dict:
    try:
        return run_decision(req.state, req.questions, req.models)
    except ValueError as exc:
        raise HTTPException(422, str(exc)) from None
    except Busy as exc:
        raise HTTPException(429, str(exc)) from None