`daily-cartoons/common.py`

90 lines
"""Shared helpers for the P24 daily-cartoon pipeline (TK-12153).
Cost guard reads ~/.claude/cost-ledger.jsonl (entries with app == APP) for today's
local date; logging goes through the cost-tracker skill's log.js. Never prints the token."""
import datetime, json, os, subprocess, sys, time, urllib.request

HERE = os.path.dirname(os.path.abspath(__file__))
SITE = os.environ.get("P24_SITE_DIR") or os.path.dirname(HERE)  # override only for tests
QUEUE = os.path.join(HERE, "queue")
APP = "p24-daily-cartoons"
LEDGER = os.path.expanduser("~/.claude/cost-ledger.jsonl")
LOGJS = os.path.expanduser("~/.claude/skills/cost-tracker/scripts/log.js")
# per-unit estimates (mirror cost-tracker pricing.json; WAN rate is a conservative assumption)
EST = {"flux": 0.003, "wan": 0.10}
API_KEYS = {"flux": ("replicate_flux_schnell", "image"), "wan": ("replicate_wan22_i2v_fast", "video")}

class CapExceeded(Exception):
    pass

def cap_usd():
    return float(os.environ.get("P24_CARTOON_CAP_USD", "1.00"))

def spent_today():
    today = datetime.date.today()
    total = 0.0
    try:
        with open(LEDGER) as f:
            for line in f:
                try:
                    e = json.loads(line)
                except Exception:
                    continue
                if e.get("app") != APP:
                    continue
                ts = e.get("ts", "")
                try:
                    d = datetime.datetime.fromisoformat(ts.replace("Z", "+00:00")).astimezone().date()
                except Exception:
                    continue
                if d == today:
                    total += float(e.get("cost_usd") or 0)
    except FileNotFoundError:
        pass
    return round(total, 6)

def guard(kind, n=1):
    """Raise CapExceeded unless today's spend + this call's estimate fits under the cap."""
    est = EST[kind] * n
    spent, cap = spent_today(), cap_usd()
    if spent + est > cap + 1e-9:
        raise CapExceeded(f"cap ${cap:.2f}: spent today ${spent:.3f} + est ${est:.3f} for {kind} would exceed it")
    print(f"  [cost] est ${est:.3f} for {kind} (spent today ${spent:.3f} / cap ${cap:.2f})")
    return est

def log_cost(kind, note, n=1):
    api, unit = API_KEYS[kind]
    out = subprocess.run(["node", LOGJS, "--api", api, "--units", f"{n}:{unit}", "--app", APP, "--note", note],
                         capture_output=True, text=True)
    if out.returncode != 0:
        print("  [cost] WARNING ledger log failed:", out.stderr.strip()[:200], file=sys.stderr)
    return EST[kind] * n

def token():
    for l in open(os.path.expanduser("~/Projects/secrets-manager/.env")):
        if l.startswith("REPLICATE_API_TOKEN="):
            return l.split("=", 1)[1].strip().strip("\"'")
    raise SystemExit("REPLICATE_API_TOKEN missing from secrets-manager/.env")

def api(url, data=None, prefer_wait=False):
    h = {"Authorization": f"Bearer {token()}", "Content-Type": "application/json"}
    if prefer_wait:
        h["Prefer"] = "wait"
    r = urllib.request.Request(url, data=json.dumps(data).encode() if data else None, headers=h)
    return json.load(urllib.request.urlopen(r, timeout=300))

def poll(pr, interval, deadline_s):
    # Replicate bills at create time: never lose the id, retry transient poll errors, bound the wait.
    end = time.time() + deadline_s
    while pr["status"] not in ("succeeded", "failed", "canceled"):
        if time.time() > end:
            raise RuntimeError(f"TIMEOUT after {deadline_s}s — recover later: GET {pr['urls']['get']}")
        time.sleep(interval)
        for a in range(5):
            try:
                pr = api(pr["urls"]["get"]); break
            except Exception as e:
                print(f"  poll error ({a+1}/5): {e}"); time.sleep(interval * (a + 1))
        else:
            raise RuntimeError(f"poll failed 5x — recover later: GET {pr['urls']['get']}")
    return pr