atlas-iac/services/hermes/scripts/batch_api.py

#!/usr/bin/env python3
"""Serve pinned batch models through a fixed, local-only Ollama endpoint."""

import json
import threading
import time
from urllib.error import HTTPError, URLError
from urllib.request import HTTPRedirectHandler, ProxyHandler, Request, build_opener


UPSTREAM = "http://ollama-batch.ai.svc.cluster.local:11434"
VERSION = "0.34.1"
PINS = {
    "qwen3.5:9b": "6488c96fa5faab64bb65cbd30d4289e20e6130ef535a93ef9a49f42eda893ea7",
    "qwen3.6:27b": "9d5803d493a991af27b9441c098aa56f2ed7bbd260877f075ec09b575c049bc3",
}
MAX_BODY = 1048576
TIMEOUT = 1800
PROTOCOL_VERSION = 2
BACKEND_API = "/api/chat"
_inference = threading.BoundedSemaphore(1)


class NoRedirect(HTTPRedirectHandler):
    """Never forward inference bodies or credentials through a redirect."""

    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


_http = build_opener(ProxyHandler({}), NoRedirect())


def _request(path, payload=None, timeout=10):
    """Read a bounded local response without consulting environment proxies."""
    data = None if payload is None else json.dumps(payload).encode()
    request = Request(UPSTREAM + path, data=data, headers={"Content-Type": "application/json"})
    with _http.open(request, timeout=timeout) as response:
        raw = response.read(4 * MAX_BODY + 1)
        if len(raw) > 4 * MAX_BODY:
            raise ValueError("upstream response too large")
        return json.loads(raw)


def catalog():
    """Fail closed if the runtime version or any approved model digest changes."""
    runtime = _request("/api/version")["version"]
    if runtime != VERSION:
        raise ValueError("runtime version mismatch")
    installed = {item["name"]: item for item in _request("/api/tags")["models"]}
    models = []
    for name, digest in PINS.items():
        item = installed.get(name, {})
        if item.get("digest", "").removeprefix("sha256:") != digest:
            raise ValueError("model digest mismatch or model unavailable")
        models.append({"model": name, "digest": digest, "size": item["size"],
                       "details": item.get("details", {})})
    return {"runtime": runtime, "placement": "titan-23/cpu", "models": models,
            "protocol_version": PROTOCOL_VERSION, "backend_api": BACKEND_API,
            "context_limits": [16384, 32768, 65536], "max_output_tokens": 16384,
            "timeout_seconds": TIMEOUT, "fallback": None}


def normalize(body):
    """Validate explicit inference settings and conservatively avoid truncation."""
    payload = json.loads(body)
    allowed = {"model", "prompt", "stream", "think", "format", "options"}
    if not isinstance(payload, dict) or set(payload) - allowed:
        raise ValueError("unsupported batch fields")
    if payload.get("model") not in PINS:
        raise ValueError("an exact approved batch model is required")
    if payload.get("stream") is not False or type(payload.get("think")) is not bool:
        raise ValueError("stream=false and explicit boolean think are required")
    prompt = payload.get("prompt")
    if not isinstance(prompt, str) or not prompt.strip():
        raise ValueError("nonempty text prompt is required")
    shape = payload.get("format")
    if shape != "json" and not isinstance(shape, dict):
        raise ValueError("JSON format or a JSON schema is required")
    if isinstance(shape, dict) and (shape.get("type") != "object" or len(json.dumps(shape)) > 65536):
        raise ValueError("bounded object JSON schema is required")
    options = payload.get("options")
    required = {"num_ctx", "num_predict", "temperature", "top_p", "top_k", "seed"}
    if not isinstance(options, dict) or set(options) != required:
        raise ValueError("explicit num_ctx, num_predict, temperature, top_p, top_k and seed are required")
    context, count = options["num_ctx"], options["num_predict"]
    if type(context) is not int or context not in (16384, 32768, 65536):
        raise ValueError("unsupported context size")
    if type(count) is not int or not 1 <= count <= 16384:
        raise ValueError("num_predict must be between 1 and 16384")
    # One UTF-8 byte per token is a deliberately conservative upper bound.
    # Reserve template overhead and the full output budget; never trim a suite.
    if len(prompt.encode("utf-8")) + count + 1024 > context:
        raise ValueError("prompt exceeds conservative context budget; use candidate batches and reconciliation")
    for key, upper in (("temperature", 2), ("top_p", 1)):
        value = options[key]
        if type(value) not in (int, float) or not 0 <= value <= upper:
            raise ValueError("invalid sampling value")
    for key, low, high in (("seed", 0, 2147483647), ("top_k", 1, 100)):
        if type(options[key]) is not int or not low <= options[key] <= high:
            raise ValueError("invalid integer sampling value")
    payload["options"] = {**options, "num_thread": 16, "num_gpu": 0}
    payload["keep_alive"] = "10m"
    return payload


def models_response():
    """Return verifiable configuration without leaking internal failure details."""
    try:
        return 200, catalog()
    except (OSError, HTTPError, URLError, ValueError, KeyError, TypeError):
        return 503, {"error": "pinned local batch models unavailable; no fallback"}


def generate(body):
    """Serialize a batch call, attaching measured wall time and exact provenance."""
    try:
        payload = normalize(body)
    except (ValueError, TypeError):
        return 400, {"error": "invalid batch request or context budget exceeded"}
    if not _inference.acquire(blocking=False):
        return 429, {"error": "local batch model busy; retry explicitly"}
    started = time.monotonic()
    try:
        provenance = catalog()
        # Ollama 0.34.1's generate API applies the JSON grammar inside thinking.
        # Chat defers that grammar until final content; do not expose thinking as
        # a substitute answer or silently disable the requested reasoning mode.
        upstream = {key: value for key, value in payload.items() if key != "prompt"}
        upstream["messages"] = [{"role": "user", "content": payload["prompt"]}]
        result = _request(BACKEND_API, upstream, timeout=TIMEOUT)
        message = result.pop("message", {})
        result["response"] = message.get("content", "")
        if result.get("model") != payload["model"] or not result.get("done"):
            raise ValueError("unexpected model or incomplete generation")
        if result.get("done_reason") == "length":
            return 422, {"error": "output budget exhausted; result is incomplete"}
        if not isinstance(result.get("response"), str):
            raise ValueError("missing output")
        try:
            json.loads(result["response"])
        except ValueError:
            return 422, {"error": "local model returned invalid structured output",
                         "wall_seconds": round(time.monotonic() - started, 3)}
        # Keep the requested rationale in the answer, not raw hidden thinking.
        result.pop("thinking", None)
        result.pop("context", None)
        result["batch_provenance"] = {
            "model_digest": PINS[payload["model"]], "runtime": provenance["runtime"],
            "placement": provenance["placement"], "options": payload["options"],
            "think": payload["think"], "wall_seconds": round(time.monotonic() - started, 3),
            "protocol_version": PROTOCOL_VERSION, "backend_api": BACKEND_API,
        }
        return 200, result
    except (OSError, HTTPError, URLError, ValueError, KeyError, TypeError):
        return 503, {"error": "pinned local batch inference failed; no fallback"}
    finally:
        _inference.release()