pigway-frame-flow/pigway_frame_flow/extensions/ai_images.py

"""Optional one-image generation, using the durable operation queue."""

import base64
import json
from io import BytesIO
from typing import Literal
from urllib.parse import urlsplit

from fastapi import Depends, Header, Request
from PIL import Image, ImageOps
from pydantic import Field

from pigway_frame_flow.core.auth import require
from pigway_frame_flow.errors import ErrorCode as E
from pigway_frame_flow.errors import FlowError
from pigway_frame_flow.extensions.api import AISecret
from pigway_frame_flow.extensions.secrets import save_secret
from pigway_frame_flow.models import Accepted, StrictModel
from pigway_frame_flow.renderers.images import MAX_IMAGE_PIXELS
from pigway_frame_flow.storage.database import encode


class ImageGenerationConfig(StrictModel):
    endpoint: str = Field(max_length=2048)
    model: str = Field(min_length=1, max_length=120)
    size: Literal["1024x1024", "1024x1536", "1536x1024"] = "1024x1536"


class ImageGenerationInput(StrictModel):
    prompt: str = Field(min_length=1, max_length=4000)


def install_ai_images(app, service):
    def config():
        row = service.db.one("SELECT payload FROM app_config WHERE key='ai_images'")
        if not row:
            raise FlowError(E.CONFIG_REQUIRED, "Configure image generation first", 409)
        return ImageGenerationConfig.model_validate_json(row["payload"])

    @app.get(
        "/api/v1/settings/ai-images",
        dependencies=[Depends(require("admin"))],
        tags=["AI"],
    )
    def get_config():
        row = service.db.one("SELECT payload FROM app_config WHERE key='ai_images'")
        return {
            "config": json.loads(row["payload"]) if row else None,
            "key_configured": (
                service.settings.config_dir / "ai-images.secret"
            ).exists(),
        }

    @app.put(
        "/api/v1/settings/ai-images",
        dependencies=[Depends(require("admin"))],
        tags=["AI"],
    )
    def set_config(payload: ImageGenerationConfig):
        url = urlsplit(payload.endpoint)
        if (
            url.scheme != "https"
            or not url.hostname
            or url.username
            or url.password
            or url.query
            or url.fragment
            or url.port not in (None, 443)
        ):
            raise FlowError(
                E.INVALID_REQUEST,
                "Image generation requires a plain public HTTPS endpoint",
            )
        if url.hostname not in service.settings.image_url_hosts:
            service.settings.image_url_hosts.append(url.hostname)
            service.db.execute(
                "INSERT OR REPLACE INTO app_config VALUES('image_url_hosts',?)",
                (encode(service.settings.image_url_hosts),),
            )
        service.db.execute(
            "INSERT OR REPLACE INTO app_config VALUES('ai_images',?)",
            (payload.model_dump_json(),),
        )
        return payload

    @app.put(
        "/api/v1/settings/ai-images/key",
        status_code=204,
        dependencies=[Depends(require("admin"))],
        tags=["AI"],
    )
    def set_key(payload: AISecret):
        key = payload.key.get_secret_value()
        save_secret(service.settings.config_dir / "ai-images.secret", key)
        service.settings.additional_secrets.append(key)

    @app.post(
        "/api/v1/ai/images",
        status_code=202,
        response_model=Accepted,
        dependencies=[Depends(require("render:write"))],
        tags=["AI"],
    )
    def generate(
        payload: ImageGenerationInput,
        request: Request,
        idempotency_key: str | None = Header(default=None),
    ):
        config()
        if not (service.settings.config_dir / "ai-images.secret").exists():
            raise FlowError(
                E.CONFIG_REQUIRED, "Configure image generation key first", 409
            )
        id = service.submit(
            "render-ai",
            payload.model_dump(),
            service.settings.device_id,
            request.state.client_id,
            request.state.request_id,
            idempotency_key,
        )
        return Accepted(operation_id=id, job_id=id, status="accepted")

    def execute(op, payload):
        service.update(op["id"], "RENDERING")
        conf = config()
        key = (service.settings.config_dir / "ai-images.secret").read_text().strip()
        service.settings.additional_secrets.append(key)
        service.db.execute(
            "UPDATE operations SET write_started=1 WHERE id=?", (op["id"],)
        )
        try:
            raw = service.assets.fetch_url(
                conf.endpoint,
                method="POST",
                headers={
                    "Authorization": "Bearer " + key,
                    "Content-Type": "application/json",
                },
                body=json.dumps(
                    {
                        "model": conf.model,
                        "prompt": payload["prompt"],
                        "size": conf.size,
                        "n": 1,
                    }
                ).encode(),
                limit=service.settings.max_bytes,
            )
            result = json.loads(raw)["data"]
            if len(result) != 1:
                raise ValueError()
            if result[0].get("b64_json"):
                image_data = base64.b64decode(result[0]["b64_json"], validate=True)
            else:
                image_data = service.assets.fetch_url(result[0]["url"])
            with Image.open(BytesIO(image_data)) as image:
                if image.width * image.height > MAX_IMAGE_PIXELS:
                    raise ValueError()
                rgba = ImageOps.exif_transpose(image).convert("RGBA")
                rgb = Image.new("RGB", rgba.size, "white")
                rgb.paste(rgba, mask=rgba.getchannel("A"))
                output = BytesIO()
                rgb.save(output, "JPEG", quality=93)
            asset = service.assets.put(output.getvalue(), "ai-art.jpg", "generated")
        except Exception:
            raise FlowError(
                E.RENDER_FAILED,
                "Image generation failed; request may have been billed. No automatic retry was sent.",
                502,
            ) from None
        service.update(op["id"], "COMPLETED", asset)
        service.db.history(op, "COMPLETED")
        service.db.event("operation.completed", {"operation_id": op["id"]})

    service.operation_handlers["render-ai"] = execute