atlas-iac/services/hermes/plugins/image-gen-broker/__init__.py

"""Hermes image backend that keeps owner OAuth outside family chat pods."""

from __future__ import annotations

import base64
import json
import os
from pathlib import Path
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

from agent.image_gen_provider import (
    DEFAULT_ASPECT_RATIO,
    ImageGenProvider,
    error_response,
    normalize_reference_images,
    resolve_aspect_ratio,
    save_b64_image,
    success_response,
)


BROKER_URL = os.environ.get(
    "HERMES_IMAGE_BROKER_URL",
    "http://hermes-image-broker.hermes.svc.cluster.local:9002",
).rstrip("/")
DEFAULT_MODEL = "atlas-image-auto-high"
ROUTE_MODELS = {
    "auto": "atlas-image-auto-high",
    "hosted": "gpt-image-2-high",
    "local": "flux-2-klein-4b-local",
}
MODELS = {
    "atlas-image-auto-high": (
        "Auto (GPT Image quality, local fallback)",
        "auto-high",
        "Hosted quality first; FLUX fallback when hosted generation fails",
    ),
    "gpt-image-2-low": ("GPT Image 2 (Fast)", "low"),
    "gpt-image-2-medium": ("GPT Image 2 (Balanced)", "medium"),
    "gpt-image-2-high": ("GPT Image 2 (Highest quality)", "high"),
    "flux-2-klein-4b-local": (
        "Local FLUX 2 Klein 4B",
        "local-high",
        "Private local generation and editing on the shared RTX 3080",
    ),
}
MAX_INPUT_BYTES = 25 << 20

IMAGE_GENERATE_PARAMETERS = {
    "type": "object",
    "properties": {
        "prompt": {
            "type": "string",
            "description": "Detailed description of the image or requested edit.",
        },
        "aspect_ratio": {
            "type": "string",
            "enum": ["landscape", "square", "portrait"],
            "default": DEFAULT_ASPECT_RATIO,
            "description": "Requested output aspect ratio.",
        },
        "image_url": {
            "type": "string",
            "description": "Optional uploaded image URL or private local path to edit.",
        },
        "reference_image_urls": {
            "type": "array",
            "items": {"type": "string"},
            "description": "Optional reference images for the generation or edit.",
        },
    },
    "required": ["prompt"],
}

IMAGE_EDIT_PARAMETERS = {
    "type": "object",
    "properties": {
        "prompt": {
            "type": "string",
            "description": "Detailed description of the requested edit.",
        },
        "aspect_ratio": {
            "type": "string",
            "enum": ["landscape", "square", "portrait"],
            "default": DEFAULT_ASPECT_RATIO,
            "description": "Requested output aspect ratio.",
        },
    },
    "required": ["prompt"],
}

LOCAL_IMAGE_SCHEMA = {
    "name": "image_generate_local",
    "description": (
        "Generate or edit an image only with the private local FLUX model on "
        "Atlas hardware. Use this tool whenever the user says local, private, "
        "FLUX, on my hardware, or otherwise explicitly rejects a hosted image "
        "provider. Never substitute the generic image_generate tool for an "
        "explicit local request. The backend is already provisioned; never ask "
        "for an endpoint or tell the user to install Diffusers or ComfyUI. "
        "For an edit follow-up, pass the newest MEDIA: path from the conversation "
        "as image_url, even when the user refers to it only as this, it, the "
        "image, or its pictured subject."
    ),
    "parameters": IMAGE_GENERATE_PARAMETERS,
}

HOSTED_IMAGE_SCHEMA = {
    "name": "image_generate_hosted",
    "description": (
        "Generate or edit an image only with hosted OpenAI GPT Image at the "
        "highest configured quality. Use this tool when the user explicitly "
        "asks for OpenAI, GPT Image, or hosted image generation. Do not use it "
        "when the user explicitly requests local or private generation. For an "
        "edit follow-up, pass the newest MEDIA: path from the conversation as "
        "image_url, even when the user refers to it only as this, it, the image, "
        "or its pictured subject."
    ),
    "parameters": IMAGE_GENERATE_PARAMETERS,
}

AUTO_EDIT_SCHEMA = {
    "name": "image_edit_latest",
    "description": (
        "Edit the newest generated image in this private conversation using "
        "the automatic image route. Use this compact tool for natural "
        "follow-ups such as turn this cat into a clown, change it, edit the "
        "image, or make the pictured subject different when the user does not "
        "name a provider. The server resolves the source image; do not copy a "
        "MEDIA path into the tool call. Hosted quality is tried first and local "
        "FLUX is the fallback."
    ),
    "parameters": IMAGE_EDIT_PARAMETERS,
}

LOCAL_EDIT_SCHEMA = {
    "name": "image_edit_latest_local",
    "description": (
        "Edit the newest generated image using only private local FLUX. Use "
        "when an edit follow-up says local, private, FLUX, or on my hardware. "
        "The server resolves the source image; pass only the edit prompt and "
        "optional aspect ratio."
    ),
    "parameters": IMAGE_EDIT_PARAMETERS,
}

HOSTED_EDIT_SCHEMA = {
    "name": "image_edit_latest_hosted",
    "description": (
        "Edit the newest generated image using only hosted OpenAI GPT Image at "
        "highest quality. Use when an edit follow-up says OpenAI, GPT Image, "
        "hosted, or highest hosted quality. The server resolves the source "
        "image; pass only the edit prompt and optional aspect ratio."
    ),
    "parameters": IMAGE_EDIT_PARAMETERS,
}


def _broker_key() -> str:
    """Load the internal relay key from its runtime-only Vault file."""
    secret_path = Path(
        os.environ.get(
            "HERMES_IMAGE_BROKER_KEY_FILE",
            "/runtime-access/chat-relay-key",
        )
    )
    try:
        value = secret_path.read_text(encoding="utf-8").strip()
        if value:
            return value
    except OSError:
        pass
    return ""


def _request(
    path: str, payload: dict[str, Any] | None, timeout: float
) -> dict[str, Any]:
    """Call the internal broker with a bounded authenticated JSON request."""
    key = _broker_key()
    if not key:
        raise RuntimeError("private image broker key is unavailable")
    data = None if payload is None else json.dumps(payload).encode("utf-8")
    request = Request(
        f"{BROKER_URL}{path}",
        data=data,
        method="GET" if data is None else "POST",
        headers={
            "Authorization": f"Bearer {key}",
            "Content-Type": "application/json",
            "Accept": "application/json",
        },
    )
    try:
        with urlopen(request, timeout=timeout) as response:  # noqa: S310 - fixed internal URL
            return json.loads(response.read())
    except HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")[:1000]
        raise RuntimeError(f"image broker returned HTTP {exc.code}: {body}") from exc
    except URLError as exc:
        raise RuntimeError(f"image broker is unavailable: {exc.reason}") from exc


def _local_image_data_url(value: str) -> str:
    """Convert a private uploaded image to a data URL for broker transport."""
    candidate = value.strip()
    # Hermes presents generated artifacts to the model as ``MEDIA:<path>``.
    # Most models pass only the path back to an edit tool, but accepting the
    # marked form makes natural follow-up edits provider-independent.
    if candidate.upper().startswith("MEDIA:"):
        candidate = candidate[6:].strip()
    lowered = candidate.lower()
    if lowered.startswith(("http://", "https://", "data:image/")):
        return candidate
    try:
        from agent.file_safety import get_read_block_error

        blocked = get_read_block_error(candidate)
        if blocked:
            raise ValueError(blocked)
    except ImportError:
        pass
    path = Path(os.path.expanduser(candidate)).resolve()
    raw = path.read_bytes()
    if not raw or len(raw) > MAX_INPUT_BYTES:
        raise ValueError("reference image must be between 1 byte and 25 MiB")
    if raw.startswith(b"\x89PNG\r\n\x1a\n"):
        mime = "image/png"
    elif raw.startswith(b"\xff\xd8\xff"):
        mime = "image/jpeg"
    elif raw.startswith((b"GIF87a", b"GIF89a")):
        mime = "image/gif"
    elif raw.startswith(b"RIFF") and raw[8:12] == b"WEBP":
        mime = "image/webp"
    else:
        raise ValueError("reference image must be PNG, JPEG, GIF, or WebP")
    return f"data:{mime};base64,{base64.b64encode(raw).decode('ascii')}"


def _latest_generated_image() -> str:
    """Return the newest generated artifact from this tenant's private cache."""
    home = Path(os.environ.get("HERMES_HOME", "/opt/data")).resolve()
    cache = (home / "cache" / "images").resolve()
    if not cache.is_dir():
        raise ValueError("no generated image is available to edit")
    candidates: list[Path] = []
    for path in cache.iterdir():
        try:
            resolved = path.resolve(strict=True)
            resolved.relative_to(cache)
        except (OSError, ValueError):
            continue
        if (
            resolved.is_file()
            and resolved.suffix.lower() in {".png", ".jpg", ".jpeg", ".gif", ".webp"}
        ):
            candidates.append(resolved)
    if not candidates:
        raise ValueError("no generated image is available to edit")
    return str(max(candidates, key=lambda item: item.stat().st_mtime_ns))


class AtlasBrokerImageProvider(ImageGenProvider):
    """High-quality GPT Image generation through the owner-only broker."""

    @property
    def name(self) -> str:
        return "atlas-broker"

    @property
    def display_name(self) -> str:
        return "Atlas private image studio"

    def is_available(self) -> bool:
        try:
            result = _request("/health", None, 5.0)
            return bool(result.get("success") and result.get("available"))
        except Exception:
            return False

    def list_models(self) -> list[dict[str, Any]]:
        models = []
        for model, values in MODELS.items():
            display, _quality, *strengths = values
            local = model == "flux-2-klein-4b-local"
            models.append(
                {
                    "id": model,
                    "display": display,
                    "speed": "~2–15min" if local else "~30s–3min",
                    "strengths": strengths[0]
                    if strengths
                    else "OpenAI GPT Image 2 quality",
                    "price": "local compute" if local else "included account capacity",
                }
            )
        return models

    def default_model(self) -> str:
        return DEFAULT_MODEL

    def capabilities(self) -> dict[str, Any]:
        return {"modalities": ["text", "image"], "max_reference_images": 16}

    def generate(
        self,
        prompt: str,
        aspect_ratio: str = DEFAULT_ASPECT_RATIO,
        *,
        image_url: str | None = None,
        reference_image_urls: list[str] | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        prompt = (prompt or "").strip()
        aspect = resolve_aspect_ratio(aspect_ratio)
        if not prompt:
            return error_response(
                error="Prompt is required",
                error_type="invalid_argument",
                provider=self.name,
                aspect_ratio=aspect,
            )
        model = str(kwargs.get("model") or DEFAULT_MODEL)
        if model not in MODELS:
            model = DEFAULT_MODEL
        try:
            primary = _local_image_data_url(image_url) if image_url else None
            references = [
                _local_image_data_url(item)
                for item in (normalize_reference_images(reference_image_urls) or [])[
                    :16
                ]
            ]
            response = _request(
                "/v1/images/generations",
                {
                    "prompt": prompt,
                    "aspect_ratio": aspect,
                    "model": model,
                    "image_url": primary,
                    "reference_image_urls": references,
                },
                1200.0,
            )
            if not response.get("success"):
                return error_response(
                    error=str(response.get("error") or "image generation failed"),
                    error_type=str(response.get("error_type") or "provider_error"),
                    provider=self.name,
                    model=model,
                    prompt=prompt,
                    aspect_ratio=aspect,
                )
            image = save_b64_image(str(response["image_b64"]), prefix=f"atlas_{model}")
            return success_response(
                image=str(image),
                model=str(response.get("model") or model),
                prompt=prompt,
                aspect_ratio=aspect,
                provider=self.name,
                modality="image" if primary or references else "text",
                extra={
                    "quality": response.get("quality", MODELS[model][1]),
                    "size": response.get("size"),
                    "route": response.get("route"),
                    "requested_model": response.get("requested_model", model),
                    "hosted_fallback_reason": response.get("hosted_fallback_reason"),
                    "input_image_count": response.get(
                        "input_image_count", len(references) + bool(primary)
                    ),
                },
            )
        except Exception as exc:
            return error_response(
                error=f"Private image studio failed: {exc}",
                error_type="broker_error",
                provider=self.name,
                model=model,
                prompt=prompt,
                aspect_ratio=aspect,
            )


def _handle_image_generate(args: dict[str, Any], route: str) -> str:
    """Dispatch one image request through a fixed route."""
    prompt = str(args.get("prompt") or "").strip()
    if not prompt:
        return json.dumps(
            {
                "success": False,
                "image": None,
                "error": "prompt is required for image generation",
                "error_type": "invalid_argument",
            }
        )
    model = ROUTE_MODELS.get(route)
    if model is None:  # pragma: no cover - routes are fixed by the wrappers
        raise ValueError(f"unsupported image route: {route}")
    provider = AtlasBrokerImageProvider()
    result = provider.generate(
        prompt,
        str(args.get("aspect_ratio") or DEFAULT_ASPECT_RATIO),
        model=model,
        image_url=args.get("image_url"),
        reference_image_urls=args.get("reference_image_urls"),
    )
    if isinstance(result, dict):
        result.setdefault("requested_route", route)
    return json.dumps(result, ensure_ascii=False)


def _handle_local_image(args: dict[str, Any], **_kwargs: Any) -> str:
    """Generate an image without allowing hosted provider substitution."""
    return _handle_image_generate(args, "local")


def _handle_hosted_image(args: dict[str, Any], **_kwargs: Any) -> str:
    """Generate an image through the highest-quality hosted route."""
    return _handle_image_generate(args, "hosted")


def _handle_latest_edit(args: dict[str, Any], route: str) -> str:
    """Edit the newest tenant artifact without model-generated path arguments."""
    try:
        image_url = _latest_generated_image()
    except ValueError as exc:
        return json.dumps(
            {
                "success": False,
                "image": None,
                "error": str(exc),
                "error_type": "missing_reference",
                "requested_route": route,
            }
        )
    return _handle_image_generate({**args, "image_url": image_url}, route)


def _handle_auto_edit(args: dict[str, Any], **_kwargs: Any) -> str:
    """Edit the newest image with hosted-first automatic failover."""
    return _handle_latest_edit(args, "auto")


def _handle_local_edit(args: dict[str, Any], **_kwargs: Any) -> str:
    """Edit the newest image without hosted provider substitution."""
    return _handle_latest_edit(args, "local")


def _handle_hosted_edit(args: dict[str, Any], **_kwargs: Any) -> str:
    """Edit the newest image only with hosted GPT Image."""
    return _handle_latest_edit(args, "hosted")


def _image_tool_available() -> bool:
    """Expose the tool when at least one broker route is healthy."""
    return AtlasBrokerImageProvider().is_available()


def register(ctx: Any) -> None:
    """Register generation routes and compact latest-image edit tools."""
    ctx.register_image_gen_provider(AtlasBrokerImageProvider())
    ctx.register_tool(
        name="image_generate_local",
        toolset="image_gen",
        schema=LOCAL_IMAGE_SCHEMA,
        handler=_handle_local_image,
        check_fn=_image_tool_available,
        requires_env=[],
        is_async=False,
        description=LOCAL_IMAGE_SCHEMA["description"],
        emoji="🎨",
    )
    ctx.register_tool(
        name="image_generate_hosted",
        toolset="image_gen",
        schema=HOSTED_IMAGE_SCHEMA,
        handler=_handle_hosted_image,
        check_fn=_image_tool_available,
        requires_env=[],
        is_async=False,
        description=HOSTED_IMAGE_SCHEMA["description"],
        emoji="🎨",
    )
    for name, schema, handler in (
        ("image_edit_latest", AUTO_EDIT_SCHEMA, _handle_auto_edit),
        ("image_edit_latest_local", LOCAL_EDIT_SCHEMA, _handle_local_edit),
        ("image_edit_latest_hosted", HOSTED_EDIT_SCHEMA, _handle_hosted_edit),
    ):
        ctx.register_tool(
            name=name,
            toolset="image_gen",
            schema=schema,
            handler=handler,
            check_fn=_image_tool_available,
            requires_env=[],
            is_async=False,
            description=schema["description"],
            emoji="🎨",
        )