"""Hugging Face Space: server-side JEV calls and real CPU ViZDoom sessions.""" from __future__ import annotations import asyncio import base64 import binascii import contextlib import io import math import os import re import secrets import time import warnings from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager from pathlib import Path from typing import Annotated, Literal from urllib.parse import urlsplit import httpx from fastapi import FastAPI, HTTPException, Request from fastapi.exceptions import RequestValidationError from fastapi.responses import FileResponse, JSONResponse from PIL import Image, ImageOps, UnidentifiedImageError from pydantic import BaseModel, ConfigDict, Field, StringConstraints, model_validator from doom_engine import DoomEpisode, SCENARIOS from arena_api import router as arena_router, worker_token ROOT = Path(__file__).resolve().parent MODEL = "autotrust/JEV-27B" # Named in every decision when the backend serves more than one model (api.scienceguru.ai: autotrust/JEV-27B-VL). BACKEND_MODEL = os.getenv("JEV_MODEL", "").strip() SESSION_LIMIT = 8 SESSION_TTL = 600 QUEUE_WAIT_SECONDS = 5 MAX_IMAGE_BYTES = 5 * 1024 * 1024 MAX_IMAGE_BASE64 = 4 * ((MAX_IMAGE_BYTES + 2) // 3) MAX_IMAGE_PIXELS = 32_000_000 MAX_IMAGE_EDGE = 1024 SAFE_ASSETS = {".css", ".js", ".svg", ".png", ".jpg", ".json", ".mp4", ".md"} Option = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=180)] class DecisionRequest(BaseModel): model_config = ConfigDict(extra="forbid") kind: Literal["choice", "noul", "score"] = "choice" state: Annotated[str, StringConstraints(max_length=6000)] = "" question: Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=1000)] options: list[Option] | None = Field(default=None, min_length=2, max_length=16) image: Annotated[str, StringConstraints(max_length=MAX_IMAGE_BASE64 + 32)] | None = None @model_validator(mode="after") def native_options(self): required = {"noul": ["false", "true"], "score": [str(i) for i in range(6)]} if self.kind in required: if self.options is None: self.options = required[self.kind] elif self.options != required[self.kind]: raise ValueError(f"{self.kind} requires options {required[self.kind]}") if self.options is None: raise ValueError("Choice requires 2 to 16 options") if len(set(self.options)) != len(self.options): raise ValueError("Options must be distinct") return self class GameRequest(BaseModel): model_config = ConfigDict(extra="forbid") scenario: Literal["basic", "defend_the_center"] = "basic" class StepRequest(BaseModel): model_config = ConfigDict(extra="forbid") controller: Literal["human", "jev"] action: Annotated[str, StringConstraints(max_length=32)] | None = None @model_validator(mode="after") def action_required(self): if self.controller == "human" and not self.action: raise ValueError("A human step requires an action") if self.controller == "jev" and self.action is not None: raise ValueError("JEV selects its own action; omit action") return self class GameSession: def __init__(self, scenario: str): self.scenario = scenario self.lock = asyncio.Lock() self.last_active = time.monotonic() self.executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="doom") self.engine: DoomEpisode | None = None async def run(self, function, *args): return await asyncio.get_running_loop().run_in_executor(self.executor, function, *args) async def start(self): self.engine = await self.run(DoomEpisode, self.scenario) async def close(self): try: if self.engine is not None: await self.run(self.engine.close) finally: self.executor.shutdown(wait=False, cancel_futures=True) def backend_settings() -> tuple[str, str]: base = os.getenv("JEV_BASE_URL", "").strip().rstrip("/") key = os.getenv("JEV_API_KEY", "").strip() parsed = urlsplit(base) # Configuration is server-only; user input can never select an outbound host. if (parsed.scheme != "https" or not parsed.netloc or parsed.username or parsed.password or parsed.query or parsed.fragment or not key): return "", "" return base, key async def reap_sessions(app: FastAPI): while True: await asyncio.sleep(30) expired = [] async with app.state.sessions_lock: for session_id, session in list(app.state.sessions.items()): if not session.lock.locked() and time.monotonic() - session.last_active > SESSION_TTL: expired.append(app.state.sessions.pop(session_id)) for session in expired: await session.close() @asynccontextmanager async def lifespan(app: FastAPI): base, key = backend_settings() app.state.configured = bool(base and key) app.state.base_url = base app.state.inference_slots = asyncio.Semaphore(4) app.state.health_lock = asyncio.Lock() app.state.health_cache = None app.state.sessions = {} app.state.sessions_lock = asyncio.Lock() async with httpx.AsyncClient( headers={"Authorization": f"Bearer {key}"} if key else {}, timeout=httpx.Timeout(120, connect=10, pool=10), limits=httpx.Limits(max_connections=5, max_keepalive_connections=5), follow_redirects=False, trust_env=False, ) as client: app.state.client = client reaper = asyncio.create_task(reap_sessions(app)) try: yield finally: reaper.cancel() with contextlib.suppress(asyncio.CancelledError): await reaper for session in list(app.state.sessions.values()): await session.close() app = FastAPI(title="JEV-27B Live Demo", lifespan=lifespan, docs_url=None, redoc_url=None) @app.exception_handler(RequestValidationError) async def invalid_request(request: Request, error: RequestValidationError): # Do not echo uploaded image data back in validation errors. return JSONResponse({"detail": [ {"loc": item["loc"], "msg": item["msg"], "type": item["type"]} for item in error.errors() ]}, status_code=422) @app.middleware("http") async def basic_limits(request: Request, call_next): if request.url.path.startswith("/api/arena/worker/"): expected = worker_token() if not expected or not secrets.compare_digest(request.headers.get("authorization", ""), "Bearer " + expected): return JSONResponse({"detail": "Worker authentication required."}, status_code=401) if request.method in {"POST", "PUT", "PATCH"}: body = bytearray() async for chunk in request.stream(): body.extend(chunk) limit = (20 * 1024 * 1024 if request.url.path.startswith("/api/arena/worker/") else MAX_IMAGE_BASE64 + 40_000 if request.url.path == "/api/decide" else 40_000) if len(body) > limit: return JSONResponse({"detail": "Request body is too large."}, status_code=413) # Starlette replays the cached body when call_next constructs its Request. request._body = bytes(body) response = await call_next(request) response.headers["X-Content-Type-Options"] = "nosniff" if request.url.path.startswith("/api/"): response.headers["Cache-Control"] = "no-store" return response @app.get("/api/health") async def health(): if not app.state.configured: return JSONResponse({"status": "not_configured", "model": MODEL, "configured": False}, status_code=503) async with app.state.health_lock: cached = app.state.health_cache if cached and time.monotonic() - cached[0] < 15: return JSONResponse(cached[1], status_code=cached[2]) ok = False try: response = await app.state.client.get(app.state.base_url + "/v1/models", timeout=8) if response.status_code == 200: data = response.json() ids = {item.get("id") for item in data.get("data", []) if isinstance(item, dict)} ok = bool(ids & {MODEL, BACKEND_MODEL or MODEL, "jev-decision"}) except (httpx.HTTPError, ValueError, TypeError, AttributeError): pass data = {"status": "ok" if ok else "unavailable", "model": MODEL, "configured": True} code = 200 if ok else 503 app.state.health_cache = (time.monotonic(), data, code) return JSONResponse(data, status_code=code) def clean_decision(data: dict, payload: dict) -> dict: """Validate provenance and shape; expose only allowlisted, bounded values.""" options = payload["options"] if data.get("model") != (BACKEND_MODEL or MODEL) or data.get("options") != options: raise ValueError("Unexpected model or options") image_used = bool(payload.get("image")) # A vision-language backend reads the image itself and reports no separate vision step. if image_used and not BACKEND_MODEL and data.get("image_used") is not True: raise ValueError("The backend did not confirm image inference") probabilities = data.get("probabilities") if not isinstance(probabilities, list) or len(probabilities) != len(options): raise ValueError("Invalid distribution") if any(type(p) not in (int, float) or not math.isfinite(p) or not 0 <= p <= 1 for p in probabilities) or abs(sum(probabilities) - 1) > 0.001: raise ValueError("Invalid probabilities") index = data.get("choice_index") if type(index) is not int or not 0 <= index < len(options): raise ValueError("Invalid choice") if probabilities[index] != max(probabilities) or data.get("choice") != options[index]: raise ValueError("Inconsistent choice") elapsed = data.get("elapsed_seconds") if type(elapsed) not in (int, float) or not math.isfinite(elapsed) or elapsed < 0: raise ValueError("Invalid timing") revision = data.get("revision", "") if not isinstance(revision, str) or not re.fullmatch(r"[a-fA-F0-9]{7,64}", revision): revision = "unavailable" result = {"kind": payload["kind"], "options": options, "probabilities": probabilities, "choice_index": index, "choice": options[index], "model": MODEL, "revision": revision, "elapsed_seconds": elapsed, "image_used": image_used} if image_used and not BACKEND_MODEL: vision = data.get("vision") if not isinstance(vision, dict): raise ValueError("Missing image understanding provenance") vision_model = vision.get("model") vision_elapsed = vision.get("elapsed_seconds") if (not isinstance(vision_model, str) or not re.fullmatch(r"[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+", vision_model) or len(vision_model) > 160 or type(vision_elapsed) not in (int, float) or not math.isfinite(vision_elapsed) or vision_elapsed < 0): raise ValueError("Invalid image understanding provenance") result["vision"] = {"model": vision_model, "elapsed_seconds": vision_elapsed} return result def prepare_image(data_url: str) -> str: """Decode and normalize a user image without fetching URLs or storing files.""" header, separator, encoded = data_url.partition(",") formats = {"data:image/jpeg;base64": "JPEG", "data:image/png;base64": "PNG", "data:image/webp;base64": "WEBP"} if not separator or header not in formats: raise HTTPException(422, "Upload a JPEG, PNG, or WebP image.") try: raw = base64.b64decode(encoded, validate=True) except (ValueError, binascii.Error): raise HTTPException(422, "The image could not be read. Choose another file.") from None if not raw or len(raw) > MAX_IMAGE_BYTES: raise HTTPException(413, "Choose an image no larger than 5 MB.") try: with warnings.catch_warnings(): warnings.simplefilter("error", Image.DecompressionBombWarning) with Image.open(io.BytesIO(raw)) as probe: if probe.format != formats[header] or getattr(probe, "is_animated", False): raise ValueError("Unsupported image format") if probe.width * probe.height > MAX_IMAGE_PIXELS: raise ValueError("Image dimensions are too large") probe.verify() with Image.open(io.BytesIO(raw)) as source: upright = ImageOps.exif_transpose(source) upright.thumbnail((MAX_IMAGE_EDGE, MAX_IMAGE_EDGE), Image.Resampling.LANCZOS) if "A" in upright.getbands() or "transparency" in upright.info: rgba = upright.convert("RGBA") normalized = Image.new("RGB", rgba.size, "white") normalized.paste(rgba, mask=rgba.getchannel("A")) else: normalized = upright.convert("RGB") output = io.BytesIO() normalized.save(output, format="JPEG", quality=90) return "data:image/jpeg;base64," + base64.b64encode(output.getvalue()).decode("ascii") except (UnidentifiedImageError, OSError, ValueError, SyntaxError, Image.DecompressionBombError, Image.DecompressionBombWarning): raise HTTPException(422, "Use a valid, still JPEG, PNG, or WebP image up to 32 megapixels.") from None def backend_body(payload: dict) -> dict: if not BACKEND_MODEL: return payload body = {key: value for key, value in payload.items() if key != "image"} body["model"] = BACKEND_MODEL if payload.get("image"): body["state"] = [payload["state"], {"image": payload["image"]}] if payload.get("state") else [{"image": payload["image"]}] return body async def infer(payload: dict) -> dict: if not app.state.configured: raise HTTPException(503, "JEV is not connected yet. The Space owner must configure its backend secrets.") try: await asyncio.wait_for(app.state.inference_slots.acquire(), QUEUE_WAIT_SECONDS) except TimeoutError: raise HTTPException(429, "JEV is busy. Please retry shortly.", headers={"Retry-After": "5"}) from None try: try: response = await app.state.client.post(app.state.base_url + "/v1/decide", json=backend_body(payload)) except httpx.TimeoutException: raise HTTPException(504, "JEV took too long to respond. The game has not advanced.") from None except httpx.HTTPError: raise HTTPException(503, "JEV is temporarily unreachable. Please retry.") from None if response.status_code == 429: raise HTTPException(429, "JEV is busy. Please retry shortly.", headers={"Retry-After": "5"}) if response.status_code != 200: raise HTTPException(502, "JEV could not complete this decision. Please retry.") try: return clean_decision(response.json(), payload) except (ValueError, TypeError, KeyError, AttributeError): raise HTTPException(502, "JEV returned an invalid decision. No action was executed.") from None finally: app.state.inference_slots.release() @app.post("/api/decide") async def decide(payload: DecisionRequest): data = payload.model_dump(exclude_none=True) if payload.image is not None: data["image"] = await asyncio.to_thread(prepare_image, payload.image) return await infer(data) @app.post("/api/games/new") async def new_game(payload: GameRequest): async with app.state.sessions_lock: if len(app.state.sessions) >= SESSION_LIMIT: raise HTTPException(429, "All game slots are occupied. Close a game or retry later.", headers={"Retry-After": "30"}) session = GameSession(payload.scenario) try: await session.start() snapshot = await session.run(session.engine.snapshot) except Exception: await session.close() raise HTTPException(503, "The game engine could not start. Please retry.") from None session_id = secrets.token_urlsafe(24) app.state.sessions[session_id] = session return {"session_id": session_id, **snapshot} @asynccontextmanager async def locked_session(session_id: str): session = app.state.sessions.get(session_id) if session is None: raise HTTPException(404, "This game has expired. Start a new game.") try: await asyncio.wait_for(session.lock.acquire(), QUEUE_WAIT_SECONDS) except TimeoutError: raise HTTPException(429, "This game is already processing an action.", headers={"Retry-After": "2"}) from None try: if app.state.sessions.get(session_id) is not session: raise HTTPException(404, "This game has expired. Start a new game.") if time.monotonic() - session.last_active > SESSION_TTL: async with app.state.sessions_lock: app.state.sessions.pop(session_id, None) await session.close() raise HTTPException(404, "This game has expired. Start a new game.") session.last_active = time.monotonic() yield session finally: session.last_active = time.monotonic() session.lock.release() @app.post("/api/games/{session_id}/step") async def step_game(session_id: str, payload: StepRequest): async with locked_session(session_id) as session: snapshot = await session.run(session.engine.snapshot) if snapshot["done"]: raise HTTPException(409, "The episode has finished. Start a new game.") actions = SCENARIOS[session.scenario]["actions"] decision = None action_id = payload.action if payload.controller == "jev": request = await session.run(session.engine.decision_payload) decision = await infer(request) action_id = actions[decision["choice_index"]]["id"] elif action_id not in {action["id"] for action in actions}: raise HTTPException(422, "That action is not available in this scenario.") try: snapshot = await session.run(session.engine.step, action_id) except Exception: raise HTTPException(503, "The game engine stopped. Start a new game.") from None response = {"session_id": session_id, **snapshot, "controller": payload.controller, "action": action_id} if decision is not None: response["decision"] = decision return response @app.delete("/api/games/{session_id}") async def delete_game(session_id: str): async with locked_session(session_id) as session: async with app.state.sessions_lock: app.state.sessions.pop(session_id, None) await session.close() return {"closed": True} app.include_router(arena_router) @app.get("/") async def home(): return FileResponse(ROOT / "index.html") @app.get("/showcase") @app.get("/index.html") async def index(): return FileResponse(ROOT / "index.html") @app.get("/play") @app.get("/play.html") async def play(): return FileResponse(ROOT / "play.html") @app.get("/{filename}") async def asset(filename: str): candidate = ROOT / filename if (filename.startswith(".") or candidate.name != filename or candidate.suffix.lower() not in SAFE_ASSETS or not candidate.is_file() or candidate.resolve().parent != ROOT): raise HTTPException(404, "Not found") return FileResponse(candidate) if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860, access_log=False)