runtime/src/decisionnode_edge/request.py

"""The request of the hosted `/v1/decide` endpoint, read into the questions the model answers.

    {"state": ..., "questions": {"<id>": {"type": "choice" | "score" | "truth" | "number", ...}}, "images": [...]}

What a request becomes (the same rules as the hosted API):

    state         a string is used as it is; an object or array becomes compact JSON with its key order
                  kept; a missing state, null and "" are the empty state.
    instructions  a string, or the compact JSON of an object or array; missing or null is the empty text.
    choice        criteria {option: description}, 1 to 255 options. A description is a string, the compact
                  JSON of an object or array, or "" for null.
    score         criteria [level, ...], 1 to 10 levels. One level is answered without the model.
    truth         optional criteria {"true": ..., "false": ...}; needs instructions or criteria.
    number        min, max and step (default 0, 255, 1): the answer is a value of that grid. Up to 256 values
                  are one internal choice over the values; up to 4,096 are answered as pages of 256 (one choice
                  over the pages, one per page) and composed. A grid of one value is answered without the model.
    images        [{"id", "data" (base64) and "media_type"}] or, on Edge, [{"id", "path"}] for a local file.
                  Every question sees every picture, in request order.

Edge has no videos, boosters, uploads, image URLs, points or boxes: a request that names them is refused.
"""

from __future__ import annotations

import json
import math
from dataclasses import dataclass, field
from decimal import Decimal, localcontext
from typing import Any

MAX_OPTIONS = 255
MAX_LEVELS = 10
MAX_IMAGES = 16
MAX_IMAGE_ID = 64
IMAGE_MEDIA_TYPES = ("image/jpeg", "image/png", "image/webp")
ONE_ROW_VALUES = 256  # the answer codes: a grid up to this size is one internal choice
MAX_GRID_VALUES = 4096
MAX_DECIMALS = 6
NUMBER_DEFAULTS = {"min": 0, "max": 255, "step": 1}
NOT_ON_EDGE = ("videos", "boosters")


class RequestError(ValueError):
    """A request the model cannot answer. `status` is the HTTP status the hosted API gives the same fault, `error_type`
    the kind of fault its error body names (the local server's, server.py)."""

    def __init__(self, message: str, status: int = 400, error_type: str = "api_usage_error"):
        super().__init__(message)
        self.status = status
        self.error_type = error_type


def text_of(value: Any) -> str:
    """A state, instruction, description or level as the prompt carries it."""
    if value is None:
        return ""
    if isinstance(value, str):
        return value
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"))


def _content(value: Any, where: str, nullable: bool = True) -> None:
    if not (isinstance(value, (str, dict, list)) or (value is None and nullable)):
        raise RequestError(f"{where} must be a string, an object or an array", 422)


# ---------------------------------------------------------------------------
# Number grids
# ---------------------------------------------------------------------------


def _decimal(value: Any) -> Decimal:
    return Decimal(repr(value)) if isinstance(value, float) else Decimal(int(value))


def _places(value: Decimal) -> int:
    """Digits after the point once trailing zeros are dropped: 0.5 has 1, 7 has 0."""
    text = format(value.normalize(), "f")
    return len(text.split(".")[1]) if "." in text else 0


def _value_text(scaled: int, d: int) -> str:
    """The grid value scaled / 10^d in its shortest plain decimal form."""
    digits = str(abs(scaled)).rjust(d + 1, "0")
    whole, frac = (digits[:-d], digits[-d:].rstrip("0")) if d else (digits, "")
    text = whole + ("." + frac if frac else "")
    return "-" + text if scaled < 0 else text


@dataclass(frozen=True)
class Grid:
    """Value i is (m + i * s) / 10^d, for i in 0 .. n - 1."""

    d: int
    m: int
    s: int
    n: int

    def text(self, i: int) -> str:
        return _value_text(self.m + i * self.s, self.d)

    def texts(self, start: int = 0, end: int | None = None) -> list[str]:
        return [self.text(i) for i in range(start, self.n if end is None else end)]

    @property
    def pages(self) -> int:
        return 1 if self.n <= ONE_ROW_VALUES else -(-self.n // ONE_ROW_VALUES)

    def page_bounds(self) -> list[tuple[int, int]]:
        """(first, end) of each page; the first n mod pages pages hold one value more."""
        q, r = divmod(self.n, self.pages)
        out, start = [], 0
        for j in range(self.pages):
            size = q + 1 if j < r else q
            out.append((start, start + size))
            start += size
        return out


def number_grid(question: dict, qid: str) -> Grid:
    values = {}
    for name in ("min", "max", "step"):
        value = question.get(name)
        if value is None:
            value = NUMBER_DEFAULTS[name]
        if isinstance(value, bool) or not isinstance(value, (int, float)) or (isinstance(value, float) and not math.isfinite(value)):
            raise RequestError(f"question {qid!r}: {name} must be a finite number", 422)
        values[name] = _decimal(value)
        if abs(values[name]) > 2**53:
            raise RequestError(f"question {qid!r}: {name} must be within plus or minus 2^53", 422)
        if name != "max" and _places(values[name]) > MAX_DECIMALS:
            raise RequestError(f"question {qid!r}: {name} has more than {MAX_DECIMALS} decimals", 422)
    if values["step"] <= 0:
        raise RequestError(f"question {qid!r}: step must be above 0", 422)
    if values["min"] > values["max"]:
        raise RequestError(f"question {qid!r}: min is above max", 422)
    d = max(_places(values["min"]), _places(values["step"]))
    with localcontext() as ctx:
        ctx.prec = 200
        if d:
            bound = Decimal(10) ** (15 - d) - Decimal(10) ** -d
            if abs(values["min"]) > bound or abs(values["max"]) > bound:
                raise RequestError(f"question {qid!r}: on a grid with {d} decimals, min and max must stay within plus or minus {bound}", 422)
        scale = Decimal(10) ** d
        m, s = int(values["min"] * scale), int(values["step"] * scale)
        top = int((values["max"] * scale).to_integral_value(rounding="ROUND_FLOOR"))
    n = (top - m) // s + 1
    if n > MAX_GRID_VALUES:
        raise RequestError(f"question {qid!r}: the grid holds {n} values, the limit is {MAX_GRID_VALUES}", 422)
    return Grid(d=d, m=m, s=s, n=n)


def number_rows(instructions: str, grid: Grid) -> list[dict]:
    """The internal choice questions of a number question: one over the values, or one over the pages and one per page."""
    if grid.pages == 1:
        keys = grid.texts()
        return [{"type": "choice", "instructions": instructions, "keys": keys, "descs": [""] * len(keys)}]
    bounds = grid.page_bounds()
    rows = [{"type": "choice", "instructions": instructions, "keys": [f"{grid.text(a)} to {grid.text(b - 1)}" for a, b in bounds], "descs": [""] * len(bounds)}]
    for a, b in bounds:
        keys = grid.texts(a, b)
        rows.append({"type": "choice", "instructions": instructions, "keys": keys, "descs": [""] * len(keys)})
    return rows


def compose(grid: Grid, answered: list[list[float]]) -> list[float]:
    """The probability of every grid value: p(page) * p(value | page) on a paged grid."""
    if grid.pages == 1:
        return list(answered[0])
    out: list[float] = []
    for j, page in enumerate(answered[1:]):
        out.extend(answered[0][j] * p for p in page)
    return out


# ---------------------------------------------------------------------------
# The request
# ---------------------------------------------------------------------------


@dataclass
class Question:
    id: str
    type: str  # choice, score, truth or number (as answered)
    rows: list[dict]  # what the model runs: one prompt per row (none for a question answered without the model)
    levels: list | None = None  # score: the levels as sent, for `legend`
    keys: list[str] | None = None  # choice: the options in request order
    grid: Grid | None = None  # number


@dataclass
class Request:
    state: str
    questions: list[Question]
    images: list[dict] = field(default_factory=list)


def parse(body: Any) -> Request:
    """Validate a `/v1/decide` request body (a dict) and read it into questions. Raises RequestError."""
    if not isinstance(body, dict):
        raise RequestError("the request must be a JSON object", 422)
    for name in NOT_ON_EDGE:
        if body.get(name):
            raise RequestError(f"`{name}` is not available on DecisionNode Edge")
    state = body.get("state")
    _content(state, "state")
    questions = body.get("questions")
    if not isinstance(questions, dict) or not questions:
        raise RequestError("questions must be an object with at least one question", 422)
    state_text = text_of(state)
    out = [_question(qid, q) for qid, q in questions.items()]
    return Request(state=state_text, questions=out, images=_images(body.get("images")))


def _question(qid: str, q: Any) -> Question:
    if not isinstance(q, dict) or "type" not in q:
        raise RequestError(f"question {qid!r} must be an object with a type", 422)
    kind = q["type"]
    if kind == "noul":  # the older name of the yes-or-no type, answered as truth
        kind = "truth"
    if kind not in ("choice", "score", "truth", "number"):
        raise RequestError(f"question {qid!r}: unknown question type {kind!r}")
    instructions, criteria = q.get("instructions"), q.get("criteria")
    _content(instructions, f"question {qid!r}: instructions")
    text = text_of(instructions)

    if kind == "number":
        if instructions is None or instructions == "":
            raise RequestError(f"Number question must have instructions: {qid}")
        grid = number_grid(q, qid)
        rows = [] if grid.n == 1 else number_rows(text, grid)
        return Question(qid, "number", rows, grid=grid)

    if kind == "choice":
        if not isinstance(criteria, dict):
            raise RequestError(f"question {qid!r}: a choice needs criteria, an object of options", 422)
        if not criteria:
            raise RequestError(f"Choice question must have at least one choice: {qid}")
        if len(criteria) > MAX_OPTIONS:
            raise RequestError(f"Too many choices: question {qid!r} has {len(criteria)} options, the limit is {MAX_OPTIONS} options per choice.")
        for key, desc in criteria.items():
            _content(desc, f"question {qid!r}: the description of {key!r}")
        keys = list(criteria)
        row = {"type": "choice", "instructions": text, "keys": keys, "descs": [text_of(d) for d in criteria.values()]}
        return Question(qid, "choice", [] if len(keys) == 1 else [row], keys=keys)

    if kind == "score":
        if not isinstance(criteria, list) or not criteria:
            raise RequestError(f"question {qid!r}: a score needs criteria, a list of 1 to {MAX_LEVELS} levels", 422)
        if len(criteria) > MAX_LEVELS:
            raise RequestError(f"Too many score levels. Must have at most {MAX_LEVELS} levels.")
        for i, level in enumerate(criteria):
            _content(level, f"question {qid!r}: level {i}", nullable=False)
        row = {"type": "score", "instructions": text, "levels": [text_of(level) for level in criteria]}
        return Question(qid, "score", [] if len(criteria) == 1 else [row], levels=list(criteria))

    # truth
    if criteria is not None and not isinstance(criteria, dict):
        raise RequestError(f"question {qid!r}: truth criteria must be an object with `true` and `false`", 422)
    if instructions is None and criteria is None:
        raise RequestError(f"Truth question must have criteria or instructions: {qid}")
    row: dict = {"type": "truth", "instructions": text}
    if criteria:
        for side in ("true", "false"):
            _content(criteria.get(side), f"question {qid!r}: criteria.{side}")
        sides = {side: text_of(criteria.get(side)) for side in ("true", "false") if text_of(criteria.get(side))}
        if sides:
            row["criteria"] = sides
    return Question(qid, "truth", [row])


def _images(images: Any) -> list[dict]:
    if images is None:
        return []
    if not isinstance(images, list):
        raise RequestError("images must be a list", 422)
    if len(images) > MAX_IMAGES:
        raise RequestError(f"Too many images: {len(images)}, the limit is {MAX_IMAGES}.")
    seen = set()
    for i, image in enumerate(images):
        if not isinstance(image, dict) or not isinstance(image.get("id"), str) or not image["id"]:
            raise RequestError(f"Image {i} needs an id.")
        image_id = image["id"]
        if len(image_id) > MAX_IMAGE_ID:
            raise RequestError(f"Image {i}: its id is longer than {MAX_IMAGE_ID} characters.")
        if image_id in seen:
            raise RequestError(f"Image id {image_id!r} is used twice.")
        seen.add(image_id)
        if image.get("url") is not None or image.get("upload_id") is not None:
            raise RequestError(f"Image {image_id!r}: DecisionNode Edge reads pictures from `data` or a local `path`, not from a URL or an upload.")
        forms = [name for name in ("data", "path") if image.get(name) is not None]
        if len(forms) != 1:
            raise RequestError(f"Image {image_id!r} needs exactly one of data and path.")
        if forms == ["data"] and image.get("media_type") not in IMAGE_MEDIA_TYPES:
            raise RequestError(f"Image {image_id!r}: media_type must be one of {', '.join(IMAGE_MEDIA_TYPES)}.")
    return images