server.py

#!/usr/bin/env python3
"""Neon web UI for batch document classification. Run: uv run server.py"""

from __future__ import annotations

import os
import tempfile
from pathlib import Path

import uvicorn
from dotenv import load_dotenv
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from markitdown import MarkItDown

import game_ai

from classify import DEFAULT_OPENROUTER_MODEL, SUPPORTED_SUFFIXES, ClassificationError, classify_document, get_headers, load_catalog

load_dotenv()

BASE = Path(__file__).parent
MAX_BYTES = 25 * 1024 * 1024
ENDPOINT = os.getenv("AZURE_ENDPOINT", "")
DEPLOYMENT = os.getenv("DEPLOYMENT_NAME", "")
OPENROUTER_MODEL = os.getenv("OPENROUTER_MODEL") or DEFAULT_OPENROUTER_MODEL
CHOICES = load_catalog(BASE / "types.json")
CONVERTER = MarkItDown(enable_plugins=False)

app = FastAPI(title="Decision Classifier")
app.mount("/static", StaticFiles(directory=BASE / "static"), name="static")


@app.on_event("startup")
def check_config() -> None:
    if not ENDPOINT or not DEPLOYMENT:
        raise SystemExit("Set AZURE_ENDPOINT and DEPLOYMENT_NAME (see .env.example).")
    get_headers()  # fail fast on auth problems


@app.get("/")
def index() -> FileResponse:
    return FileResponse(BASE / "static" / "index.html")


@app.get("/game")
def game() -> FileResponse:
    return FileResponse(BASE / "static" / "game.html")


class GameState(BaseModel):
    state: str
    model: str = "decision"


@app.post("/api/game/decide")
def game_decide(body: GameState) -> dict:
    # Plain `def` so the blocking model call runs in a worker thread.
    try:
        if body.model == "jev":
            key = os.getenv("OPENROUTER_API_KEY")
            if not key:
                raise HTTPException(400, "OPENROUTER_API_KEY is not set in .env")
            return game_ai.decide_openrouter(body.state[:4000], OPENROUTER_MODEL, key)
        if body.model == "decision":
            return game_ai.decide_foundry(body.state[:4000], ENDPOINT, DEPLOYMENT)
    except ClassificationError as exc:
        raise HTTPException(502, str(exc)) from exc
    raise HTTPException(400, f"Unknown model {body.model!r}")


@app.get("/api/config")
def config() -> dict:
    models = [
        {"id": "decision", "label": "Decision-1 (Foundry)", "detail": DEPLOYMENT, "available": True},
        {"id": "jev", "label": "Jev (OpenRouter)", "detail": OPENROUTER_MODEL, "available": bool(os.getenv("OPENROUTER_API_KEY"))},
    ]
    return {"types": list(CHOICES), "models": models, "suffixes": sorted(SUPPORTED_SUFFIXES)}


@app.post("/api/classify")
def classify(file: UploadFile = File(...), model: str = Form("decision")) -> JSONResponse:
    # Plain `def`: FastAPI runs it in a worker thread, so conversion + the API call don't block other uploads.
    openrouter = None
    if model == "jev":
        key = os.getenv("OPENROUTER_API_KEY")  # loaded from .env at startup
        if not key:
            raise HTTPException(400, "OPENROUTER_API_KEY is not set in .env")
        openrouter = (OPENROUTER_MODEL, key)
    elif model != "decision":
        raise HTTPException(400, f"Unknown model {model!r}")
    name = Path(file.filename or "upload").name
    suffix = Path(name).suffix.lower()
    if suffix not in SUPPORTED_SUFFIXES:
        raise HTTPException(415, f"Unsupported file type {suffix or '(none)'}")
    data = file.file.read(MAX_BYTES + 1)
    if len(data) > MAX_BYTES:
        raise HTTPException(413, "File larger than 25 MB")
    with tempfile.TemporaryDirectory() as tmp:
        path = Path(tmp) / f"doc{suffix}"
        path.write_bytes(data)
        try:
            result = classify_document(path, CONVERTER, CHOICES, ENDPOINT, DEPLOYMENT, openrouter=openrouter)
        except ClassificationError as exc:
            raise HTTPException(502, str(exc)) from exc
        except Exception as exc:
            raise HTTPException(422, f"Could not read document: {exc}") from exc
    return JSONResponse({
        "file": name,
        "type_id": result["type_id"],
        "probability": result.get("top_probability"),
        "probabilities": result["decision"].get("probabilities"),
        "model": model,
    })


if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=int(os.getenv("PORT", "8000")))