life-echo/api/services/memoir_images/provider.py

import base64
import hmac
import logging
import os
import time
import uuid
from hashlib import sha1

import httpx

logger = logging.getLogger(__name__)

_SIZE_TO_ASPECT_RATIO = {
    "1024x1024": "square",
    "768x1024": "portrait",
    "1024x768": "landscape",
    "1280x720": "landscape",
    "720x1280": "portrait",
}


class LiblibImageProvider:
    """Liblib (https://openapi.liblibai.cloud) image generation adapter."""

    def __init__(
        self,
        http_client: httpx.Client | None = None,
        access_key: str | None = None,
        secret_key: str | None = None,
        base_url: str | None = None,
        template_uuid: str | None = None,
    ):
        self.http_client = http_client or httpx.Client(timeout=120)
        self.access_key = access_key or os.getenv("LIBLIB_ACCESS_KEY", "")
        self.secret_key = secret_key or os.getenv("LIBLIB_SECRET_KEY", "")
        self.base_url = (base_url or os.getenv("LIBLIB_BASE_URL", "https://openapi.liblibai.cloud")).rstrip("/")
        self.template_uuid = template_uuid or os.getenv(
            "LIBLIB_TEMPLATE_UUID", "5d7e67009b344550bc1aa6ccbfa1d7f4"
        )

    # ------------------------------------------------------------------
    # Signature helpers
    # ------------------------------------------------------------------

    def _sign(self, uri: str) -> str:
        """Build a full URL with Liblib HMAC-SHA1 query-string auth."""
        timestamp = str(int(time.time() * 1000))
        nonce = str(uuid.uuid4())
        content = "&".join((uri, timestamp, nonce))
        digest = hmac.new(self.secret_key.encode(), content.encode(), sha1).digest()
        signature = base64.urlsafe_b64encode(digest).rstrip(b"=").decode()
        return (
            f"{self.base_url}{uri}"
            f"?AccessKey={self.access_key}"
            f"&Signature={signature}"
            f"&Timestamp={timestamp}"
            f"&SignatureNonce={nonce}"
        )

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def submit_generation(self, prompt: str, size: str, style: str) -> dict:
        uri = "/api/generate/webui/text2img/ultra"
        url = self._sign(uri)

        aspect_ratio = _SIZE_TO_ASPECT_RATIO.get(size, "square")

        body = {
            "templateUuid": self.template_uuid,
            "generateParams": {
                "prompt": prompt,
                "aspectRatio": aspect_ratio,
                "imgCount": 1,
                "steps": 30,
            },
        }
        response = self.http_client.post(
            url,
            headers={"Content-Type": "application/json"},
            json=body,
        )
        response.raise_for_status()
        data = response.json()

        if data.get("code") != 0:
            raise RuntimeError(f"Liblib submit failed: {data.get('msg', data)}")

        generate_uuid = data["data"]["generateUuid"]
        return {"status": "processing", "job_id": generate_uuid, "image_url": None}

    def poll_until_complete(
        self, job: dict, poll_interval_seconds: int, max_attempts: int
    ) -> dict:
        uri = "/api/generate/webui/status"

        for attempt in range(max_attempts):
            url = self._sign(uri)
            response = self.http_client.post(
                url,
                headers={"Content-Type": "application/json"},
                json={"generateUuid": job["job_id"]},
            )
            response.raise_for_status()
            data = response.json()

            if data.get("code") != 0:
                raise RuntimeError(f"Liblib status query failed: {data.get('msg', data)}")

            result = data.get("data", {})
            status = result.get("generateStatus")

            if status == 5:  # success
                images = result.get("images") or []
                if images:
                    return {
                        "status": "completed",
                        "image_url": images[0]["imageUrl"],
                        "job_id": job["job_id"],
                    }
                raise RuntimeError(f"Liblib returned success but no images for {job['job_id']}")

            if status == 6:  # failed
                raise RuntimeError(f"Liblib generation failed: {result.get('generateMsg', 'unknown')}")

            if status == 7:  # timeout on Liblib side
                raise TimeoutError(f"Liblib generation timed out on server side: {job['job_id']}")

            logger.debug(
                "Liblib poll attempt %d/%d, status=%s, job=%s",
                attempt + 1, max_attempts, status, job["job_id"],
            )
            time.sleep(poll_interval_seconds)

        raise TimeoutError(f"Liblib image generation timed out after {max_attempts} attempts for {job['job_id']}")

    def download_image(self, job: dict) -> bytes:
        response = self.http_client.get(job["image_url"])
        response.raise_for_status()
        return response.content