Company_Sprite_Gen/pixellab_api.py

"""
Minimal PixelLab v2 client used by pixellab_library.py.

- Token from PIXELLAB_API_KEY (environment or .env next to this file).
- Retries rate limits (429/529) with backoff, polls async jobs (edit endpoints return a job id),
  and caps concurrent jobs across ALL processes at the account limit (PIXELLAB_CONCURRENCY,
  default 10) with one file lock per slot, so parallel workers and agents queue instead of 429.
- Appends one line per paid call to pixellab_src/ledger.jsonl so spend can be tracked across
  runs and parallel workers: `python pixellab_library.py cost`.
"""
import base64
import fcntl
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"
LEDGER = Path(__file__).with_name("pixellab_src") / "ledger.jsonl"
LOCK_DIR = LEDGER.with_name(".api_slots")
CONCURRENCY = int(os.environ.get("PIXELLAB_CONCURRENCY", "10"))
RETRIES = 10
JOB_POLL_SECONDS = 4
JOB_TIMEOUT_SECONDS = 900


class ApiError(Exception):
    pass


def load_key():
    key = os.environ.get("PIXELLAB_API_KEY")
    env = Path(__file__).with_name(".env")
    if not key and env.exists():
        for line in env.read_text().splitlines():
            name, sep, value = line.strip().partition("=")
            if sep and name.strip() == "PIXELLAB_API_KEY":
                key = value.strip().strip("'\"")
    if not key:
        sys.exit("PIXELLAB_API_KEY not set (environment or .env)")
    return key


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


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


def find_image(obj):
    """First image in a response: a Base64Image dict or a data: URL, at any depth."""
    if isinstance(obj, dict):
        if isinstance(obj.get("base64"), str):
            return decode(obj)
        obj = list(obj.values())
    if isinstance(obj, list):
        for v in obj:
            img = find_image(v)
            if img is not None:
                return img
    if isinstance(obj, str) and obj.startswith("data:image"):
        return decode({"base64": obj})
    return None


def redact(payload):
    return {k: ("<image>" if isinstance(v, dict) and v.get("type") == "base64" else v) for k, v in payload.items()}


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(RETRIES):
            try:
                r = self.session.request(method, f"{API}/{path}", json=payload, timeout=300)
            except requests.RequestException as e:
                if attempt == RETRIES - 1:
                    raise ApiError(str(e))
                time.sleep(5 + attempt * 5)
                continue
            if r.status_code in (429, 529, 502, 503) and attempt < RETRIES - 1:
                time.sleep(min(60, 6 * (attempt + 1)) + random.uniform(0, 4))
                continue
            if not r.ok:
                raise ApiError(f"{path}: {r.status_code} {r.text[:400]}")
            return r.json()

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

    def run(self, endpoint, payload, label):
        """POST, wait for the job if there is one, return (image, usage). Logs the call to the ledger.

        Holds one of CONCURRENCY slot locks for the whole job (POST + polling)."""
        LOCK_DIR.mkdir(parents=True, exist_ok=True)
        while True:
            for i in random.sample(range(CONCURRENCY), CONCURRENCY):
                f = (LOCK_DIR / f"slot{i}").open("a")
                try:
                    fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
                except BlockingIOError:
                    f.close()
                    continue
                try:
                    return self._run(endpoint, payload, label)
                finally:
                    fcntl.flock(f, fcntl.LOCK_UN)
                    f.close()
            time.sleep(1 + random.random())

    def _run(self, endpoint, payload, label):
        resp = self.call("POST", endpoint, payload)
        usage = resp.get("usage")
        job = resp.get("background_job_id")
        deadline = time.time() + JOB_TIMEOUT_SECONDS
        while job:
            if time.time() > deadline:
                raise ApiError(f"job {job} still running after {JOB_TIMEOUT_SECONDS}s")
            time.sleep(JOB_POLL_SECONDS)
            status = self.call("GET", f"background-jobs/{job}")
            if status["status"] == "failed":
                raise ApiError(f"job {job} failed: {json.dumps(status.get('last_response'))[:400]}")
            if status["status"] == "completed":
                resp, usage, job = status.get("last_response") or {}, status.get("usage") or usage, None
        img = find_image({k: v for k, v in resp.items() if k != "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, "ok": img is not None}) + "\n")
        if img is None:
            raise ApiError(f"no image in response: {json.dumps(redact(resp))[:400]}")
        return img, usage