import asyncio import json from pathlib import Path from fastapi import FastAPI, File, HTTPException, UploadFile, WebSocket, WebSocketDisconnect from fastapi.responses import FileResponse, JSONResponse from fastapi.staticfiles import StaticFiles from . import db, registry, settings as settings_mod, store from .runner import runner app = FastAPI(title="MODELBEAST") WEB_DIST = db.ROOT / "web" / "dist" INBOX = db.DATA / "inbox" @app.on_event("startup") async def startup(): runner.operators = registry.load_operators() con = db.connect() app.state.con = con runner.get_settings = lambda: settings_mod.get_all(con) asyncio.create_task(runner.worker()) asyncio.create_task(watch_inbox()) # -- operators ---------------------------------------------------------------- @app.get("/api/operators") def list_operators(): return [{k: v for k, v in op.items() if k != "dir"} for op in runner.operators.values()] # -- settings ----------------------------------------------------------------- @app.get("/api/settings") def get_settings(): s = settings_mod.get_all(app.state.con) view = settings_mod.public_view(s) view["_secret_keys"] = sorted(settings_mod.SECRET_KEYS) view["_env_keys"] = sorted(settings_mod.ENV_MAP.keys()) # env var names currently satisfied — lets the UI gate operators view["_env_set"] = sorted(env for key, env in settings_mod.ENV_MAP.items() if s.get(key)) return view @app.put("/api/settings") def put_settings(payload: dict): # ignore masked secret placeholders so re-saving the form doesn't wipe a key # (an empty string is a deliberate clear, not a placeholder) clean = {k: v for k, v in payload.items() if not (k in settings_mod.SECRET_KEYS and v and set(str(v)) <= {"•"})} settings_mod.set_many(app.state.con, clean) return get_settings() # -- assets --------------------------------------------------------------------- @app.get("/api/assets") def list_assets(): return store.list_assets(app.state.con) @app.post("/api/assets") async def upload_asset(file: UploadFile = File(...)): data = await file.read() if not data: raise HTTPException(400, "empty file") return store.register_upload(app.state.con, file.filename or "unnamed", data) @app.delete("/api/assets/{asset_id}") def delete_asset(asset_id: str): if not store.delete_asset(app.state.con, asset_id): raise HTTPException(404, "no such asset") return {"ok": True} @app.get("/api/assets/{asset_id}/file") def asset_file(asset_id: str, member: str | None = None): asset = store.get_asset(app.state.con, asset_id) if not asset: raise HTTPException(404, "no such asset") path = Path(asset["path"]) if member: target = (path / member).resolve() if not str(target).startswith(str(path.resolve())) or not target.is_file(): raise HTTPException(404, "no such member") return FileResponse(target) if path.is_dir(): files = sorted(p.name for p in path.iterdir() if p.is_file()) return JSONResponse({"folder": True, "files": files}) return FileResponse(path, filename=asset["name"]) # -- jobs ----------------------------------------------------------------------- @app.get("/api/jobs") def list_jobs(): return runner.list_jobs(app.state.con) @app.post("/api/jobs") async def create_job(payload: dict): operator = payload.get("operator") if operator not in runner.operators: raise HTTPException(400, f"unknown operator: {operator}") asset_ids = payload.get("asset_ids") if not asset_ids: asset_ids = [payload["asset_id"]] if payload.get("asset_id") else [] for aid in asset_ids: if not store.get_asset(app.state.con, aid): raise HTTPException(400, f"unknown asset: {aid}") params = payload.get("params") or {} job = runner.create_job(app.state.con, operator, asset_ids, params) await runner.broadcast({"type": "job", "job": job}) return job @app.get("/api/jobs/{job_id}") def get_job(job_id: str): job = runner.get_job(app.state.con, job_id) if not job: raise HTTPException(404, "no such job") return job @app.post("/api/jobs/{job_id}/cancel") async def cancel_job(job_id: str): if not await runner.cancel(app.state.con, job_id): raise HTTPException(400, "job not cancellable") return {"ok": True} @app.post("/api/jobs/{job_id}/retry") async def retry_job(job_id: str): job = runner.get_job(app.state.con, job_id) if not job: raise HTTPException(404, "no such job") new = runner.create_job(app.state.con, job["operator"], job.get("asset_ids") or [], job["params"]) await runner.broadcast({"type": "job", "job": new}) return new @app.delete("/api/jobs/{job_id}") def delete_job(job_id: str): if not runner.delete_job(app.state.con, job_id): raise HTTPException(400, "job running or missing") return {"ok": True} # -- inbox watch folder --------------------------------------------------------- async def watch_inbox(): """Ingest files dropped into data/inbox/ once their size is stable (avoids grabbing partial copies). Lets UE/Blender/rsync drop outputs in for pickup.""" INBOX.mkdir(parents=True, exist_ok=True) seen: dict[str, int] = {} while True: await asyncio.sleep(3) try: for p in INBOX.iterdir(): if p.name.startswith(".") or not p.exists(): continue size = p.stat().st_size if p.is_file() else sum( f.stat().st_size for f in p.rglob("*") if f.is_file()) if seen.get(str(p)) == size and size > 0: asset = store.register_file(app.state.con, p, move=True, meta={"source": "inbox"}) seen.pop(str(p), None) await runner.broadcast({"type": "assets_changed", "assets": [asset["id"]]}) else: seen[str(p)] = size except Exception as e: print(f"[inbox] {e}") # -- websocket --------------------------------------------------------------- @app.websocket("/ws") async def ws(websocket: WebSocket): await websocket.accept() runner.subscribers.add(websocket) try: while True: await websocket.receive_text() except WebSocketDisconnect: runner.subscribers.discard(websocket) # -- frontend (built) ----------------------------------------------------------- if WEB_DIST.exists(): app.mount("/", StaticFiles(directory=WEB_DIST, html=True), name="web")