diff --git a/hermes_cli/config.py b/hermes_cli/config.py index 7b381392092..77cffed56a6 100644 --- a/hermes_cli/config.py +++ b/hermes_cli/config.py @@ -2532,6 +2532,14 @@ OPTIONAL_ENV_VARS = { "password": True, "category": "tool", }, + "KREA_API_KEY": { + "description": "Krea API key for Krea 2 image generation (Medium + Large)", + "prompt": "Krea API key", + "url": "https://www.krea.ai/settings/api-tokens", + "tools": ["image_generate"], + "password": True, + "category": "tool", + }, "VOICE_TOOLS_OPENAI_KEY": { "description": "OpenAI API key for voice transcription (Whisper) and OpenAI TTS", "prompt": "OpenAI API Key (for Whisper STT + TTS)", diff --git a/plugins/image_gen/krea/__init__.py b/plugins/image_gen/krea/__init__.py new file mode 100644 index 00000000000..552f2ae71fe --- /dev/null +++ b/plugins/image_gen/krea/__init__.py @@ -0,0 +1,548 @@ +"""Krea image generation backend. + +Exposes Krea's `Krea 2` foundation image model family — Krea 2 Medium and +Krea 2 Large — as an :class:`ImageGenProvider` implementation. + +Krea's API is asynchronous: the generate endpoint returns a ``job_id`` +that you poll at ``GET /jobs/{job_id}``. This provider hides that +roundtrip behind the synchronous ``generate()`` contract: submit, poll +every 2s with light backoff, materialise the result URL to local cache, +return the success/error dict like every other backend. + +Selection precedence (first hit wins): + +1. ``KREA_IMAGE_MODEL`` env var (escape hatch for scripts / tests) +2. ``image_gen.krea.model`` in ``config.yaml`` +3. ``image_gen.model`` in ``config.yaml`` (when it's one of our IDs) +4. :data:`DEFAULT_MODEL` — ``krea-2-medium`` (Krea's "start here" recommendation) + +Docs: https://docs.krea.ai/developers/krea-2/overview +API: https://docs.krea.ai/api-reference/krea/krea-2-large +""" + +from __future__ import annotations + +import logging +import os +import time +from typing import Any, Dict, List, Optional, Tuple + +import requests + +from agent.image_gen_provider import ( + DEFAULT_ASPECT_RATIO, + ImageGenProvider, + error_response, + resolve_aspect_ratio, + save_url_image, + success_response, +) + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +BASE_URL = "https://api.krea.ai" + +# Map our short model IDs to Krea's URL path segment. +_MODELS: Dict[str, Dict[str, Any]] = { + "krea-2-medium": { + "display": "Krea 2 Medium", + "speed": "~15-25s", + "strengths": "Illustration, anime, painting, expressive styles. Faster + cheaper.", + "price": "$0.030 (text) / $0.035 (style refs) / $0.040 (moodboards)", + "path": "medium", + }, + "krea-2-large": { + "display": "Krea 2 Large", + "speed": "~25-60s", + "strengths": "Photorealism, raw textured looks (motion blur, grain), expressive styles.", + "price": "$0.060 (text) / $0.065 (style refs) / $0.070 (moodboards)", + "path": "large", + }, +} + +DEFAULT_MODEL = "krea-2-medium" + +# Hermes uses 3 abstract aspect ratios. Map to Krea's enum (which is wider). +# Krea accepts: 1:1, 4:3, 3:2, 16:9, 2.35:1, 4:5, 2:3, 9:16 +_ASPECT_MAP = { + "landscape": "16:9", + "square": "1:1", + "portrait": "9:16", +} + +# Only resolution Krea currently supports. +DEFAULT_RESOLUTION = "1K" + +# Valid creativity levels per Krea docs. Default is "medium". +_VALID_CREATIVITY = {"raw", "low", "medium", "high"} + +# Polling cadence. Krea recommends 2-5s; we start at 2s and back off to 5s +# for long jobs (Large can take ~1min). Total ceiling matches Krea's +# hosted-tool timeout of 3 minutes. +_POLL_INITIAL_INTERVAL = 2.0 +_POLL_MAX_INTERVAL = 5.0 +_POLL_BACKOFF = 1.3 +_POLL_TIMEOUT_SECONDS = 180.0 + +# HTTP statuses worth retrying during the poll loop. Everything else (401, +# 402, 403, 404, other 4xx) is a permanent failure — surface it immediately +# instead of burning the 180s deadline retrying a request that will never +# succeed. +_RETRYABLE_POLL_STATUSES = frozenset({408, 409, 425, 429, 500, 502, 503, 504}) + +_TERMINAL_STATES = {"completed", "failed", "cancelled"} + + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- + + +def _load_krea_config() -> Dict[str, Any]: + """Read ``image_gen.krea`` (with fallthrough to ``image_gen``) from config.yaml.""" + try: + from hermes_cli.config import load_config + + cfg = load_config() + section = cfg.get("image_gen") if isinstance(cfg, dict) else None + return section if isinstance(section, dict) else {} + except Exception as exc: # noqa: BLE001 + logger.debug("Could not load image_gen config: %s", exc) + return {} + + +def _resolve_model() -> Tuple[str, Dict[str, Any]]: + """Decide which model to use and return ``(model_id, meta)``.""" + env_override = os.environ.get("KREA_IMAGE_MODEL") + if env_override and env_override in _MODELS: + return env_override, _MODELS[env_override] + + cfg = _load_krea_config() + krea_cfg = cfg.get("krea") if isinstance(cfg.get("krea"), dict) else {} + candidate: Optional[str] = None + if isinstance(krea_cfg, dict): + value = krea_cfg.get("model") + if isinstance(value, str) and value in _MODELS: + candidate = value + if candidate is None: + top = cfg.get("model") + if isinstance(top, str) and top in _MODELS: + candidate = top + + if candidate is not None: + return candidate, _MODELS[candidate] + + return DEFAULT_MODEL, _MODELS[DEFAULT_MODEL] + + +def _resolve_creativity(value: Optional[str]) -> str: + """Coerce ``creativity`` kwarg to a valid Krea value (default ``medium``).""" + if isinstance(value, str): + v = value.strip().lower() + if v in _VALID_CREATIVITY: + return v + cfg = _load_krea_config() + krea_cfg = cfg.get("krea") if isinstance(cfg.get("krea"), dict) else {} + cfg_value = krea_cfg.get("creativity") if isinstance(krea_cfg, dict) else None + if isinstance(cfg_value, str) and cfg_value.strip().lower() in _VALID_CREATIVITY: + return cfg_value.strip().lower() + return "medium" + + +# --------------------------------------------------------------------------- +# Provider +# --------------------------------------------------------------------------- + + +class KreaImageGenProvider(ImageGenProvider): + """Krea ``Krea 2`` foundation image model backend (Medium + Large).""" + + @property + def name(self) -> str: + return "krea" + + @property + def display_name(self) -> str: + return "Krea" + + def is_available(self) -> bool: + return bool(os.environ.get("KREA_API_KEY")) + + def list_models(self) -> List[Dict[str, Any]]: + return [ + { + "id": model_id, + "display": meta["display"], + "speed": meta["speed"], + "strengths": meta["strengths"], + "price": meta["price"], + } + for model_id, meta in _MODELS.items() + ] + + def default_model(self) -> Optional[str]: + return DEFAULT_MODEL + + def get_setup_schema(self) -> Dict[str, Any]: + return { + "name": "Krea", + "badge": "paid", + "tag": "Krea 2 foundation model — Medium ($0.03) + Large ($0.06). Strong style transfer + moodboards.", + "env_vars": [ + { + "key": "KREA_API_KEY", + "prompt": "Krea API key", + "url": "https://www.krea.ai/settings/api-tokens", + }, + ], + } + + # ------------------------------------------------------------------ + # generate() + # ------------------------------------------------------------------ + + def generate( + self, + prompt: str, + aspect_ratio: str = DEFAULT_ASPECT_RATIO, + **kwargs: Any, + ) -> Dict[str, Any]: + prompt = (prompt or "").strip() + aspect = resolve_aspect_ratio(aspect_ratio) + krea_ar = _ASPECT_MAP.get(aspect, "1:1") + + if not prompt: + return error_response( + error="Prompt is required and must be a non-empty string", + error_type="invalid_argument", + provider="krea", + aspect_ratio=aspect, + ) + + api_key = os.environ.get("KREA_API_KEY") + if not api_key: + return error_response( + error=( + "KREA_API_KEY not set. Run `hermes tools` → Image " + "Generation → Krea to configure, or get a key at " + "https://www.krea.ai/settings/api-tokens." + ), + error_type="auth_required", + provider="krea", + aspect_ratio=aspect, + ) + + model_id, meta = _resolve_model() + creativity = _resolve_creativity(kwargs.get("creativity")) + + payload: Dict[str, Any] = { + "prompt": prompt, + "aspect_ratio": krea_ar, + "resolution": DEFAULT_RESOLUTION, + "creativity": creativity, + } + + # Optional forward-compat passthroughs — the Krea API accepts these + # but they're not required and most agent calls won't supply them. + seed = kwargs.get("seed") + if isinstance(seed, int): + payload["seed"] = seed + + styles = kwargs.get("styles") + if isinstance(styles, list) and styles: + payload["styles"] = styles + + image_style_references = kwargs.get("image_style_references") + if isinstance(image_style_references, list) and image_style_references: + # Krea caps at 10 refs per request. + payload["image_style_references"] = image_style_references[:10] + + moodboards = kwargs.get("moodboards") + if isinstance(moodboards, list) and moodboards: + # Krea currently caps at 1 moodboard per request. + payload["moodboards"] = moodboards[:1] + + headers = { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + "User-Agent": "Hermes-Agent/1.0 (krea-image-gen)", + } + + # 1. Submit job. + submit_url = f"{BASE_URL}/generate/image/krea/krea-2/{meta['path']}" + try: + response = requests.post( + submit_url, + headers=headers, + json=payload, + timeout=30, + ) + response.raise_for_status() + except requests.HTTPError as exc: + resp = exc.response + status = resp.status_code if resp is not None else 0 + try: + body = resp.json() if resp is not None else {} + err_msg = ( + body.get("error", {}).get("message") + if isinstance(body.get("error"), dict) + else body.get("message") or body.get("detail") + ) or (resp.text[:300] if resp is not None else str(exc)) + except Exception: # noqa: BLE001 + err_msg = resp.text[:300] if resp is not None else str(exc) + logger.error("Krea submit failed (%d): %s", status, err_msg) + return error_response( + error=f"Krea image generation failed ({status}): {err_msg}", + error_type="api_error", + provider="krea", + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + ) + except requests.Timeout: + return error_response( + error="Krea submit timed out (30s)", + error_type="timeout", + provider="krea", + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + ) + except requests.ConnectionError as exc: + return error_response( + error=f"Krea connection error: {exc}", + error_type="connection_error", + provider="krea", + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + ) + + try: + submit_body = response.json() + except Exception as exc: # noqa: BLE001 + return error_response( + error=f"Krea returned invalid JSON on submit: {exc}", + error_type="invalid_response", + provider="krea", + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + ) + + job_id = submit_body.get("job_id") + if not isinstance(job_id, str) or not job_id: + return error_response( + error="Krea submit response missing job_id", + error_type="invalid_response", + provider="krea", + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + ) + + # 2. Poll for completion. + job_url = f"{BASE_URL}/jobs/{job_id}" + poll_headers = { + "Authorization": f"Bearer {api_key}", + "User-Agent": "Hermes-Agent/1.0 (krea-image-gen)", + } + interval = _POLL_INITIAL_INTERVAL + deadline = time.monotonic() + _POLL_TIMEOUT_SECONDS + last_status: Optional[str] = None + + while True: + time.sleep(interval) + interval = min(interval * _POLL_BACKOFF, _POLL_MAX_INTERVAL) + + try: + poll_resp = requests.get(job_url, headers=poll_headers, timeout=30) + poll_resp.raise_for_status() + except requests.HTTPError as exc: + resp = exc.response + status = resp.status_code if resp is not None else 0 + logger.error("Krea poll failed (%d) for job %s", status, job_id) + # Fail fast for non-retryable statuses (auth/billing/not-found, + # other permanent 4xx) so callers don't wait the full 180s + # deadline on a request that will never succeed. Only retry + # transient statuses such as 408/409/425/429/5xx. + if status not in _RETRYABLE_POLL_STATUSES or time.monotonic() >= deadline: + return error_response( + error=f"Krea poll failed ({status}) for job {job_id}", + error_type="api_error", + provider="krea", + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + ) + # Otherwise keep trying — transient 5xx (and a few retryable + # 4xx like 408/409/425/429) are common on async jobs. + continue + except (requests.Timeout, requests.ConnectionError) as exc: + logger.warning("Krea poll transient error for job %s: %s", job_id, exc) + if time.monotonic() >= deadline: + return error_response( + error=f"Krea poll timed out for job {job_id}: {exc}", + error_type="timeout", + provider="krea", + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + ) + continue + + try: + job = poll_resp.json() + except Exception as exc: # noqa: BLE001 + logger.warning("Krea poll returned invalid JSON for job %s: %s", job_id, exc) + if time.monotonic() >= deadline: + return error_response( + error=f"Krea poll returned invalid JSON: {exc}", + error_type="invalid_response", + provider="krea", + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + ) + continue + + status_str = job.get("status") if isinstance(job, dict) else None + if isinstance(status_str, str): + last_status = status_str + if status_str in _TERMINAL_STATES: + break + + # ``completed_at`` is a backstop terminal marker even when the + # ``status`` enum is unfamiliar (Krea adds new pending states + # over time — backlogged/scheduled/sampling — and we don't + # want to mis-handle a future one). + if isinstance(job, dict) and job.get("completed_at"): + break + + if time.monotonic() >= deadline: + return error_response( + error=( + f"Krea job {job_id} did not complete within " + f"{int(_POLL_TIMEOUT_SECONDS)}s (last status: {last_status or 'unknown'})" + ), + error_type="timeout", + provider="krea", + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + ) + + # 3. Terminal — extract result. + if not isinstance(job, dict): + return error_response( + error="Krea returned non-dict job body", + error_type="invalid_response", + provider="krea", + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + ) + + if last_status == "failed": + err = (job.get("result") or {}).get("error") if isinstance(job.get("result"), dict) else None + return error_response( + error=f"Krea job {job_id} failed: {err or 'unknown error'}", + error_type="api_error", + provider="krea", + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + ) + + if last_status == "cancelled": + return error_response( + error=f"Krea job {job_id} was cancelled", + error_type="cancelled", + provider="krea", + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + ) + + # Successful path — pull URL out of the result. + result = job.get("result") + if not isinstance(result, dict): + return error_response( + error="Krea job completed but result was missing", + error_type="empty_response", + provider="krea", + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + ) + + # Per Krea's job-lifecycle docs the completed payload exposes + # ``result.urls`` (an array). Fall back to a single ``url`` field + # for forward/backward compatibility. + image_url: Optional[str] = None + urls = result.get("urls") + if isinstance(urls, list) and urls: + for candidate in urls: + if isinstance(candidate, str) and candidate.strip(): + image_url = candidate.strip() + break + if image_url is None: + single = result.get("url") + if isinstance(single, str) and single.strip(): + image_url = single.strip() + + if image_url is None: + return error_response( + error="Krea result contained no image URL", + error_type="empty_response", + provider="krea", + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + ) + + # Materialise locally — Krea result URLs may expire, mirroring + # what we do for xAI / OpenAI URL responses (#26942). + try: + saved_path = save_url_image(image_url, prefix=f"krea_{model_id}") + except Exception as exc: # noqa: BLE001 + logger.warning( + "Krea image URL %s could not be cached (%s); falling back to bare URL.", + image_url, + exc, + ) + image_ref = image_url + else: + image_ref = str(saved_path) + + extra: Dict[str, Any] = { + "krea_aspect_ratio": krea_ar, + "resolution": DEFAULT_RESOLUTION, + "creativity": creativity, + "job_id": job_id, + } + if isinstance(job.get("completed_at"), str): + extra["completed_at"] = job["completed_at"] + + return success_response( + image=image_ref, + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + provider="krea", + extra=extra, + ) + + +# --------------------------------------------------------------------------- +# Plugin entry point +# --------------------------------------------------------------------------- + + +def register(ctx) -> None: + """Plugin entry point — wire ``KreaImageGenProvider`` into the registry.""" + ctx.register_image_gen_provider(KreaImageGenProvider()) diff --git a/plugins/image_gen/krea/plugin.yaml b/plugins/image_gen/krea/plugin.yaml new file mode 100644 index 00000000000..bc650dc5222 --- /dev/null +++ b/plugins/image_gen/krea/plugin.yaml @@ -0,0 +1,7 @@ +name: krea +version: 1.0.0 +description: "Krea image generation backend (Krea 2 Large + Krea 2 Medium foundation models)." +author: NousResearch +kind: backend +requires_env: + - KREA_API_KEY diff --git a/tests/plugins/image_gen/test_krea_provider.py b/tests/plugins/image_gen/test_krea_provider.py new file mode 100644 index 00000000000..cc9dcd5a6b0 --- /dev/null +++ b/tests/plugins/image_gen/test_krea_provider.py @@ -0,0 +1,625 @@ +#!/usr/bin/env python3 +"""Tests for Krea image generation provider.""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(autouse=True) +def _fake_api_key(monkeypatch): + """Ensure KREA_API_KEY is set for all tests.""" + monkeypatch.setenv("KREA_API_KEY", "test-key-12345") + + +def _completed_job(url: str = "https://krea.cdn/img.png") -> dict: + return { + "job_id": "00000000-0000-0000-0000-000000000abc", + "status": "completed", + "created_at": "2026-05-27T00:00:00Z", + "completed_at": "2026-05-27T00:00:30Z", + "result": {"urls": [url]}, + } + + +def _submit_response(job_id: str = "00000000-0000-0000-0000-000000000abc"): + resp = MagicMock() + resp.status_code = 200 + resp.raise_for_status = MagicMock() + resp.json.return_value = { + "job_id": job_id, + "status": "queued", + "created_at": "2026-05-27T00:00:00Z", + "completed_at": None, + "result": None, + } + return resp + + +def _poll_response(body: dict): + resp = MagicMock() + resp.status_code = 200 + resp.raise_for_status = MagicMock() + resp.json.return_value = body + return resp + + +# --------------------------------------------------------------------------- +# Provider class tests +# --------------------------------------------------------------------------- + + +class TestKreaImageGenProvider: + def test_name(self): + from plugins.image_gen.krea import KreaImageGenProvider + + assert KreaImageGenProvider().name == "krea" + + def test_display_name(self): + from plugins.image_gen.krea import KreaImageGenProvider + + assert KreaImageGenProvider().display_name == "Krea" + + def test_is_available_with_key(self, monkeypatch): + monkeypatch.setenv("KREA_API_KEY", "sk-test") + from plugins.image_gen.krea import KreaImageGenProvider + + assert KreaImageGenProvider().is_available() is True + + def test_is_available_without_key(self, monkeypatch): + monkeypatch.delenv("KREA_API_KEY", raising=False) + from plugins.image_gen.krea import KreaImageGenProvider + + assert KreaImageGenProvider().is_available() is False + + def test_list_models(self): + from plugins.image_gen.krea import KreaImageGenProvider + + models = KreaImageGenProvider().list_models() + ids = {m["id"] for m in models} + assert {"krea-2-medium", "krea-2-large"} <= ids + # Each entry carries the picker fields the registry expects. + for m in models: + assert m["display"] + assert m["speed"] + assert m["strengths"] + assert m["price"] + + def test_default_model_is_medium(self): + from plugins.image_gen.krea import KreaImageGenProvider + + assert KreaImageGenProvider().default_model() == "krea-2-medium" + + def test_get_setup_schema(self): + from plugins.image_gen.krea import KreaImageGenProvider + + schema = KreaImageGenProvider().get_setup_schema() + assert schema["name"] == "Krea" + assert schema["badge"] == "paid" + env_vars = schema["env_vars"] + assert len(env_vars) == 1 + assert env_vars[0]["key"] == "KREA_API_KEY" + assert "krea.ai" in env_vars[0]["url"] + + +# --------------------------------------------------------------------------- +# Model resolution +# --------------------------------------------------------------------------- + + +class TestModelResolution: + def test_default(self): + from plugins.image_gen.krea import _resolve_model + + model_id, meta = _resolve_model() + assert model_id == "krea-2-medium" + assert meta["path"] == "medium" + + def test_env_override_large(self, monkeypatch): + monkeypatch.setenv("KREA_IMAGE_MODEL", "krea-2-large") + from plugins.image_gen.krea import _resolve_model + + model_id, meta = _resolve_model() + assert model_id == "krea-2-large" + assert meta["path"] == "large" + + def test_env_override_unknown_falls_back_to_default(self, monkeypatch): + monkeypatch.setenv("KREA_IMAGE_MODEL", "krea-2-xxl-fake") + from plugins.image_gen.krea import _resolve_model + + model_id, _ = _resolve_model() + assert model_id == "krea-2-medium" + + def test_creativity_default(self): + from plugins.image_gen.krea import _resolve_creativity + + assert _resolve_creativity(None) == "medium" + + def test_creativity_valid(self): + from plugins.image_gen.krea import _resolve_creativity + + assert _resolve_creativity("HIGH") == "high" + assert _resolve_creativity(" raw ") == "raw" + + def test_creativity_invalid(self): + from plugins.image_gen.krea import _resolve_creativity + + assert _resolve_creativity("ultra") == "medium" + + +# --------------------------------------------------------------------------- +# Generate — main flow +# --------------------------------------------------------------------------- + + +class TestGenerate: + def test_missing_api_key(self, monkeypatch): + monkeypatch.delenv("KREA_API_KEY", raising=False) + from plugins.image_gen.krea import KreaImageGenProvider + + result = KreaImageGenProvider().generate(prompt="test") + assert result["success"] is False + assert "KREA_API_KEY" in result["error"] + assert result["error_type"] == "auth_required" + + def test_empty_prompt(self): + from plugins.image_gen.krea import KreaImageGenProvider + + result = KreaImageGenProvider().generate(prompt=" ") + assert result["success"] is False + assert result["error_type"] == "invalid_argument" + + def test_successful_generation(self): + """Happy path: submit → one poll → completed → URL downloaded.""" + from plugins.image_gen.krea import KreaImageGenProvider + + submit = _submit_response() + poll = _poll_response(_completed_job("https://krea.cdn/result.png")) + + with patch("plugins.image_gen.krea.requests.post", return_value=submit) as mock_post, \ + patch("plugins.image_gen.krea.requests.get", return_value=poll) as mock_get, \ + patch( + "plugins.image_gen.krea.save_url_image", + return_value=Path("/tmp/krea_krea-2-medium_test.png"), + ) as mock_save, \ + patch("plugins.image_gen.krea.time.sleep"): # skip real waits + result = KreaImageGenProvider().generate(prompt="A cinematic lamp") + + assert result["success"] is True + assert result["image"] == "/tmp/krea_krea-2-medium_test.png" + assert result["provider"] == "krea" + assert result["model"] == "krea-2-medium" + assert result["aspect_ratio"] == "landscape" + assert result["job_id"] == "00000000-0000-0000-0000-000000000abc" + assert result["resolution"] == "1K" + assert result["creativity"] == "medium" + # Submit hit the medium endpoint + post_url = mock_post.call_args[0][0] + assert post_url.endswith("/generate/image/krea/krea-2/medium") + # Poll hit /jobs/{job_id} + poll_url = mock_get.call_args[0][0] + assert "/jobs/00000000-0000-0000-0000-000000000abc" in poll_url + # URL was materialised once + mock_save.assert_called_once() + + def test_large_model_routes_to_large_endpoint(self, monkeypatch): + monkeypatch.setenv("KREA_IMAGE_MODEL", "krea-2-large") + from plugins.image_gen.krea import KreaImageGenProvider + + submit = _submit_response() + poll = _poll_response(_completed_job()) + + with patch("plugins.image_gen.krea.requests.post", return_value=submit) as mock_post, \ + patch("plugins.image_gen.krea.requests.get", return_value=poll), \ + patch( + "plugins.image_gen.krea.save_url_image", + return_value=Path("/tmp/x.png"), + ), \ + patch("plugins.image_gen.krea.time.sleep"): + KreaImageGenProvider().generate(prompt="test") + + post_url = mock_post.call_args[0][0] + assert post_url.endswith("/generate/image/krea/krea-2/large") + + def test_aspect_ratio_mapping(self): + """Hermes 'square' must map to Krea '1:1' in the wire payload.""" + from plugins.image_gen.krea import KreaImageGenProvider + + submit = _submit_response() + poll = _poll_response(_completed_job()) + + with patch("plugins.image_gen.krea.requests.post", return_value=submit) as mock_post, \ + patch("plugins.image_gen.krea.requests.get", return_value=poll), \ + patch( + "plugins.image_gen.krea.save_url_image", + return_value=Path("/tmp/x.png"), + ), \ + patch("plugins.image_gen.krea.time.sleep"): + KreaImageGenProvider().generate(prompt="test", aspect_ratio="square") + + payload = mock_post.call_args.kwargs["json"] + assert payload["aspect_ratio"] == "1:1" + assert payload["resolution"] == "1K" + + def test_auth_header(self): + from plugins.image_gen.krea import KreaImageGenProvider + + submit = _submit_response() + poll = _poll_response(_completed_job()) + + with patch("plugins.image_gen.krea.requests.post", return_value=submit) as mock_post, \ + patch("plugins.image_gen.krea.requests.get", return_value=poll), \ + patch( + "plugins.image_gen.krea.save_url_image", + return_value=Path("/tmp/x.png"), + ), \ + patch("plugins.image_gen.krea.time.sleep"): + KreaImageGenProvider().generate(prompt="test") + + headers = mock_post.call_args.kwargs["headers"] + assert headers["Authorization"] == "Bearer test-key-12345" + assert headers["Content-Type"] == "application/json" + + def test_passthrough_seed_styles_moodboards(self): + from plugins.image_gen.krea import KreaImageGenProvider + + submit = _submit_response() + poll = _poll_response(_completed_job()) + + with patch("plugins.image_gen.krea.requests.post", return_value=submit) as mock_post, \ + patch("plugins.image_gen.krea.requests.get", return_value=poll), \ + patch( + "plugins.image_gen.krea.save_url_image", + return_value=Path("/tmp/x.png"), + ), \ + patch("plugins.image_gen.krea.time.sleep"): + KreaImageGenProvider().generate( + prompt="test", + seed=42, + styles=[{"id": "lora-1", "strength": 0.7}], + moodboards=[{"url": "https://x.com/mood.png"}, {"url": "https://x.com/mood2.png"}], + image_style_references=[{"url": f"https://x.com/{i}.png"} for i in range(15)], + creativity="high", + ) + + payload = mock_post.call_args.kwargs["json"] + assert payload["seed"] == 42 + assert payload["styles"] == [{"id": "lora-1", "strength": 0.7}] + assert len(payload["moodboards"]) == 1 # capped at 1 + assert len(payload["image_style_references"]) == 10 # capped at 10 + assert payload["creativity"] == "high" + + def test_unknown_kwargs_ignored(self): + """Forward-compat: unknown kwargs must not break generate().""" + from plugins.image_gen.krea import KreaImageGenProvider + + submit = _submit_response() + poll = _poll_response(_completed_job()) + + with patch("plugins.image_gen.krea.requests.post", return_value=submit), \ + patch("plugins.image_gen.krea.requests.get", return_value=poll), \ + patch( + "plugins.image_gen.krea.save_url_image", + return_value=Path("/tmp/x.png"), + ), \ + patch("plugins.image_gen.krea.time.sleep"): + result = KreaImageGenProvider().generate( + prompt="test", + fictional_param="should be ignored", + num_images=4, + ) + + assert result["success"] is True + + +# --------------------------------------------------------------------------- +# Generate — error paths +# --------------------------------------------------------------------------- + + +class TestGenerateErrors: + def test_submit_http_error(self): + import requests as req_lib + from plugins.image_gen.krea import KreaImageGenProvider + + resp = req_lib.Response() + resp.status_code = 401 + resp._content = b'{"error": {"message": "Invalid API key"}}' + resp.headers["Content-Type"] = "application/json" + resp.raise_for_status = MagicMock( + side_effect=req_lib.HTTPError(response=resp) + ) + + with patch("plugins.image_gen.krea.requests.post", return_value=resp): + result = KreaImageGenProvider().generate(prompt="test") + + assert result["success"] is False + assert result["error_type"] == "api_error" + assert "401" in result["error"] + assert "Invalid API key" in result["error"] + + def test_submit_timeout(self): + import requests as req_lib + from plugins.image_gen.krea import KreaImageGenProvider + + with patch( + "plugins.image_gen.krea.requests.post", side_effect=req_lib.Timeout() + ): + result = KreaImageGenProvider().generate(prompt="test") + + assert result["success"] is False + assert result["error_type"] == "timeout" + + def test_submit_connection_error(self): + import requests as req_lib + from plugins.image_gen.krea import KreaImageGenProvider + + with patch( + "plugins.image_gen.krea.requests.post", + side_effect=req_lib.ConnectionError("dns nope"), + ): + result = KreaImageGenProvider().generate(prompt="test") + + assert result["success"] is False + assert result["error_type"] == "connection_error" + + def test_submit_missing_job_id(self): + from plugins.image_gen.krea import KreaImageGenProvider + + bad_submit = MagicMock() + bad_submit.status_code = 200 + bad_submit.raise_for_status = MagicMock() + bad_submit.json.return_value = {"status": "queued"} + + with patch("plugins.image_gen.krea.requests.post", return_value=bad_submit): + result = KreaImageGenProvider().generate(prompt="test") + + assert result["success"] is False + assert result["error_type"] == "invalid_response" + assert "job_id" in result["error"] + + def test_job_failed(self): + from plugins.image_gen.krea import KreaImageGenProvider + + failed = { + "job_id": "abc", + "status": "failed", + "completed_at": "2026-05-27T00:01:00Z", + "result": {"error": "NSFW content"}, + } + + submit = _submit_response() + with patch("plugins.image_gen.krea.requests.post", return_value=submit), \ + patch( + "plugins.image_gen.krea.requests.get", + return_value=_poll_response(failed), + ), \ + patch("plugins.image_gen.krea.time.sleep"): + result = KreaImageGenProvider().generate(prompt="test") + + assert result["success"] is False + assert result["error_type"] == "api_error" + assert "NSFW" in result["error"] + + def test_job_cancelled(self): + from plugins.image_gen.krea import KreaImageGenProvider + + cancelled = { + "job_id": "abc", + "status": "cancelled", + "completed_at": "2026-05-27T00:01:00Z", + "result": {}, + } + + with patch("plugins.image_gen.krea.requests.post", return_value=_submit_response()), \ + patch( + "plugins.image_gen.krea.requests.get", + return_value=_poll_response(cancelled), + ), \ + patch("plugins.image_gen.krea.time.sleep"): + result = KreaImageGenProvider().generate(prompt="test") + + assert result["success"] is False + assert result["error_type"] == "cancelled" + + def test_completed_but_missing_urls(self): + from plugins.image_gen.krea import KreaImageGenProvider + + completed_empty = { + "job_id": "abc", + "status": "completed", + "completed_at": "2026-05-27T00:01:00Z", + "result": {"urls": []}, + } + + with patch("plugins.image_gen.krea.requests.post", return_value=_submit_response()), \ + patch( + "plugins.image_gen.krea.requests.get", + return_value=_poll_response(completed_empty), + ), \ + patch("plugins.image_gen.krea.time.sleep"): + result = KreaImageGenProvider().generate(prompt="test") + + assert result["success"] is False + assert result["error_type"] == "empty_response" + + def test_url_download_failure_falls_back_to_bare_url(self): + """Mirror of xAI behaviour — if local cache fails, return the URL.""" + import requests as req_lib + from plugins.image_gen.krea import KreaImageGenProvider + + url = "https://krea.cdn/expired-soon.png" + submit = _submit_response() + poll = _poll_response(_completed_job(url)) + + with patch("plugins.image_gen.krea.requests.post", return_value=submit), \ + patch("plugins.image_gen.krea.requests.get", return_value=poll), \ + patch( + "plugins.image_gen.krea.save_url_image", + side_effect=req_lib.HTTPError("404"), + ), \ + patch("plugins.image_gen.krea.time.sleep"): + result = KreaImageGenProvider().generate(prompt="test") + + assert result["success"] is True + assert result["image"] == url + + def test_polling_picks_up_completed_at_with_unknown_status(self): + """``completed_at`` set + unrecognised pending status → still terminal.""" + from plugins.image_gen.krea import KreaImageGenProvider + + # Use a status value that is NOT in our terminal set ("intermediate-complete") + # but with completed_at populated — Krea's spec says completed_at is the + # canonical terminal marker. + oddball = { + "job_id": "abc", + "status": "intermediate-complete", + "completed_at": "2026-05-27T00:01:00Z", + "result": {"urls": ["https://krea.cdn/done.png"]}, + } + + with patch("plugins.image_gen.krea.requests.post", return_value=_submit_response()), \ + patch( + "plugins.image_gen.krea.requests.get", + return_value=_poll_response(oddball), + ), \ + patch( + "plugins.image_gen.krea.save_url_image", + return_value=Path("/tmp/x.png"), + ), \ + patch("plugins.image_gen.krea.time.sleep"): + result = KreaImageGenProvider().generate(prompt="test") + + assert result["success"] is True + + +class TestPollRetryPolicy: + """Polling fail-fast on permanent 4xx, retry on transient 5xx/429.""" + + def _http_error_response(self, status: int): + import requests as req_lib + + resp = req_lib.Response() + resp.status_code = status + resp._content = b'{"error": "boom"}' + resp.headers["Content-Type"] = "application/json" + resp.raise_for_status = MagicMock( + side_effect=req_lib.HTTPError(response=resp) + ) + return resp + + def test_poll_fails_fast_on_401(self): + """Auth failure mid-poll should not wait the 180s deadline.""" + from plugins.image_gen.krea import KreaImageGenProvider + + bad_poll = self._http_error_response(401) + + with patch("plugins.image_gen.krea.requests.post", return_value=_submit_response()), \ + patch("plugins.image_gen.krea.requests.get", return_value=bad_poll) as mock_get, \ + patch("plugins.image_gen.krea.time.sleep"): + result = KreaImageGenProvider().generate(prompt="test") + + assert result["success"] is False + assert result["error_type"] == "api_error" + assert "401" in result["error"] + # One call — no retry on permanent auth failure. + assert mock_get.call_count == 1 + + def test_poll_fails_fast_on_404(self): + """Missing job (404) should surface immediately, not retry for 180s.""" + from plugins.image_gen.krea import KreaImageGenProvider + + bad_poll = self._http_error_response(404) + + with patch("plugins.image_gen.krea.requests.post", return_value=_submit_response()), \ + patch("plugins.image_gen.krea.requests.get", return_value=bad_poll) as mock_get, \ + patch("plugins.image_gen.krea.time.sleep"): + result = KreaImageGenProvider().generate(prompt="test") + + assert result["success"] is False + assert result["error_type"] == "api_error" + assert "404" in result["error"] + assert mock_get.call_count == 1 + + def test_poll_fails_fast_on_403(self): + """Billing/permission failure (403) should not retry.""" + from plugins.image_gen.krea import KreaImageGenProvider + + bad_poll = self._http_error_response(403) + + with patch("plugins.image_gen.krea.requests.post", return_value=_submit_response()), \ + patch("plugins.image_gen.krea.requests.get", return_value=bad_poll) as mock_get, \ + patch("plugins.image_gen.krea.time.sleep"): + result = KreaImageGenProvider().generate(prompt="test") + + assert result["success"] is False + assert mock_get.call_count == 1 + + def test_poll_retries_on_503_then_succeeds(self): + """Transient 5xx should retry and eventually surface a completion.""" + from plugins.image_gen.krea import KreaImageGenProvider + + flaky = self._http_error_response(503) + good = _poll_response(_completed_job("https://krea.cdn/ok.png")) + + with patch("plugins.image_gen.krea.requests.post", return_value=_submit_response()), \ + patch( + "plugins.image_gen.krea.requests.get", + side_effect=[flaky, flaky, good], + ) as mock_get, \ + patch( + "plugins.image_gen.krea.save_url_image", + return_value=Path("/tmp/x.png"), + ), \ + patch("plugins.image_gen.krea.time.sleep"): + result = KreaImageGenProvider().generate(prompt="test") + + assert result["success"] is True + assert mock_get.call_count == 3 + + def test_poll_retries_on_429(self): + """Rate-limit (429) is in the retryable set.""" + from plugins.image_gen.krea import KreaImageGenProvider + + rate_limited = self._http_error_response(429) + good = _poll_response(_completed_job("https://krea.cdn/ok.png")) + + with patch("plugins.image_gen.krea.requests.post", return_value=_submit_response()), \ + patch( + "plugins.image_gen.krea.requests.get", + side_effect=[rate_limited, good], + ) as mock_get, \ + patch( + "plugins.image_gen.krea.save_url_image", + return_value=Path("/tmp/x.png"), + ), \ + patch("plugins.image_gen.krea.time.sleep"): + result = KreaImageGenProvider().generate(prompt="test") + + assert result["success"] is True + assert mock_get.call_count == 2 + + +# --------------------------------------------------------------------------- +# Registration +# --------------------------------------------------------------------------- + + +class TestRegistration: + def test_register(self): + from plugins.image_gen.krea import KreaImageGenProvider, register + + mock_ctx = MagicMock() + register(mock_ctx) + mock_ctx.register_image_gen_provider.assert_called_once() + provider = mock_ctx.register_image_gen_provider.call_args[0][0] + assert isinstance(provider, KreaImageGenProvider) + assert provider.name == "krea"