rune.py

"""Shared by every Rune Labs Space: the Rune client, its concurrency cap, and image handling.

deploy.py copies this file into each Space. Edit it here, never in a Space.
"""
import base64
import binascii
import http.client
import io
import json
import logging
import os
import threading
import time
import urllib.error
import urllib.request

from PIL import Image

MODEL = "rune-v3"
MAX_SIDE = 1024
MAX_TEXT = 4000
MAX_DATA_URL = 6 * 1024 * 1024
MAX_PIXELS = 25_000_000  # checked from the header, before any pixel is decoded
BUSY = "Rune is busy right now. Try again in a few seconds."
FAILED = "Rune could not answer that one. Try again in a moment."

log = logging.getLogger("rune")

# One cap for every caller in the process, Gradio handler or FastAPI route alike.
_slots = threading.BoundedSemaphore(2)
SLOT_WAIT = 30  # seconds a caller waits for a slot before "busy"; never holds a server thread forever


class RuneError(Exception):
    """Rune did not answer. str() is safe to show to a user."""


class Busy(RuneError):
    pass


def request_body(state, questions, images=None):
    if isinstance(state, str):
        state = state[:MAX_TEXT]
    body = {"model": MODEL, "state": state, "questions": questions}
    if images:
        body["images"] = images
    return body


def decide(body, *, opener=urllib.request.urlopen, sleep=time.sleep):
    req = urllib.request.Request(
        os.environ["RUNE_URL"], json.dumps(body).encode(),
        {"Authorization": "Bearer " + os.environ["RUNE_API_KEY"], "Content-Type": "application/json"})
    if not _slots.acquire(timeout=SLOT_WAIT):
        log.warning("rune busy: no free slot")
        raise Busy(BUSY)
    try:
        for attempt in (0, 1):
            try:
                with opener(req, timeout=60) as r:
                    reply = json.load(r)
                if not isinstance(reply, dict) or not isinstance(reply.get("answers"), dict):
                    raise ValueError("reply has no answers")
                return reply["answers"]
            except urllib.error.HTTPError as e:
                e.close()
                if e.code not in (429, 503):
                    log.warning("rune call failed: HTTP %s", e.code)
                    raise RuneError(FAILED) from None
                if attempt == 0:
                    sleep(1.5)
            except (OSError, http.client.HTTPException, ValueError) as e:
                # Logged by type only: the request carries user content, the log never does.
                log.warning("rune call failed: %s", type(e).__name__)
                raise RuneError(FAILED) from None
    finally:
        _slots.release()
    log.warning("rune busy after one retry")
    raise Busy(BUSY)


def image_data_url(img):
    img = img.copy()
    img.thumbnail((MAX_SIDE, MAX_SIDE))
    png = img.mode in ("RGBA", "LA", "P")
    buf = io.BytesIO()
    if png:
        img.save(buf, "PNG")
    else:
        img.convert("RGB").save(buf, "JPEG", quality=90)
    return f"data:image/{'png' if png else 'jpeg'};base64," + base64.b64encode(buf.getvalue()).decode()


def from_data_url(url):
    """Decode an untrusted data URL from a browser into a resized image. Raises ValueError."""
    if len(url) > MAX_DATA_URL or not url.startswith(("data:image/png;base64,", "data:image/jpeg;base64,")):
        raise ValueError("expected a PNG or JPEG data URL under 6 MB")
    try:
        img = Image.open(io.BytesIO(base64.b64decode(url.split(",", 1)[1], validate=True)), formats=["PNG", "JPEG"])
        if img.width * img.height > MAX_PIXELS:
            raise ValueError("image too large")
        img.load()
    except (binascii.Error, OSError, Image.DecompressionBombError) as e:
        raise ValueError("not a readable image") from e
    img.thumbnail((MAX_SIDE, MAX_SIDE))
    return img


def as_label(answer, names=None):
    """Rune answer -> {label: probability}, the shape gr.Label and the live pages draw.

    names maps a choice's option keys to what to show instead, usually its criteria.
    """
    if "noul" in answer:
        return {"yes": answer["noul"], "no": 1 - answer["noul"]}
    names = names or answer.get("legend") or {}
    return {names.get(k, k): p for k, p in answer["probabilities"].items()}


def curl(body):
    shown = dict(body)
    if "images" in shown:
        shown["images"] = [u[:40] + "…" for u in shown["images"]]
    return (f"curl {os.environ.get('RUNE_URL', '')} \\\n  -H 'Authorization: Bearer $RUNE_API_KEY' \\\n"
            f"  -H 'Content-Type: application/json' \\\n  -d '{json.dumps(shown, ensure_ascii=False)}'")