live.py

"""FastAPI app for the live Rune Labs demos: a static page plus /decide.

deploy.py copies this file into the Docker Spaces. Edit it here, never in a Space.
"""
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

import rune

MAX_BODY = rune.MAX_DATA_URL + 64 * 1024


class Ask(BaseModel):
    image: str
    before: str | None = None  # an optional earlier frame, sent first, for "compare the two" questions
    context: dict = {}


def make_app(build_questions, static_dir, prepare=None, extend=None):
    """build_questions(context) -> (state, questions) for one decision about the posted image(s).

    context["images"] is set here to how many images came (1, or 2 with a "before" image). A ValueError from
    build_questions is a 400. prepare(image) -> image draws on each image before Rune sees it; extend(app)
    adds routes, ahead of the static page that would otherwise catch them."""
    app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None)

    @app.middleware("http")
    async def refuse_oversized(request: Request, call_next):
        # Checked before FastAPI reads and parses the body; a data URL plus the JSON around it.
        size = request.headers.get("content-length")
        if request.method == "POST" and (size is None or not size.isdigit() or int(size) > MAX_BODY):
            return JSONResponse({"detail": "Request too large."}, status_code=413 if size else 411)
        return await call_next(request)

    @app.post("/decide")
    def ask(a: Ask):
        try:
            imgs = [rune.from_data_url(u) for u in (a.before, a.image) if u is not None]
            state, questions = build_questions({**a.context, "images": len(imgs)})
        except ValueError as e:
            raise HTTPException(400, str(e))
        if prepare:
            imgs = [prepare(i) for i in imgs]
        body = rune.request_body(state, questions, [rune.image_data_url(i) for i in imgs])
        try:
            answers = rune.decide(body)
        except rune.Busy as e:
            raise HTTPException(503, str(e))
        except rune.RuneError as e:
            raise HTTPException(502, str(e))
        return {"labels": {k: rune.as_label(v) for k, v in answers.items()}}

    if extend:
        extend(app)
    app.mount("/", StaticFiles(directory=static_dir, html=True), name="static")
    return app