modelbeast/mb

#!/usr/bin/env python3
"""mb — headless CLI for MODELBEAST. Zero dependencies (python3 stdlib only),
works from any machine that can reach the server (tailnet included).

  export MB_HOST=http://100.89.131.57:8777   # default: http://localhost:8777

  mb ops [-v]                       list operators (+params with -v)
  mb assets [--kind image]          list assets
  mb upload FILE [FILE...]          upload files, print asset ids
  mb run OP [--asset ID | --file F] [-p k=v ...] [--params-json J]
            [--wait] [--follow] [--download DIR]
  mb job ID                         job status json
  mb jobs [-n 20]                   recent jobs
  mb log ID [-f]                    print job log (-f = follow to completion)
  mb wait ID [--download DIR]       block until job finishes (exit 0/2/3)
  mb get ASSET_ID [-o PATH]         download an asset (folders: all members)
  mb rm-asset ID / mb rm-job ID     delete
  mb retry ID / mb cancel ID        job control
  mb settings                       show settings (secrets masked)
  mb set KEY VALUE                  set a setting (e.g. keys — value visible in
                                    shell history; prefer the web UI for secrets)

Typical: mb run flux_local -p prompt="a brass astrolabe" -p seed=7 --wait --download out/
Exit codes: 0 ok/done, 1 usage/HTTP error, 2 job error, 3 job cancelled.
"""
import argparse
import json
import mimetypes
import os
import sys
import time
import urllib.error
import urllib.request
import uuid
from pathlib import Path

HOST = os.environ.get("MB_HOST", "http://localhost:8777").rstrip("/")


def http(method, path, body=None):
    req = urllib.request.Request(HOST + path, method=method)
    data = None
    if body is not None:
        data = json.dumps(body).encode()
        req.add_header("Content-Type", "application/json")
    try:
        with urllib.request.urlopen(req, data=data, timeout=120) as r:
            ct = r.headers.get("content-type", "")
            raw = r.read()
            return json.loads(raw) if "json" in ct else raw
    except urllib.error.HTTPError as e:
        sys.exit(f"HTTP {e.code} {method} {path}: {e.read().decode()[:500]}")
    except urllib.error.URLError as e:
        sys.exit(f"cannot reach {HOST} ({e.reason}) — is the server up? set MB_HOST?")


def upload_file(path):
    p = Path(path)
    if not p.is_file():
        sys.exit(f"no such file: {p}")
    boundary = uuid.uuid4().hex
    ctype = mimetypes.guess_type(p.name)[0] or "application/octet-stream"
    body = (f"--{boundary}\r\nContent-Disposition: form-data; name=\"file\"; "
            f"filename=\"{p.name}\"\r\nContent-Type: {ctype}\r\n\r\n").encode()
    body += p.read_bytes() + f"\r\n--{boundary}--\r\n".encode()
    req = urllib.request.Request(HOST + "/api/assets", data=body, method="POST")
    req.add_header("Content-Type", f"multipart/form-data; boundary={boundary}")
    try:
        with urllib.request.urlopen(req, timeout=600) as r:
            return json.loads(r.read())
    except urllib.error.HTTPError as e:
        sys.exit(f"upload failed: HTTP {e.code}: {e.read().decode()[:300]}")


def coerce(op, key, value):
    """Coerce k=v strings using the operator's JSON schema, else best-effort JSON."""
    prop = (op.get("params_schema", {}).get("properties", {}) or {}).get(key, {})
    t = prop.get("type")
    try:
        if t == "integer":
            return int(value)
        if t == "number":
            return float(value)
        if t == "boolean":
            return value.lower() in ("1", "true", "yes", "on")
        if t == "string":
            return value
        return json.loads(value)
    except (ValueError, json.JSONDecodeError):
        return value


def outputs_of(job_id):
    return [a for a in http("GET", "/api/assets") if a.get("parent_job") == job_id]


def download_asset(asset, dest_dir):
    dest_dir = Path(dest_dir)
    dest_dir.mkdir(parents=True, exist_ok=True)
    url = f"{HOST}/api/assets/{asset['id']}/file"
    with urllib.request.urlopen(url, timeout=600) as r:
        ct = r.headers.get("content-type", "")
        data = r.read()
    if "json" in ct:
        listing = json.loads(data)
        if listing.get("folder"):
            sub = dest_dir / asset["name"]
            sub.mkdir(exist_ok=True)
            for member in listing["files"]:
                murl = url + "?member=" + urllib.request.quote(member)
                with urllib.request.urlopen(murl, timeout=600) as r:
                    (sub / member).write_bytes(r.read())
            print(f"  ↓ {sub} ({len(listing['files'])} files)")
            return
    out = dest_dir / asset["name"]
    out.write_bytes(data)
    print(f"  ↓ {out} ({len(data)} bytes)")


def wait_for(job_id, follow=False, download=None):
    printed = 0
    last_status = None
    while True:
        j = http("GET", f"/api/jobs/{job_id}")
        if follow:
            log = j.get("log") or ""
            if len(log) > printed:
                sys.stdout.write(log[printed:])
                sys.stdout.flush()
                printed = len(log)
        elif j["status"] != last_status:
            print(f"[{job_id}] {j['status']}")
            last_status = j["status"]
        if j["status"] in ("done", "error", "cancelled"):
            if j["status"] == "error":
                print(f"ERROR: {j.get('error')}\n--- last log ---\n{(j.get('log') or '')[-1500:]}",
                      file=sys.stderr)
                sys.exit(2)
            if j["status"] == "cancelled":
                sys.exit(3)
            outs = outputs_of(job_id)
            print(f"done — {len(outs)} output(s):")
            for a in outs:
                print(f"  {a['id']}  {a['kind']:<14} {a['name']}")
            if download:
                for a in outs:
                    download_asset(a, download)
            return
        time.sleep(2)


def main():
    ap = argparse.ArgumentParser(prog="mb", add_help=True)
    sub = ap.add_subparsers(dest="cmd", required=True)

    sub.add_parser("ops").add_argument("-v", action="store_true")
    p = sub.add_parser("assets"); p.add_argument("--kind")
    p = sub.add_parser("upload"); p.add_argument("files", nargs="+")
    p = sub.add_parser("run")
    p.add_argument("operator")
    p.add_argument("--asset"); p.add_argument("--file")
    p.add_argument("-p", "--param", action="append", default=[])
    p.add_argument("--params-json")
    p.add_argument("--wait", action="store_true")
    p.add_argument("--follow", action="store_true")
    p.add_argument("--download")
    p = sub.add_parser("job"); p.add_argument("id")
    p = sub.add_parser("jobs"); p.add_argument("-n", type=int, default=20)
    p = sub.add_parser("log"); p.add_argument("id"); p.add_argument("-f", action="store_true")
    p = sub.add_parser("wait"); p.add_argument("id"); p.add_argument("--download")
    p = sub.add_parser("get"); p.add_argument("id"); p.add_argument("-o", default=".")
    p = sub.add_parser("rm-asset"); p.add_argument("id")
    p = sub.add_parser("rm-job"); p.add_argument("id")
    p = sub.add_parser("retry"); p.add_argument("id")
    p = sub.add_parser("cancel"); p.add_argument("id")
    sub.add_parser("settings")
    p = sub.add_parser("set"); p.add_argument("key"); p.add_argument("value")
    a = ap.parse_args()

    if a.cmd == "ops":
        for op in http("GET", "/api/operators"):
            gate = f"  [needs {','.join(op['requires_env'])}]" if op.get("requires_env") else ""
            inputs = ",".join(op["accepts"]) or "none (prompt-driven)"
            print(f"{op['id']:<20} {op.get('category','?'):<11} in:{inputs:<18} "
                  f"lane:{op.get('resources','cpu')}{gate}")
            if a.v:
                for k, d in (op.get("params_schema", {}).get("properties", {}) or {}).items():
                    dv = f" (default {d['default']})" if "default" in d else ""
                    en = f" one of {d['enum']}" if "enum" in d else ""
                    print(f"    -p {k}=<{d.get('type','any')}>{en}{dv}  {d.get('description','')}")
    elif a.cmd == "assets":
        for x in http("GET", "/api/assets"):
            if a.kind and x["kind"] != a.kind:
                continue
            print(f"{x['id']}  {x['kind']:<14} {x['size']:>10}  {x['name']}")
    elif a.cmd == "upload":
        for f in a.files:
            asset = upload_file(f)
            print(f"{asset['id']}  {asset['kind']}  {asset['name']}")
    elif a.cmd == "run":
        ops = {o["id"]: o for o in http("GET", "/api/operators")}
        if a.operator not in ops:
            sys.exit(f"unknown operator '{a.operator}'. Run: mb ops")
        op = ops[a.operator]
        asset_id = a.asset
        if a.file:
            asset_id = upload_file(a.file)["id"]
            print(f"uploaded → {asset_id}")
        if op["accepts"] and not asset_id:
            sys.exit(f"{a.operator} needs an input ({','.join(op['accepts'])}): --asset ID or --file F")
        params = {}
        if a.params_json:
            params.update(json.loads(a.params_json))
        for kv in a.param:
            if "=" not in kv:
                sys.exit(f"bad -p '{kv}' (want k=v)")
            k, v = kv.split("=", 1)
            params[k] = coerce(op, k, v)
        job = http("POST", "/api/jobs",
                   {"operator": a.operator, "asset_id": asset_id, "params": params})
        print(f"job {job['id']} queued")
        if a.wait or a.follow or a.download:
            wait_for(job["id"], follow=a.follow, download=a.download)
    elif a.cmd == "job":
        print(json.dumps(http("GET", f"/api/jobs/{a.id}"), indent=2))
    elif a.cmd == "jobs":
        for j in http("GET", "/api/jobs")[: a.n]:
            print(f"{j['id']}  {j['status']:<9} {j['operator']}")
    elif a.cmd == "log":
        if a.f:
            wait_for(a.id, follow=True)
        else:
            print(http("GET", f"/api/jobs/{a.id}").get("log") or "(empty)")
    elif a.cmd == "wait":
        wait_for(a.id, download=a.download)
    elif a.cmd == "get":
        assets = {x["id"]: x for x in http("GET", "/api/assets")}
        if a.id not in assets:
            sys.exit("no such asset")
        download_asset(assets[a.id], a.o)
    elif a.cmd == "rm-asset":
        http("DELETE", f"/api/assets/{a.id}"); print("deleted")
    elif a.cmd == "rm-job":
        http("DELETE", f"/api/jobs/{a.id}"); print("deleted")
    elif a.cmd == "retry":
        j = http("POST", f"/api/jobs/{a.id}/retry"); print(f"new job {j['id']}")
    elif a.cmd == "cancel":
        http("POST", f"/api/jobs/{a.id}/cancel"); print("cancel requested")
    elif a.cmd == "settings":
        print(json.dumps(http("GET", "/api/settings"), indent=2))
    elif a.cmd == "set":
        http("PUT", "/api/settings", {a.key: a.value}); print("saved")


if __name__ == "__main__":
    main()