hermes-agent/plugins/image_gen/xai/__init__.py

"""xAI image generation backend.

Exposes xAI's ``grok-imagine-image`` models as an :class:`ImageGenProvider`:
text-to-image (``/v1/images/generations``, several aspect ratios, 1K/2K) and
image editing (``/v1/images/edits``), base64 output saved to cache.

Selection precedence (first hit wins): ``model`` kwarg → ``XAI_IMAGE_MODEL``
env → ``image_gen.xai.model`` → :data:`DEFAULT_MODEL`.
"""

from __future__ import annotations

import logging
import os
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import requests

from agent.image_gen_provider import (
    DEFAULT_ASPECT_RATIO,
    ImageGenProvider,
    resolve_aspect_ratio,
    success_response,
)
from plugins.image_gen._common import (
    catalog_rows,
    collect_source_images,
    error_factory,
    load_image_gen_config,
    materialize_image,
    post_json,
)
from tools.xai_http import (
    build_xai_storage_options,
    hermes_xai_user_agent,
    maybe_mark_xai_storage_notice_seen,
    read_xai_imagine_storage_config,
    resolve_xai_http_credentials,
    xai_storage_notice_text,
)

logger = logging.getLogger(__name__)

_MODELS: Dict[str, Dict[str, Any]] = {
    "grok-imagine-image": {
        "display": "Grok Imagine Image",
        "speed": "~5-10s",
        "strengths": "Fast, high-quality",
    },
    "grok-imagine-image-2.0": {
        "display": "Grok Imagine Image 2.0",
        "speed": "~10-20s",
        "strengths": "Typography/layout-aware; legible small text; strongest quality.",
    },
    "grok-imagine-image-quality": {
        "display": "Grok Imagine Image (Quality)",
        "speed": "~10-20s",
        "strengths": "Higher fidelity / detail; slower than the standard model.",
    },
}

DEFAULT_MODEL = "grok-imagine-image"
# xAI documents the quality model as the edit-capable baseline.
_EDIT_FALLBACK_MODEL = "grok-imagine-image-quality"

# Live catalog cache: (models_dict, fetched_monotonic). xAI's
# ``/image-generation-models`` endpoint is the source of truth so newly
# released Imagine models appear in the picker without a code change; the
# static ``_MODELS`` table is the offline fallback and supplies curated
# speed/strengths text for the models we know about.
_LIVE_CACHE: Optional[Tuple[Dict[str, Dict[str, Any]], float]] = None
_LIVE_CACHE_TTL = 300.0
_LIVE_TIMEOUT = 10.0

# xAI aspect ratios (more options than FAL/OpenAI)
_XAI_ASPECT_RATIOS = {
    "landscape": "16:9",
    "square": "1:1",
    "portrait": "9:16",
    "4:3": "4:3",
    "3:4": "3:4",
    "3:2": "3:2",
    "2:3": "2:3",
}
_XAI_RESOLUTIONS = {"1k", "2k"}
DEFAULT_RESOLUTION = "1k"
_MAX_SOURCE_IMAGES = 3
_REQUEST_TIMEOUT = 120


def _base_url(creds: Dict[str, Any]) -> str:
    return str(creds.get("base_url") or "https://api.x.ai/v1").strip().rstrip("/")


def _fetch_live_models() -> Dict[str, Dict[str, Any]]:
    """Fetch image models from xAI's ``/image-generation-models`` endpoint.

    Returns ``{model_id: {"input_modalities": [...], "aliases": [...]}}``.
    Raises on any failure — callers treat that as "use the static table".
    """
    creds = resolve_xai_http_credentials()
    api_key = str(creds.get("api_key") or "").strip()
    if not api_key:
        raise RuntimeError("no xAI credentials")
    response = requests.get(
        f"{_base_url(creds)}/image-generation-models",
        headers={"Authorization": f"Bearer {api_key}", "User-Agent": hermes_xai_user_agent()},
        timeout=_LIVE_TIMEOUT,
    )
    response.raise_for_status()
    payload = response.json()
    entries = payload.get("models") or payload.get("data") or []
    out: Dict[str, Dict[str, Any]] = {}
    for entry in entries:
        if not isinstance(entry, dict):
            continue
        model_id = entry.get("id") or entry.get("name")
        if not isinstance(model_id, str) or not model_id.strip():
            continue
        out[model_id.strip()] = {
            "input_modalities": entry.get("input_modalities") or [],
            "aliases": entry.get("aliases") or [],
        }
    return out


def _live_models() -> Dict[str, Dict[str, Any]]:
    """Cached live catalog (``{}`` when unreachable)."""
    global _LIVE_CACHE
    import time

    if _LIVE_CACHE is not None and time.monotonic() - _LIVE_CACHE[1] < _LIVE_CACHE_TTL:
        return _LIVE_CACHE[0]
    try:
        live = _fetch_live_models()
    except Exception as exc:  # noqa: BLE001 - offline/unauth → static fallback
        logger.debug("xAI live image model catalog unavailable: %s", exc)
        live = {}
    _LIVE_CACHE = (live, time.monotonic())
    return live


def _catalog() -> Dict[str, Dict[str, Any]]:
    """Merged model catalog: live endpoint IDs + curated static metadata.

    Known models keep their curated display/speed/strengths; models xAI ships
    later still show up (with generic metadata) so users can pick them the day
    they launch. Curated entries the live list momentarily omits are kept;
    static table alone when the API is unreachable.
    """
    live = _live_models()
    if not live:
        return dict(_MODELS)
    merged: Dict[str, Dict[str, Any]] = {}
    for model_id in live:
        meta = _MODELS.get(model_id) or {
            "display": model_id,
            "speed": "",
            "strengths": "New xAI Imagine model (from live xAI catalog)",
        }
        merged[model_id] = dict(meta)
        merged[model_id]["input_modalities"] = live[model_id].get("input_modalities") or []
    for model_id, meta in _MODELS.items():
        merged.setdefault(model_id, dict(meta))
    return merged


def _load_xai_config() -> Dict[str, Any]:
    return load_image_gen_config("xai")


def _configured_model() -> Optional[str]:
    value = _load_xai_config().get("model")
    return value if isinstance(value, str) else None


def _resolve_model(caller_model: Optional[str] = None) -> Tuple[str, Dict[str, Any]]:
    """Return ``(model_id, meta)``: caller kwarg → ``XAI_IMAGE_MODEL`` → config → default.

    Every candidate is validated against the merged live+static catalog, so a
    newly released xAI model is selectable the day it appears in the live catalog.
    """
    catalog = _catalog()
    for candidate in (caller_model, os.environ.get("XAI_IMAGE_MODEL"), _configured_model()):
        if candidate and candidate in catalog:
            return candidate, catalog[candidate]
    return DEFAULT_MODEL, catalog.get(DEFAULT_MODEL, _MODELS[DEFAULT_MODEL])


def _resolve_edit_model(caller_model: Optional[str] = None) -> str:
    """Model for ``/v1/images/edits``: an explicitly selected model that accepts
    image input is honored; otherwise the documented quality baseline."""
    catalog = _catalog()
    explicit = caller_model or os.environ.get("XAI_IMAGE_MODEL") or _configured_model()
    if explicit and explicit in catalog:
        if "image" in (catalog[explicit].get("input_modalities") or []):
            return explicit
    return _EDIT_FALLBACK_MODEL


def _resolve_resolution() -> str:
    res = _load_xai_config().get("resolution")
    return res if isinstance(res, str) and res in _XAI_RESOLUTIONS else DEFAULT_RESOLUTION


def _xai_image_field(source: str) -> Dict[str, str]:
    """Build the xAI ``image`` field for an edit request.

    ``/v1/images/edits`` accepts a public HTTPS URL or a base64 data URI; local
    file paths are read and encoded into a ``data:`` URI.
    """
    source = source.strip()
    if source.lower().startswith(("http://", "https://", "data:")):
        return {"url": source, "type": "image_url"}
    import base64

    # Shared credential-read guard before reading local bytes (same boundary
    # the OpenAI / OpenRouter / Codex image providers apply).
    from agent.file_safety import raise_if_read_blocked

    raise_if_read_blocked(source)
    with open(os.path.expanduser(source), "rb") as fh:  # windows-footgun: ok
        raw = fh.read()
    ext = (os.path.splitext(source)[1].lstrip(".") or "png").lower()
    if ext == "jpg":
        ext = "jpeg"
    b64 = base64.b64encode(raw).decode("utf-8")
    return {"url": f"data:image/{ext};base64,{b64}", "type": "image_url"}


class XAIImageGenProvider(ImageGenProvider):
    """xAI ``grok-imagine-image`` backend."""

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

    @property
    def display_name(self) -> str:
        return "xAI (Grok)"

    def is_available(self) -> bool:
        return bool(resolve_xai_http_credentials().get("api_key"))

    def list_models(self) -> List[Dict[str, Any]]:
        return catalog_rows(_catalog(), ("display", "speed", "strengths"))

    def get_setup_schema(self) -> Dict[str, Any]:
        # Auth resolution is delegated to the shared ``xai_grok`` post_setup hook
        # (``hermes_cli/tools_config.py``), identical to the TTS / video-gen
        # entries, so every xAI service shows the same OAuth-or-API-key choice.
        storage_notice = xai_storage_notice_text("image_gen")
        tag = (
            "grok-imagine-image - text-to-image & image editing; uses xAI "
            "Grok OAuth or XAI_API_KEY"
        )
        if storage_notice:
            tag += f". {storage_notice}"
        return {
            "name": "xAI Grok Imagine (image)",
            "badge": "paid",
            "tag": tag,
            "env_vars": [],
            "post_setup": "xai_grok",
        }

    def capabilities(self) -> Dict[str, Any]:
        # /v1/images/edits accepts up to 3 total source images.
        return {
            "modalities": ["text", "image"],
            "max_reference_images": 2,
            "max_source_images": _MAX_SOURCE_IMAGES,
        }

    def generate(
        self,
        prompt: str,
        aspect_ratio: str = DEFAULT_ASPECT_RATIO,
        *,
        image_url: Optional[str] = None,
        reference_image_urls: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> Dict[str, Any]:
        """Text-to-image via ``/v1/images/generations``, or image editing via
        ``/v1/images/edits`` when source images are supplied.

        Editing uses a JSON body (xAI does not support the OpenAI SDK's
        multipart ``images.edit()``) and an image-input-capable model.
        """
        creds = resolve_xai_http_credentials()
        api_key = str(creds.get("api_key") or "").strip()
        provider_name = str(creds.get("provider") or "xai").strip() or "xai"
        if not api_key:
            return error_factory(provider_name, aspect_ratio)(
                "No xAI credentials found. Configure xAI OAuth in `hermes model` or set XAI_API_KEY.",
                "missing_api_key",
            )

        model_id, meta = _resolve_model(kwargs.get("model"))
        aspect = resolve_aspect_ratio(aspect_ratio)
        xai_res = _resolve_resolution()

        source_images = collect_source_images(image_url, reference_image_urls)
        edit_fail = error_factory(provider_name, aspect, model=_EDIT_FALLBACK_MODEL, prompt=prompt)
        if len(source_images) > _MAX_SOURCE_IMAGES:
            return edit_fail(
                f"xAI image editing supports at most {_MAX_SOURCE_IMAGES} source images",
                "too_many_references",
            )
        for index, source in enumerate(source_images):
            field = "image_url" if index == 0 and image_url and image_url.strip() == source else "reference_image_urls"
            if not source.lower().startswith(("http://", "https://", "data:")):
                if not Path(source).expanduser().is_file():
                    return edit_fail(
                        f"{field} must be a public HTTPS URL or data URI "
                        "(e.g. the `image`/`public_url` from a prior Imagine result)",
                        "invalid_image_url",
                    )
        is_edit = bool(source_images)

        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "User-Agent": hermes_xai_user_agent(),
        }
        base_url = _base_url(creds)
        storage_options = build_xai_storage_options(
            "image_gen", filename_prefix="hermes-xai-image", extension="png",
        )
        storage_notice = maybe_mark_xai_storage_notice_seen("image_gen")
        storage_cfg = read_xai_imagine_storage_config("image_gen")

        if is_edit:
            # An explicit user selection that accepts image input (e.g.
            # grok-imagine-image-2.0) is honored; otherwise the quality baseline.
            model_id = _resolve_edit_model(kwargs.get("model"))
            try:
                image_fields = [_xai_image_field(source) for source in source_images]
            except Exception as exc:
                return edit_fail(
                    f"Could not load source image for editing: {exc}", "io_error", model=model_id,
                )
            payload: Dict[str, Any] = {"model": model_id, "prompt": prompt}
            if len(image_fields) == 1:
                payload["image"] = image_fields[0]
            else:
                payload["images"] = image_fields
            endpoint_url = f"{base_url}/images/edits"
        else:
            payload = {
                "model": model_id,
                "prompt": prompt,
                "aspect_ratio": _XAI_ASPECT_RATIOS.get(aspect, "1:1"),
                "resolution": xai_res,
            }
            endpoint_url = f"{base_url}/images/generations"
        if storage_options is not None:
            payload["storage_options"] = storage_options

        fail = error_factory(provider_name, aspect, model=model_id, prompt=prompt)
        result, failure = post_json(
            endpoint_url, headers=headers, payload=payload, timeout=_REQUEST_TIMEOUT, label="xAI",
        )
        if failure:
            if failure.kind == "http":
                logger.error("xAI image gen failed (%d): %s", failure.status, failure.message)
            return fail(failure.error, failure.error_type)

        # xAI returns data[0].b64_json, data[0].url, and optionally
        # data[0].file_output when storage_options were requested.
        data = result.get("data", [])
        if not data:
            return fail("xAI returned no image data", "empty_response")

        first = data[0]
        file_output = first.get("file_output") if isinstance(first, dict) else None
        file_output = file_output if isinstance(file_output, dict) else {}
        public_url = file_output.get("public_url") if isinstance(file_output.get("public_url"), str) else None

        if public_url:
            image_ref = public_url
        else:
            # grok-imagine-image URLs (``imgen.x.ai/xai-tmp-*``) 404 within
            # minutes; materialise locally so downstream consumers (Telegram
            # send_photo …) get a stable path.
            image_ref, err = materialize_image(
                first.get("b64_json"), first.get("url"),
                prefix=f"xai_{model_id}", label="xAI", provider="xai",
                model=model_id, prompt=prompt, aspect=aspect, log=logger,
            )
            if err:
                return err

        extra: Dict[str, Any] = {"storage_enabled": bool(storage_cfg["enabled"])}
        if not is_edit:
            extra["resolution"] = xai_res
        if storage_notice:
            extra["storage_notice"] = storage_notice
        if public_url:
            extra["public_url"] = public_url
        for key in (
            "filename", "expires_at", "public_url_expires_at", "public_url_error", "storage_error",
        ):
            if key in file_output:
                extra[key] = file_output[key]
        if result.get("usage"):
            extra["usage"] = result["usage"]

        return success_response(
            image=image_ref,
            model=model_id,
            prompt=prompt,
            aspect_ratio=aspect,
            provider="xai",
            modality="image" if is_edit else "text",
            extra=extra,
        )


def register(ctx: Any) -> None:
    """Register this provider with the image gen registry."""
    ctx.register_image_gen_provider(XAIImageGenProvider())