runtime/src/decisionnode_edge/server.py

"""A local HTTP server for a DecisionNode Edge model, on the standard library only.

    python -m decisionnode_edge serve path/to/model [--host 127.0.0.1] [--port 8000] [--api-key FILE] [--backend ...]

    POST /v1/systemone   the /v1/systemone format used by public decision-model leaderboards:
                         {"model", "state", "questions", "assets"}; noul, choice, score and box questions; a noul is
                         answered {"type": "noul", "noul": p}, a choice and a score as on /v1/decide but with lists
                         that sum to exactly 1.00, a box as {"type": "box", "box": {"top", "bottom", "left", "right"},
                         "confidence"} (systemone.py)
    POST /v1/decide      the DecisionNode API's call: the runtime's own request and answers (EdgeModel.decide)
    GET  /v1/models      {"models": [{"name", "description", "release_date"}]}
    GET  /healthz        {"status": "ok", "model", "backend"} (no key needed)

Pictures: `images` as the runtime takes them, base64 `data` only (a local `path` is refused over HTTP: a caller must
never make the server read its files), and on /v1/systemone the format's own: the `assets` map, {"<ref>": {"type":
"image", "format": "png" | "jpg" | "jpeg", "data_base64": ...}}, whose references in the text become `Picture <n>`,
`image_url` content parts and `images` as data URLs (systemone.translate).

Point and box questions are taken on /v1/decide. On /v1/systemone a box is taken (in the format's shape, its
`confidence` the model folder's box calibration where there is one, box_calibration.json, else `present`), and point
and number questions are refused.

Errors carry a `detail` body: {"detail": {"error_type", "message"}} (a 422 as a list of field errors). A
request over the model's length limit says "maximum context length", a choice over 255 options "options per choice",
and on /v1/systemone a score over 10 levels "a score takes 2 to 10 levels": the words leaderboard harnesses look for to
count a request as beyond the model rather than as a failure. On /v1/systemone these three come in the hosted route's
bodies (systemone.mark_capacity).

The model answers one request at a time (a lock); the server listens on the loopback address unless told otherwise.
`model` is required on /v1/systemone, as on the hosted API, and any name is accepted: the answer names the model that
answered.
"""

from __future__ import annotations

import argparse
import hmac
import json
import logging
import os
import sys
import threading
import time
import traceback
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Any

from . import request as requests, systemone
from .request import RequestError

MAX_BODY_BYTES = 64 << 20  # 16 pictures as base64 with room to spare
SYSTEMONE_TYPES = ("noul", "choice", "score", "box")
AUTH_MESSAGE = "Cannot authenticate with the server. Please check your API key and try again."
ROUTES = {"/v1/systemone": "POST", "/v1/decide": "POST", "/v1/models": "GET", "/healthz": "GET"}

log = logging.getLogger("decisionnode_edge.server")


class HTTPError(Exception):
    def __init__(self, status: int, error_type: str, message: str, headers: dict | None = None):
        super().__init__(message)
        self.status, self.error_type, self.message, self.headers = status, error_type, message, headers or {}


def error_body(status: int, error_type: str, message: str, errors: list[dict] | None = None) -> dict:
    if status == 422:
        return {"detail": errors or [{"type": "value_error", "loc": ["body"], "msg": message, "input": None}]}
    return {"detail": {"error_type": error_type, "message": message}}


def _images_over_http(images: Any) -> list:
    if images is None:
        return []
    if not isinstance(images, list):
        raise RequestError("images must be a list", 422)
    for image in images:
        if isinstance(image, dict) and image.get("path") is not None:
            raise RequestError("Over HTTP a picture is sent as base64 `data` (or an asset), never as a local `path`.")
    return list(images)


def box_calibration(model: Any, path: str | None = None, sha256: str | None = None) -> systemone.BoxCalibration | None:
    """The box calibration /v1/systemone reads: `path` (or "none" for none), else the model folder's
    box_calibration.json when there is one. `sha256` (else `box_calibration_sha256` in decisionnode_edge.json), when
    set, must be the file's. A file that does not check is a BoxCalibrationError: the server does not start on it."""
    if path == "none":
        return None
    sha256 = sha256 if sha256 is not None else model.config.get("box_calibration_sha256") or ""
    if path is None:
        path = os.path.join(model.path, systemone.BOX_CALIBRATION_FILE)
        if not os.path.isfile(path):
            if sha256:
                raise systemone.BoxCalibrationError(f"{path}: the model names a box calibration that is not there")
            return None
    return systemone.BoxCalibration.load(path, sha256)


class Service:
    """The routes, without the socket: handle(method, path, headers, body) gives (status, JSON body, extra headers).
    `calibration`: the box calibration of /v1/systemone (default: box_calibration(model))."""

    def __init__(self, model: Any, api_key: str | None = None, description: str | None = None, release_date: str | None = None, calibration: Any = "model"):
        self.model = model
        self.api_key = api_key
        self.name = model.config["name"]
        self.description = description or model.config.get("description") or f"{self.name}: typed, calibrated decisions, run locally"
        self.release_date = release_date or model.config.get("release_date")
        self.calibration = box_calibration(model) if calibration == "model" else calibration
        self.lock = threading.Lock()

    def handle(self, method: str, path: str, headers: dict[str, str], body: bytes) -> tuple[int, dict, dict]:
        try:
            return 200, self._route(method, path.split("?", 1)[0], headers, body), {}
        except HTTPError as error:
            return error.status, error_body(error.status, error.error_type, error.message), error.headers
        except RequestError as error:
            if path.split("?", 1)[0] == "/v1/systemone":
                error = systemone.mark_capacity(error)
            if error.detail is not None:
                return error.status, {"detail": error.detail}, {}
            return error.status, error_body(error.status, error.error_type, str(error), error.errors), {}
        except Exception:  # a fault of the model or the backend, never the caller's
            log.error("the model could not answer:\n%s", traceback.format_exc())
            return 500, error_body(500, "internal_error", "The model could not answer this request."), {}

    def _route(self, method: str, path: str, headers: dict[str, str], body: bytes) -> dict:
        if path not in ROUTES:
            raise HTTPError(404, "not_found_error", "Not found.")
        if method != ROUTES[path]:
            raise HTTPError(405, "api_usage_error", "Method not allowed.", {"Allow": ROUTES[path]})
        if path == "/healthz":
            return {"status": "ok", "model": self.name, "backend": self.model.backend}
        self._authenticate(headers)
        if path == "/v1/models":
            return {"models": [{"name": self.name, "description": self.description, "release_date": self.release_date}]}
        try:
            request = json.loads(body)
        except (UnicodeDecodeError, json.JSONDecodeError) as error:
            raise RequestError(f"the body is not JSON: {error}", 422) from None
        if path == "/v1/decide":
            if isinstance(request, dict) and request.get("images") is not None:
                request = {**request, "images": _images_over_http(request["images"])}
            with self.lock:
                return self.model.decide(request)
        return self._systemone(request)

    def _authenticate(self, headers: dict[str, str]) -> None:
        if self.api_key is None:
            return
        given = headers.get("authorization", "")
        scheme, _, token = given.partition(" ")
        if scheme.lower() != "bearer" or not hmac.compare_digest(token.strip().encode(), self.api_key.encode()):
            raise HTTPError(401, "authentication_error", AUTH_MESSAGE, {"WWW-Authenticate": "Bearer"})

    def _systemone(self, request: Any) -> dict:
        if not isinstance(request, dict):
            raise RequestError("the request must be a JSON object", 422)
        if not isinstance(request.get("model"), str):
            raise RequestError("model: Field required, a string", 422)
        questions = request.get("questions")
        if isinstance(questions, dict):
            for qid, question in questions.items():
                kind = question.get("type") if isinstance(question, dict) else None
                if kind in ("number", "point"):
                    raise RequestError(f"question {qid!r}: the {kind} type is not available on /v1/systemone; send it to /v1/decide")
                if isinstance(question, dict) and "type" in question and kind not in SYSTEMONE_TYPES:
                    raise RequestError(f"question {qid!r}: unknown question type {kind!r}")
        _images_over_http(request.get("images"))
        body = {name: request[name] for name in ("state", "questions", "images", "assets") if name in request}
        body, done = systemone.translate(body, max_images=requests.MAX_IMAGES)
        parsed = requests.parse(body)
        systemone.normalised_boxes(q.place for q in parsed.questions if q.place is not None)  # the format's edges are fractions
        with self.lock:
            response = self.model._decide(parsed, exact_sum=True)
        answers = {}
        for qid, answer in response["answers"].items():  # a noul is answered under its own name, as the format has it
            if answer["type"] == "truth":
                answer = {"type": "noul", "noul": answer["truth"]}
            elif answer["type"] == "score" and qid in done.legends and len(done.legends[qid]) == len(answer["legend"]):
                answer = {**answer, "legend": {str(i): level for i, level in enumerate(done.legends[qid])}}  # the levels as sent
            answers[qid] = answer
        return systemone.box_answers({"model": response["model"], "answers": answers, "usage": response["usage"]}, self.calibration)


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"  # keep-alive: every response carries Content-Length
    disable_nagle_algorithm = True  # headers and body are two writes: without it a kept-alive client waits for a delayed ACK
    server_version = "decisionnode-edge"
    service: Service = None  # set by make_server

    def _answer(self, method: str) -> None:
        started = time.perf_counter()
        body = b""
        if method == "POST":
            length = self.headers.get("Content-Length")
            if length is None or not length.isdigit():
                return self._send(411, error_body(411, "api_usage_error", "Content-Length is required."), {"Connection": "close"}, started)
            if int(length) > MAX_BODY_BYTES:
                return self._send(413, error_body(413, "api_usage_error", "Request body too large."), {"Connection": "close"}, started)
            body = self.rfile.read(int(length))
        headers = {k.lower(): v for k, v in self.headers.items()}
        status, payload, extra = self.service.handle(method, self.path, headers, body)
        self._send(status, payload, extra, started)

    def _send(self, status: int, payload: dict, extra: dict, started: float) -> None:
        data = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(data)))
        self.send_header("X-Latency-Ms", f"{(time.perf_counter() - started) * 1000:.2f}")
        for name, value in extra.items():
            self.send_header(name, value)
        self.end_headers()
        self.wfile.write(data)
        if extra.get("Connection") == "close":
            self.close_connection = True

    def do_GET(self) -> None:  # noqa: N802 (the standard library's names)
        self._answer("GET")

    def do_POST(self) -> None:  # noqa: N802
        self._answer("POST")

    def do_PUT(self) -> None:  # noqa: N802
        self._answer("PUT")

    def do_DELETE(self) -> None:  # noqa: N802
        self._answer("DELETE")

    def log_message(self, format: str, *args: Any) -> None:  # noqa: A002
        log.info("%s %s", self.address_string(), format % args)


def make_server(model: Any, host: str = "127.0.0.1", port: int = 8000, api_key: str | None = None, **service: Any) -> ThreadingHTTPServer:
    handler = type("EdgeHandler", (Handler,), {"service": Service(model, api_key=api_key, **service)})
    server = ThreadingHTTPServer((host, port), handler)
    server.daemon_threads = True
    return server


def main(argv: list[str] | None = None) -> int:
    from . import EdgeModel

    parser = argparse.ArgumentParser(prog="python -m decisionnode_edge serve", description="Serve a DecisionNode Edge model over HTTP: /v1/systemone, /v1/decide, /v1/models.")
    parser.add_argument("model", help="the model directory")
    parser.add_argument("--host", default="127.0.0.1", help="the address to listen on (default 127.0.0.1, this machine only)")
    parser.add_argument("--port", type=int, default=8000)
    parser.add_argument("--api-key", default=None, help="a file holding the key callers must send as `Authorization: Bearer <key>`")
    parser.add_argument("--device", default=None)
    parser.add_argument("--dtype", default=None)
    parser.add_argument("--backend", default=None, help="auto, torch, vllm, openvino, mlx or llamacpp (default: auto)")
    parser.add_argument("--artifacts", default=None, help="the directory python -m decisionnode_edge.build wrote")
    parser.add_argument("--description", default=None, help="the description GET /v1/models gives")
    parser.add_argument("--release-date", default=None, help="the release date GET /v1/models gives (YYYY-MM-DD)")
    parser.add_argument("--box-calibration", default=None, help="the box calibration of /v1/systemone (default: the model folder's box_calibration.json when there is one; none for none)")
    parser.add_argument("--box-calibration-sha256", default=None, help="the sha256 the box calibration file must have")
    args = parser.parse_args(argv)
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s")
    api_key = None
    if args.api_key:
        with open(args.api_key, encoding="utf-8") as f:
            api_key = f.read().strip()
        if not api_key:
            raise SystemExit(f"{args.api_key} is empty")
    model = EdgeModel.load(args.model, device=args.device, dtype=args.dtype, backend=args.backend, artifacts=args.artifacts)
    calibration = box_calibration(model, args.box_calibration, args.box_calibration_sha256)
    model.decide({"state": "The color is red.", "questions": {"warmup": {"type": "choice", "instructions": "Which color is named?", "criteria": {"red": "red", "blue": "blue"}}}})
    server = make_server(model, args.host, args.port, api_key, description=args.description, release_date=args.release_date, calibration=calibration)
    log.info("serving %s (%s backend) on http://%s:%d%s", model.config["name"], model.backend, args.host, server.server_address[1], ", key required" if api_key else "")
    if calibration is not None:
        log.info("/v1/systemone box confidence: the box calibration %s (sha256 %s)", calibration.model or "of this model", calibration.sha256)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        server.server_close()
    return 0


if __name__ == "__main__":
    sys.exit(main())