mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-05-29 06:31:32 +00:00
feat(image_gen): add Krea provider plugin (Krea 2 Medium + Large) (#33236)
* feat(image_gen): add Krea provider plugin (Krea 2 Medium + Large)
New built-in image_gen backend wrapping Krea's Krea 2 foundation
image model family. Auto-discovered like the other image_gen plugins
and appears in 'hermes tools' → Image Generation → Krea.
Krea's API is asynchronous — submit returns a job_id, poll /jobs/{id}
until terminal. The provider hides that behind the synchronous
ImageGenProvider.generate() contract: submit, poll every 2s with
light backoff (max 5s), 3-minute ceiling matching Krea's hosted-tool
timeout. Result URL is materialised to $HERMES_HOME/cache/images/
to avoid CDN-expiry 404s downstream (same fix as xAI #26942).
Models:
- krea-2-medium (default — Krea's 'start here' recommendation)
- krea-2-large
Aspect ratios map landscape→16:9, square→1:1, portrait→9:16.
Resolution: 1K (Krea's only current option).
Kwarg passthrough: seed, creativity (raw/low/medium/high), styles,
image_style_references (capped 10), moodboards (capped 1) — matches
Krea's per-request limits. Unknown kwargs are ignored.
Config knobs (config.yaml):
image_gen.provider: krea
image_gen.krea.model: krea-2-medium | krea-2-large
image_gen.krea.creativity: raw | low | medium | high
Env overrides: KREA_API_KEY (required), KREA_IMAGE_MODEL.
KREA_API_KEY is registered in OPTIONAL_ENV_VARS so 'hermes setup'
prompts for it.
31 new tests; image_gen suite + picker + tools_config: 211/211.
* fix(image_gen/krea): address review feedback
- Update KREA_API_KEY setup URL to the canonical token-creation page
(https://www.krea.ai/app/api/tokens). The previous URL returned 404.
- Fail fast on non-retryable HTTP statuses during poll. The previous
loop retried every HTTPError for the full 180s deadline, so an auth
(401), billing (402), forbidden (403), or not-found (404) response
would make image_generate hang for three minutes. Only retry
transient statuses (408/409/425/429/5xx); surface everything else
immediately.
- Add 5 tests covering fail-fast on 401/403/404 and retry on 429/503.
* fix(krea): point users at the real API token dashboard URL
Three call sites linked users to dashboard pages that don't exist:
- hermes_cli/config.py: https://www.krea.ai/app/api/tokens
- plugins/image_gen/krea/__init__.py get_setup_schema: https://www.krea.ai/api-keys
- plugins/image_gen/krea/__init__.py auth_required error: https://www.krea.ai/api-keys
Per Krea's own docs (https://docs.krea.ai/developers/api-keys-and-billing),
the real dashboard URL is https://www.krea.ai/settings/api-tokens. All three
sites now point there.
This commit is contained in:
parent
eccbbe4b1b
commit
9919caff46
4 changed files with 1188 additions and 0 deletions
|
|
@ -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)",
|
||||
|
|
|
|||
548
plugins/image_gen/krea/__init__.py
Normal file
548
plugins/image_gen/krea/__init__.py
Normal file
|
|
@ -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())
|
||||
7
plugins/image_gen/krea/plugin.yaml
Normal file
7
plugins/image_gen/krea/plugin.yaml
Normal file
|
|
@ -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
|
||||
625
tests/plugins/image_gen/test_krea_provider.py
Normal file
625
tests/plugins/image_gen/test_krea_provider.py
Normal file
|
|
@ -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"
|
||||
Loading…
Add table
Add a link
Reference in a new issue