src/shared/image_client.py

"""Shared OpenRouter image-generation adapter (dedicated /api/v1/images API).

Same safety posture as ai_client: secrets only from env, bounded retries for
transient failures, telemetry without prompt content, user-safe errors.
LOREIFY_FAKE_AI=1 returns an embedded placeholder image (local testing only).
"""

import base64
import logging
import os
import time

import httpx

from .ai_client import AIUnavailable, USER_ERROR, NO_KEY_ERROR

log = logging.getLogger("lf.image")

# 8x8 solid-gold PNG placeholder for LOREIFY_FAKE_AI local pipeline tests.
_FAKE_PNG = base64.b64decode(
    "iVBORw0KGgoAAAANSUhEUgAAAAgAAAAICAYAAADED76LAAAAFklEQVR4nGP8v4rhPwMRgIkY"
    "RaMKKVMIAJVzAx2Vr/xnAAAAAElFTkSuQmCC"
)


def _cfg():
    return {
        "api_key": os.environ.get("OPENROUTER_API_KEY", ""),
        "base_url": os.environ.get("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1"),
        "model": os.environ.get("OPENROUTER_IMAGE_MODEL", "krea/krea-2-large"),
        "timeout": float(os.environ.get("OPENROUTER_IMAGE_TIMEOUT_SECONDS", "120")),
    }


def generate_image(prompt: str) -> tuple[bytes, str]:
    """Generate one image. Returns (image_bytes, media_type).

    Raises AIUnavailable with a user-safe message on any failure.
    """
    if os.environ.get("LOREIFY_FAKE_AI") == "1":
        time.sleep(0.4)
        return _FAKE_PNG, "image/png"

    cfg = _cfg()
    if not cfg["api_key"]:
        raise AIUnavailable(NO_KEY_ERROR)

    headers = {
        "Authorization": f"Bearer {cfg['api_key']}",
        "HTTP-Referer": "https://loreify.ai",
        "X-Title": "Loreify D&D Tools",
        "Content-Type": "application/json",
    }
    # krea models advertise no extra supported_parameters — keep the body minimal.
    body = {"model": cfg["model"], "prompt": prompt}

    delay = 2.0
    last_exc = None
    for attempt in range(3):
        start = time.time()
        try:
            resp = httpx.post(f"{cfg['base_url']}/images", headers=headers,
                              json=body, timeout=cfg["timeout"])
        except (httpx.TimeoutException, httpx.TransportError) as exc:
            last_exc = exc
            log.info("image call attempt=%d error_class=%s", attempt, type(exc).__name__)
            time.sleep(delay)
            delay *= 2
            continue
        latency = time.time() - start
        if resp.status_code == 429 or resp.status_code >= 500:
            last_exc = AIUnavailable(USER_ERROR)
            retry_after = resp.headers.get("retry-after")
            wait = min(float(retry_after), 20.0) if retry_after and retry_after.isdigit() else delay
            log.info("image call attempt=%d status=%d retry_in=%.1f", attempt, resp.status_code, wait)
            time.sleep(wait)
            delay *= 2
            continue
        if resp.status_code != 200:
            log.info("image call status=%d error_class=http", resp.status_code)
            raise AIUnavailable(USER_ERROR)
        payload = resp.json()
        data = payload.get("data") or []
        if not data or not data[0].get("b64_json"):
            log.info("image call ok but empty data")
            raise AIUnavailable(USER_ERROR)
        usage = payload.get("usage") or {}
        log.info(
            "image ok provider=openrouter model=%s latency=%.2fs tokens=%s cost=%s",
            cfg["model"], latency, usage.get("total_tokens"), usage.get("cost"),
        )
        media_type = data[0].get("media_type") or "image/png"
        try:
            return base64.b64decode(data[0]["b64_json"]), media_type
        except Exception:
            raise AIUnavailable(USER_ERROR)
    raise last_exc if isinstance(last_exc, AIUnavailable) else AIUnavailable(USER_ERROR)