brandling-safari/tools/pixellab_client.py

"""Minimal PixelLab v2 client for the Brandling Safari art tools.

- Token from PIXELLAB_API_KEY (environment, or a .env in this repo or its parent folder).
- Retries rate limits / transient errors, polls background jobs.
- Appends one line per paid call to art_src/pixellab_ledger.jsonl (`python tools/pixellab_client.py`
  prints the spend so far and the account balance).
"""
import base64
import io
import json
import os
import random
import sys
import time
from pathlib import Path

import requests
from PIL import Image

API = "https://api.pixellab.ai/v2"
ROOT = Path(__file__).resolve().parent.parent
LEDGER = ROOT / "art_src" / "pixellab_ledger.jsonl"
PENDING = ROOT / "art_src" / ".pending"
POLL_SECONDS = 5
JOB_TIMEOUT_SECONDS = 1200


class ApiError(Exception):
    pass


def load_key(name="PIXELLAB_API_KEY"):
    if os.environ.get(name):
        return os.environ[name]
    for env in (ROOT / ".env", ROOT.parent / ".env"):
        if env.exists():
            for line in env.read_text().splitlines():
                k, sep, v = line.strip().partition("=")
                if sep and k.strip() == name:
                    return v.strip().strip("'\"")
    sys.exit(f"{name} not set (environment, {ROOT / '.env'} or {ROOT.parent / '.env'})")


def b64(img: Image.Image) -> dict:
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return {"type": "base64", "base64": base64.b64encode(buf.getvalue()).decode(), "format": "png"}


def decode(field) -> Image.Image:
    data = field["base64"] if isinstance(field, dict) else field
    if data.startswith("data:"):
        data = data.split(",", 1)[1]
    return Image.open(io.BytesIO(base64.b64decode(data))).convert("RGBA")


def images_in(obj):
    """Every image in a response, depth-first in document order."""
    out = []
    if isinstance(obj, dict):
        if obj.get("type") == "base64" and "base64" in obj:
            return [decode(obj)]
        if "base64" in obj and isinstance(obj["base64"], str) and len(obj["base64"]) > 64:
            return [decode(obj)]
        for v in obj.values():
            out += images_in(v)
    elif isinstance(obj, list):
        for v in obj:
            out += images_in(v)
    elif isinstance(obj, str) and obj.startswith("data:image"):
        out.append(decode(obj))
    return out


def redact(obj):
    if isinstance(obj, dict):
        if "base64" in obj:
            return "<image>"
        return {k: redact(v) for k, v in obj.items()}
    if isinstance(obj, list):
        return [redact(v) for v in obj]
    return obj


class Client:
    def __init__(self):
        self.session = requests.Session()
        self.session.headers["Authorization"] = f"Bearer {load_key()}"

    def call(self, method, path, payload=None):
        for attempt in range(8):
            try:
                r = self.session.request(method, f"{API}/{path}", json=payload, timeout=300)
            except requests.RequestException as e:
                if attempt == 7:
                    raise ApiError(str(e))
                time.sleep(5 + attempt * 5)
                continue
            if r.status_code in (429, 502, 503, 529) and attempt < 7:
                time.sleep(min(60, 6 * (attempt + 1)) + random.uniform(0, 3))
                continue
            if not r.ok:
                raise ApiError(f"{method} {path}: {r.status_code} {r.text[:600]}")
            return r.json() if r.content else {}

    def balance(self):
        return self.call("GET", "balance")

    def wait_job(self, job_id):
        deadline = time.time() + JOB_TIMEOUT_SECONDS
        while True:
            if time.time() > deadline:
                raise ApiError(f"job {job_id} still running after {JOB_TIMEOUT_SECONDS}s")
            status = self.call("GET", f"background-jobs/{job_id}")
            if status.get("status") == "failed":
                raise ApiError(f"job {job_id} failed: {json.dumps(redact(status))[:600]}")
            if status.get("status") == "completed":
                return status
            time.sleep(POLL_SECONDS)

    def run(self, endpoint, payload, label):
        """POST and wait for every background job it starts. Returns (first response, job results).

        The job ids are written to art_src/.pending/<label>.json before waiting, so if this process
        is interrupted the next run with the same label picks the paid-for jobs up instead of
        starting (and paying for) new ones."""
        pending = PENDING / (label.replace("/", "__") + ".json")
        resp = None
        if pending.exists():
            saved = json.loads(pending.read_text())
            if saved.get("endpoint") == endpoint:
                resp = saved["response"]
                print(f"[resume] {label}: waiting for {len(saved.get('jobs', []))} earlier job(s)", flush=True)
        if resp is None:
            resp = self.call("POST", endpoint, payload)
        jobs = list(resp.get("background_job_ids") or [])
        if resp.get("background_job_id"):
            jobs.append(resp["background_job_id"])
        if jobs:
            PENDING.mkdir(parents=True, exist_ok=True)
            pending.write_text(json.dumps({"endpoint": endpoint, "jobs": jobs, "response": redact(resp)}))
        results, usages = [resp], [resp.get("usage")]
        for job in jobs:
            st = self.wait_job(job)
            results.append(st.get("last_response") or {})
            usages.append(st.get("usage"))
        self.log(label, endpoint, [u for u in usages if u])
        pending.unlink(missing_ok=True)
        return resp, results[1:] if jobs else results

    def log(self, label, endpoint, usage):
        LEDGER.parent.mkdir(parents=True, exist_ok=True)
        with LEDGER.open("a") as f:
            f.write(json.dumps({"time": round(time.time()), "label": label, "endpoint": endpoint, "usage": usage}) + "\n")


def spend():
    total, by = 0.0, {}
    if LEDGER.exists():
        for line in LEDGER.read_text().splitlines():
            d = json.loads(line)
            g = sum(float(u.get("generations", 0) or 0) for u in d["usage"] if isinstance(u, dict))
            total += g
            by[d["endpoint"]] = by.get(d["endpoint"], 0.0) + g
    return total, by


if __name__ == "__main__":
    total, by = spend()
    print(f"ledger: {total:.0f} generations", json.dumps(by))
    print(json.dumps(Client().balance()))