feat(v3.12.0): media pipeline, AI image generation, capability probe, companion v2.9.0
Three-month batch sync from internal repo (~80 commits) covering Tracks F.5a, F.7e, F.8, F.17, F.18, F.X. WordPress media pipeline - Pillow-based optimization, AI image generation (OpenAI / Stability / Replicate / Google Nano Banana / OpenRouter), chunked + resumable uploads, bulk delete/reassign, idempotent retries. Capability discovery (F.7e) - Per-site credential probe + adapters for WordPress / WooCommerce / Gitea, tier-fit unions granted ∪ roles, capability badge UI with HTMX partial re-check, install hint in every companion-unreachable error. Companion plugin overhaul - Renamed wordpress-plugin/airano-mcp-seo-bridge → wordpress-plugin/airano-mcp-bridge. - Eight new endpoints: /capabilities, /bulk-meta, /export, /cache-purge, /transient-flush, /site-health, /audit-hook, /upload-and-attach. - wp.org Plugin Check pass: i18n, WP_Filesystem, scheme allowlist on audit-hook URL. Other - Gitea ergonomics (F.17): batch files, tree, search, compare, releases, fork. - Opportunistic bcrypt upgrade for legacy SHA-256 admin keys (F.8). - n8n refactor: structured errors, capability probe, missing tools backfilled. - Idempotency-Key dedup for AI media upload retries; WP client fast-fails on unreachable sites. Docs - README + CLAUDE.md drop the fixed "633 tools" claim. The total grows with each release; per-plugin approximations + dashboard-surfaced counts replace it. - Tools/Tests badges removed in favour of "Plugins: 10". Deployment - PyPI mirror chain, optional BUILD_HTTP_PROXY, Alpine→Yandex apk mirror, Debian-slim Plan-B Dockerfile, mirror.gcr.io variant. CI - Black + Ruff clean on Python 3.12; pytest tests/ green. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
33
plugins/ai_image/__init__.py
Normal file
33
plugins/ai_image/__init__.py
Normal file
@@ -0,0 +1,33 @@
|
||||
"""AI image generation provider library (F.5a.4).
|
||||
|
||||
Not a registered MCP plugin — this package exposes provider
|
||||
implementations and a lookup registry consumed by media-upload tools in
|
||||
other plugins (currently ``wordpress_generate_and_upload_image``).
|
||||
|
||||
Typical usage::
|
||||
|
||||
from plugins.ai_image.registry import get_provider
|
||||
|
||||
provider = get_provider("openai")
|
||||
result = await provider.generate(
|
||||
api_key=key, prompt="a red cube", size="1024x1024"
|
||||
)
|
||||
image_bytes, mime, meta = result.bytes, result.mime, result.meta
|
||||
"""
|
||||
|
||||
from plugins.ai_image.providers.base import (
|
||||
BaseImageProvider,
|
||||
GenerationRequest,
|
||||
GenerationResult,
|
||||
ProviderError,
|
||||
)
|
||||
from plugins.ai_image.registry import get_provider, list_providers
|
||||
|
||||
__all__ = [
|
||||
"BaseImageProvider",
|
||||
"GenerationRequest",
|
||||
"GenerationResult",
|
||||
"ProviderError",
|
||||
"get_provider",
|
||||
"list_providers",
|
||||
]
|
||||
1
plugins/ai_image/providers/__init__.py
Normal file
1
plugins/ai_image/providers/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""AI image provider implementations (F.5a.4)."""
|
||||
72
plugins/ai_image/providers/base.py
Normal file
72
plugins/ai_image/providers/base.py
Normal file
@@ -0,0 +1,72 @@
|
||||
"""Base types for AI image providers (F.5a.4).
|
||||
|
||||
Each concrete provider subclasses :class:`BaseImageProvider` and implements
|
||||
``generate()``. Providers return raw image bytes + MIME so the caller (e.g.
|
||||
the WordPress upload pipeline) can reuse the existing raw-upload path.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
|
||||
class ProviderError(Exception):
|
||||
"""Stable-coded error for provider failures.
|
||||
|
||||
Codes follow the F.5a error taxonomy (``PROVIDER_QUOTA``,
|
||||
``PROVIDER_AUTH``, ``PROVIDER_BAD_REQUEST``, ``PROVIDER_UNAVAILABLE``,
|
||||
``PROVIDER_TIMEOUT``). ``details`` is JSON-serialisable.
|
||||
"""
|
||||
|
||||
def __init__(self, code: str, message: str, details: dict[str, Any] | None = None):
|
||||
super().__init__(message)
|
||||
self.code = code
|
||||
self.message = message
|
||||
self.details = details or {}
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {"error_code": self.code, "message": self.message, "details": self.details}
|
||||
|
||||
|
||||
@dataclass
|
||||
class GenerationRequest:
|
||||
"""Normalised generation parameters.
|
||||
|
||||
Not every provider uses every field — unknown fields are ignored.
|
||||
"""
|
||||
|
||||
prompt: str
|
||||
size: str = "1024x1024"
|
||||
quality: str = "standard"
|
||||
model: str | None = None
|
||||
negative_prompt: str | None = None
|
||||
extras: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GenerationResult:
|
||||
"""Raw bytes + metadata returned by a provider."""
|
||||
|
||||
data: bytes
|
||||
mime: str
|
||||
filename: str
|
||||
meta: dict[str, Any] = field(default_factory=dict)
|
||||
cost_usd: float | None = None
|
||||
|
||||
|
||||
class BaseImageProvider(ABC):
|
||||
"""Abstract provider: turns ``(api_key, request)`` into image bytes."""
|
||||
|
||||
name: str = "base"
|
||||
|
||||
@abstractmethod
|
||||
async def generate(self, api_key: str, request: GenerationRequest) -> GenerationResult:
|
||||
"""Call the provider API and return image bytes + metadata.
|
||||
|
||||
Implementations should raise :class:`ProviderError` with a stable
|
||||
code for predictable client handling. Network retries for transient
|
||||
failures (429 / 5xx) should happen inside the implementation.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
184
plugins/ai_image/providers/openai.py
Normal file
184
plugins/ai_image/providers/openai.py
Normal file
@@ -0,0 +1,184 @@
|
||||
"""OpenAI image-generation provider (DALL-E 3 + gpt-image-1).
|
||||
|
||||
DALL-E 3 returns a time-limited URL (~1h TTL). gpt-image-1 can return
|
||||
``b64_json`` directly. In both cases we return the raw bytes immediately
|
||||
to the caller so the 1h URL expiry is never a concern downstream.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
|
||||
from plugins.ai_image.providers.base import (
|
||||
BaseImageProvider,
|
||||
GenerationRequest,
|
||||
GenerationResult,
|
||||
ProviderError,
|
||||
)
|
||||
|
||||
_logger = logging.getLogger("mcphub.ai_image.openai")
|
||||
|
||||
_API_URL = "https://api.openai.com/v1/images/generations"
|
||||
_DEFAULT_MODEL = "dall-e-3"
|
||||
_RETRY_STATUS = {429, 500, 502, 503, 504}
|
||||
_MAX_RETRIES = 3
|
||||
_REQUEST_TIMEOUT = 90
|
||||
|
||||
# Rough per-image pricing (USD) for audit logging. These are documented
|
||||
# public list prices at the time of writing and may drift — treat as
|
||||
# ballpark for cost dashboards, not billing.
|
||||
_COST_TABLE: dict[tuple[str, str, str], float] = {
|
||||
("dall-e-3", "1024x1024", "standard"): 0.040,
|
||||
("dall-e-3", "1024x1024", "hd"): 0.080,
|
||||
("dall-e-3", "1792x1024", "standard"): 0.080,
|
||||
("dall-e-3", "1024x1792", "standard"): 0.080,
|
||||
("dall-e-3", "1792x1024", "hd"): 0.120,
|
||||
("dall-e-3", "1024x1792", "hd"): 0.120,
|
||||
("gpt-image-1", "1024x1024", "low"): 0.011,
|
||||
("gpt-image-1", "1024x1024", "medium"): 0.042,
|
||||
("gpt-image-1", "1024x1024", "high"): 0.167,
|
||||
}
|
||||
|
||||
|
||||
def _estimate_cost(model: str, size: str, quality: str) -> float | None:
|
||||
return _COST_TABLE.get((model, size, quality))
|
||||
|
||||
|
||||
class OpenAIProvider(BaseImageProvider):
|
||||
name = "openai"
|
||||
|
||||
async def generate(self, api_key: str, request: GenerationRequest) -> GenerationResult:
|
||||
if not api_key:
|
||||
raise ProviderError("PROVIDER_AUTH", "OpenAI API key is missing.")
|
||||
|
||||
model = request.model or _DEFAULT_MODEL
|
||||
payload: dict[str, Any] = {
|
||||
"model": model,
|
||||
"prompt": request.prompt,
|
||||
"size": request.size,
|
||||
"n": 1,
|
||||
}
|
||||
if model == "dall-e-3":
|
||||
payload["quality"] = request.quality or "standard"
|
||||
payload["response_format"] = "url"
|
||||
else:
|
||||
payload["response_format"] = "b64_json"
|
||||
if request.quality:
|
||||
payload["quality"] = request.quality
|
||||
|
||||
data, meta = await self._post_with_retry(api_key, payload)
|
||||
|
||||
items = data.get("data") or []
|
||||
if not items:
|
||||
raise ProviderError("PROVIDER_BAD_RESPONSE", "OpenAI returned no images.")
|
||||
first = items[0]
|
||||
|
||||
if first.get("b64_json"):
|
||||
image_bytes = base64.b64decode(first["b64_json"])
|
||||
elif first.get("url"):
|
||||
image_bytes = await _fetch_url(first["url"])
|
||||
else:
|
||||
raise ProviderError(
|
||||
"PROVIDER_BAD_RESPONSE",
|
||||
"OpenAI response contained no b64_json or url.",
|
||||
)
|
||||
|
||||
meta = dict(meta)
|
||||
if first.get("revised_prompt"):
|
||||
meta["revised_prompt"] = first["revised_prompt"]
|
||||
meta["model"] = model
|
||||
meta["size"] = request.size
|
||||
|
||||
return GenerationResult(
|
||||
data=image_bytes,
|
||||
mime="image/png",
|
||||
filename=f"openai-{model}.png",
|
||||
meta=meta,
|
||||
cost_usd=_estimate_cost(model, request.size, request.quality or "standard"),
|
||||
)
|
||||
|
||||
async def _post_with_retry(
|
||||
self, api_key: str, payload: dict[str, Any]
|
||||
) -> tuple[dict[str, Any], dict[str, Any]]:
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
last_error: str = ""
|
||||
delay = 1.0
|
||||
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
for attempt in range(1, _MAX_RETRIES + 1):
|
||||
try:
|
||||
async with session.post(_API_URL, json=payload, headers=headers) as resp:
|
||||
if resp.status == 200:
|
||||
body = await resp.json()
|
||||
return body, {"attempt": attempt, "status": 200}
|
||||
text = await resp.text()
|
||||
last_error = text[:500]
|
||||
if resp.status == 401:
|
||||
raise ProviderError(
|
||||
"PROVIDER_AUTH",
|
||||
"OpenAI rejected the API key (401).",
|
||||
{"body": last_error},
|
||||
)
|
||||
if resp.status == 400:
|
||||
raise ProviderError(
|
||||
"PROVIDER_BAD_REQUEST",
|
||||
f"OpenAI rejected request (400): {last_error}",
|
||||
{"body": last_error},
|
||||
)
|
||||
if resp.status == 429 and attempt < _MAX_RETRIES:
|
||||
_logger.warning("OpenAI 429, retry %d", attempt)
|
||||
await asyncio.sleep(delay)
|
||||
delay *= 2
|
||||
continue
|
||||
if resp.status in _RETRY_STATUS and attempt < _MAX_RETRIES:
|
||||
_logger.warning("OpenAI %d, retry %d", resp.status, attempt)
|
||||
await asyncio.sleep(delay)
|
||||
delay *= 2
|
||||
continue
|
||||
if resp.status == 429:
|
||||
raise ProviderError(
|
||||
"PROVIDER_QUOTA",
|
||||
"OpenAI quota/rate-limit hit after retries.",
|
||||
{"status": 429, "body": last_error},
|
||||
)
|
||||
raise ProviderError(
|
||||
"PROVIDER_UNAVAILABLE",
|
||||
f"OpenAI upstream error HTTP {resp.status}.",
|
||||
{"status": resp.status, "body": last_error},
|
||||
)
|
||||
except TimeoutError as exc:
|
||||
last_error = f"timeout: {exc}"
|
||||
if attempt < _MAX_RETRIES:
|
||||
await asyncio.sleep(delay)
|
||||
delay *= 2
|
||||
continue
|
||||
raise ProviderError(
|
||||
"PROVIDER_TIMEOUT",
|
||||
"OpenAI request timed out after retries.",
|
||||
) from exc
|
||||
raise ProviderError(
|
||||
"PROVIDER_UNAVAILABLE",
|
||||
f"OpenAI call failed after {_MAX_RETRIES} attempts: {last_error}",
|
||||
)
|
||||
|
||||
|
||||
async def _fetch_url(url: str) -> bytes:
|
||||
"""Fetch a DALL-E URL immediately. URLs expire in ~1h, so no caching."""
|
||||
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
async with session.get(url) as resp:
|
||||
if resp.status != 200:
|
||||
raise ProviderError(
|
||||
"PROVIDER_BAD_RESPONSE",
|
||||
f"Failed to fetch OpenAI image URL (HTTP {resp.status}).",
|
||||
{"status": resp.status},
|
||||
)
|
||||
return await resp.read()
|
||||
474
plugins/ai_image/providers/openrouter.py
Normal file
474
plugins/ai_image/providers/openrouter.py
Normal file
@@ -0,0 +1,474 @@
|
||||
"""OpenRouter image-generation provider.
|
||||
|
||||
OpenRouter is an aggregator that routes requests to many model vendors.
|
||||
For image generation, the usable path is their chat-completions endpoint
|
||||
with ``modalities=["image","text"]`` — models like
|
||||
``google/gemini-2.5-flash-image`` (a.k.a. "Nano Banana") return the
|
||||
generated image as a ``data:``-prefixed base64 URL inside the message
|
||||
``images`` array; other models may return ``image_url``-style pointers.
|
||||
|
||||
Why this provider is worth adding:
|
||||
|
||||
* Unlocks Gemini image models without needing a Google Cloud project.
|
||||
* Lets users concentrate AI spend on a single OpenRouter key rather
|
||||
than managing separate OpenAI / Stability / Replicate accounts.
|
||||
* Supported models (as of F.X.fix 2026-04-18):
|
||||
- ``google/gemini-2.5-flash-image`` (default, GA)
|
||||
- ``google/gemini-2.5-flash-image-preview`` (DEPRECATED; 404 on
|
||||
fresh OpenRouter accounts — kept as a recognised alias so we can
|
||||
emit ``PROVIDER_MODEL_DEPRECATED`` with a clear hint)
|
||||
- Any other OpenRouter model that returns image parts in
|
||||
``message.images[]`` — the parser is tolerant of newer entries.
|
||||
|
||||
The default model is picked to cover the most common use case (hero
|
||||
images for WordPress / WooCommerce posts); callers can override via
|
||||
``GenerationRequest.model``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
|
||||
from plugins.ai_image.providers.base import (
|
||||
BaseImageProvider,
|
||||
GenerationRequest,
|
||||
GenerationResult,
|
||||
ProviderError,
|
||||
)
|
||||
|
||||
_logger = logging.getLogger("mcphub.ai_image.openrouter")
|
||||
|
||||
_API_URL = "https://openrouter.ai/api/v1/chat/completions"
|
||||
_MODELS_URL = "https://openrouter.ai/api/v1/models"
|
||||
# GA endpoint. The preview alias (...-image-preview) was deprecated after
|
||||
# Google promoted the model to GA — new OpenRouter accounts get 404 on
|
||||
# the preview ID.
|
||||
_DEFAULT_MODEL = "google/gemini-2.5-flash-image"
|
||||
_DEPRECATED_MODELS: dict[str, str] = {
|
||||
"google/gemini-2.5-flash-image-preview": "google/gemini-2.5-flash-image",
|
||||
}
|
||||
_RETRY_STATUS = {429, 500, 502, 503, 504}
|
||||
_MAX_RETRIES = 3
|
||||
_REQUEST_TIMEOUT = 120
|
||||
|
||||
# USD-per-image price table. Figures are conservative upper bounds from
|
||||
# public OpenRouter pricing (2026-04); operators who negotiate custom
|
||||
# rates can override via ``OPENROUTER_PRICING_OVERRIDE`` (JSON map of
|
||||
# model_id -> price USD).
|
||||
_MODEL_PRICING: dict[str, float] = {
|
||||
"google/gemini-2.5-flash-image": 0.00039,
|
||||
"google/gemini-2.5-flash-image-preview": 0.00039,
|
||||
"google/imagen-3.0-generate": 0.04,
|
||||
"google/imagen-3.0-fast": 0.02,
|
||||
"openai/dall-e-3": 0.04,
|
||||
"openai/dall-e-3-hd": 0.08,
|
||||
"openai/gpt-image-1": 0.017,
|
||||
"black-forest-labs/flux-1.1-pro": 0.04,
|
||||
"black-forest-labs/flux-pro": 0.055,
|
||||
"black-forest-labs/flux-schnell": 0.003,
|
||||
"stability-ai/stable-diffusion-3.5-large": 0.065,
|
||||
}
|
||||
|
||||
|
||||
def _pricing_table() -> dict[str, float]:
|
||||
"""Return the effective price table, merging env override if present."""
|
||||
import json as _json
|
||||
|
||||
override = os.environ.get("OPENROUTER_PRICING_OVERRIDE", "").strip()
|
||||
if not override:
|
||||
return _MODEL_PRICING
|
||||
try:
|
||||
parsed = _json.loads(override)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
_logger.warning("OPENROUTER_PRICING_OVERRIDE is not valid JSON: %s", exc)
|
||||
return _MODEL_PRICING
|
||||
if not isinstance(parsed, dict):
|
||||
return _MODEL_PRICING
|
||||
merged = dict(_MODEL_PRICING)
|
||||
for k, v in parsed.items():
|
||||
try:
|
||||
merged[str(k)] = float(v)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
return merged
|
||||
|
||||
|
||||
def _cost_for(model: str) -> float | None:
|
||||
"""Return USD cost for a known model, else None (logs at debug)."""
|
||||
price = _pricing_table().get(model)
|
||||
if price is None:
|
||||
_logger.debug("openrouter: no pricing entry for model %r", model)
|
||||
return price
|
||||
|
||||
|
||||
# Module-level cache for list_image_models() — 1h TTL is plenty; model
|
||||
# catalog drift is measured in days, not minutes.
|
||||
_MODELS_CACHE_TTL_SECONDS = 3600
|
||||
_models_cache: dict[str, Any] = {"fetched_at": 0.0, "payload": None}
|
||||
|
||||
|
||||
class OpenRouterProvider(BaseImageProvider):
|
||||
name = "openrouter"
|
||||
|
||||
async def generate(self, api_key: str, request: GenerationRequest) -> GenerationResult:
|
||||
if not api_key:
|
||||
raise ProviderError("PROVIDER_AUTH", "OpenRouter API key is missing.")
|
||||
|
||||
requested_model = request.model or _DEFAULT_MODEL
|
||||
replacement = _DEPRECATED_MODELS.get(requested_model)
|
||||
if replacement is not None:
|
||||
raise ProviderError(
|
||||
"PROVIDER_MODEL_DEPRECATED",
|
||||
(
|
||||
f"OpenRouter model '{requested_model}' is deprecated. "
|
||||
f"Use '{replacement}' instead."
|
||||
),
|
||||
{"requested_model": requested_model, "replacement_model": replacement},
|
||||
)
|
||||
model = requested_model
|
||||
|
||||
# Chat-completions shape with image modality declared. The
|
||||
# prompt is sent as the single user message; size hints are
|
||||
# advisory — Gemini picks an output size internally.
|
||||
user_content: str = request.prompt
|
||||
if request.negative_prompt:
|
||||
user_content = f"{user_content}\n\nAvoid: {request.negative_prompt}"
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": user_content}],
|
||||
"modalities": ["image", "text"],
|
||||
}
|
||||
|
||||
body = await self._post_with_retry(api_key, payload)
|
||||
|
||||
image_url = _extract_image_url(body)
|
||||
if not image_url:
|
||||
raise ProviderError(
|
||||
"PROVIDER_BAD_RESPONSE",
|
||||
"OpenRouter response contained no image data. Is the "
|
||||
f"model '{model}' configured for image output?",
|
||||
{"model": model},
|
||||
)
|
||||
|
||||
image_bytes, mime = await _materialise_image(image_url)
|
||||
if not image_bytes:
|
||||
raise ProviderError(
|
||||
"PROVIDER_BAD_RESPONSE",
|
||||
"OpenRouter image URL could not be fetched / decoded.",
|
||||
{"model": model},
|
||||
)
|
||||
|
||||
# Derive a useful filename suffix from the mime type.
|
||||
ext = {"image/png": "png", "image/jpeg": "jpg", "image/webp": "webp"}.get(mime, "png")
|
||||
meta: dict[str, Any] = {"model": model, "size": request.size}
|
||||
# Attach OpenRouter usage if present so audit logs can correlate.
|
||||
usage = body.get("usage") or {}
|
||||
if isinstance(usage, dict):
|
||||
for k in ("prompt_tokens", "completion_tokens", "total_tokens"):
|
||||
if k in usage:
|
||||
meta[k] = usage[k]
|
||||
|
||||
return GenerationResult(
|
||||
data=image_bytes,
|
||||
mime=mime,
|
||||
filename=f"openrouter-{model.replace('/', '-')}.{ext}",
|
||||
meta=meta,
|
||||
cost_usd=_cost_for(model),
|
||||
)
|
||||
|
||||
async def _post_with_retry(self, api_key: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
# OpenRouter recommends identifying the calling app for
|
||||
# observability. Using a stable string so admins can spot
|
||||
# traffic from MCPHub in their OpenRouter dashboard.
|
||||
"HTTP-Referer": "https://github.com/airano-ir/mcphub",
|
||||
"X-Title": "MCPHub",
|
||||
}
|
||||
last_error: str = ""
|
||||
delay = 1.0
|
||||
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
for attempt in range(1, _MAX_RETRIES + 1):
|
||||
try:
|
||||
async with session.post(_API_URL, json=payload, headers=headers) as resp:
|
||||
if resp.status == 200:
|
||||
return await resp.json()
|
||||
text = await resp.text()
|
||||
last_error = text[:500]
|
||||
if resp.status == 401:
|
||||
raise ProviderError(
|
||||
"PROVIDER_AUTH",
|
||||
"OpenRouter rejected the API key (401).",
|
||||
{"body": last_error},
|
||||
)
|
||||
if resp.status == 400:
|
||||
raise ProviderError(
|
||||
"PROVIDER_BAD_REQUEST",
|
||||
f"OpenRouter rejected request (400): {last_error}",
|
||||
{"body": last_error},
|
||||
)
|
||||
if resp.status in _RETRY_STATUS and attempt < _MAX_RETRIES:
|
||||
_logger.warning(
|
||||
"OpenRouter %d, retry %d/%d", resp.status, attempt, _MAX_RETRIES
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
delay *= 2
|
||||
continue
|
||||
if resp.status == 429:
|
||||
raise ProviderError(
|
||||
"PROVIDER_QUOTA",
|
||||
"OpenRouter quota/rate-limit hit after retries.",
|
||||
{"status": 429, "body": last_error},
|
||||
)
|
||||
raise ProviderError(
|
||||
"PROVIDER_UNAVAILABLE",
|
||||
f"OpenRouter upstream error HTTP {resp.status}.",
|
||||
{"status": resp.status, "body": last_error},
|
||||
)
|
||||
except TimeoutError as exc:
|
||||
last_error = f"timeout: {exc}"
|
||||
if attempt < _MAX_RETRIES:
|
||||
await asyncio.sleep(delay)
|
||||
delay *= 2
|
||||
continue
|
||||
raise ProviderError(
|
||||
"PROVIDER_TIMEOUT",
|
||||
"OpenRouter request timed out after retries.",
|
||||
) from exc
|
||||
raise ProviderError(
|
||||
"PROVIDER_UNAVAILABLE",
|
||||
f"OpenRouter call failed after {_MAX_RETRIES} attempts: {last_error}",
|
||||
)
|
||||
|
||||
async def list_image_models(
|
||||
self, api_key: str | None = None, *, force: bool = False
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Return OpenRouter catalog entries whose modality includes image.
|
||||
|
||||
The catalog drifts slowly (days) so we cache the whole filtered
|
||||
list at module scope for 1h. ``api_key`` is optional — the
|
||||
``/v1/models`` endpoint is public, but authenticated callers see
|
||||
per-account availability flags. On any upstream error the
|
||||
previous cache (if any) is returned; otherwise an empty list.
|
||||
"""
|
||||
now = time.time()
|
||||
payload = _models_cache.get("payload")
|
||||
fetched_at = _models_cache.get("fetched_at", 0.0)
|
||||
if not force and payload is not None and (now - fetched_at) < _MODELS_CACHE_TTL_SECONDS:
|
||||
return list(payload)
|
||||
|
||||
headers = {
|
||||
"HTTP-Referer": "https://github.com/airano-ir/mcphub",
|
||||
"X-Title": "MCPHub",
|
||||
}
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)
|
||||
try:
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
async with session.get(_MODELS_URL, headers=headers) as resp:
|
||||
if resp.status != 200:
|
||||
_logger.warning("openrouter /v1/models HTTP %s", resp.status)
|
||||
return list(payload or [])
|
||||
body = await resp.json()
|
||||
except Exception as exc: # noqa: BLE001
|
||||
_logger.warning("openrouter /v1/models fetch failed: %s", exc)
|
||||
return list(payload or [])
|
||||
|
||||
raw_models = body.get("data") or []
|
||||
pricing = _pricing_table()
|
||||
filtered: list[dict[str, Any]] = []
|
||||
for item in raw_models:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
if not _model_is_image(item):
|
||||
continue
|
||||
model_id = str(item.get("id") or "")
|
||||
if not model_id or model_id in _DEPRECATED_MODELS:
|
||||
continue
|
||||
filtered.append(
|
||||
{
|
||||
"id": model_id,
|
||||
"name": item.get("name") or model_id,
|
||||
"description": (item.get("description") or "")[:400],
|
||||
"context_length": item.get("context_length"),
|
||||
"input_modalities": _modalities_of(item, "input"),
|
||||
"output_modalities": _modalities_of(item, "output"),
|
||||
"price_per_image_usd": pricing.get(model_id),
|
||||
}
|
||||
)
|
||||
|
||||
filtered.sort(key=lambda m: m["id"])
|
||||
_models_cache["fetched_at"] = now
|
||||
_models_cache["payload"] = filtered
|
||||
return list(filtered)
|
||||
|
||||
|
||||
def _modalities_of(item: dict[str, Any], side: str) -> list[str]:
|
||||
"""Extract input/output modality list from an OpenRouter model entry."""
|
||||
arch = item.get("architecture") or {}
|
||||
if isinstance(arch, dict):
|
||||
key = "input_modalities" if side == "input" else "output_modalities"
|
||||
mods = arch.get(key)
|
||||
if isinstance(mods, list):
|
||||
return [str(m) for m in mods if isinstance(m, str)]
|
||||
modality = arch.get("modality")
|
||||
if isinstance(modality, str):
|
||||
parts = modality.split("->")
|
||||
if len(parts) == 2:
|
||||
idx = 0 if side == "input" else 1
|
||||
return [p.strip() for p in parts[idx].split("+") if p.strip()]
|
||||
return []
|
||||
|
||||
|
||||
def _model_is_image(item: dict[str, Any]) -> bool:
|
||||
"""True when an OpenRouter model entry emits image output."""
|
||||
out = _modalities_of(item, "output")
|
||||
if any("image" in m.lower() for m in out):
|
||||
return True
|
||||
# Legacy catalog rows without structured modalities — fall back to a
|
||||
# name/description heuristic so we don't silently drop valid models.
|
||||
blob = " ".join(str(item.get(k, "") or "") for k in ("id", "name", "description")).lower()
|
||||
return "image" in blob and "gen" in blob
|
||||
|
||||
|
||||
def _extract_image_url(body: dict[str, Any]) -> str | None:
|
||||
"""Pull the first image URL out of an OpenRouter chat response.
|
||||
|
||||
The wire format varies by model vendor. We recognise:
|
||||
|
||||
* ``choices[0].message.images[i].image_url.url`` — a string that
|
||||
may be a ``data:image/...;base64,...`` URL or an https URL.
|
||||
* ``choices[0].message.content`` when it is a list of parts, each
|
||||
of shape ``{type: "image_url", image_url: {url: ...}}`` —
|
||||
OpenAI-compatible multimodal reply shape.
|
||||
|
||||
Returns None if no usable reference was found.
|
||||
"""
|
||||
choices = body.get("choices") or []
|
||||
if not isinstance(choices, list) or not choices:
|
||||
return None
|
||||
message = (choices[0] or {}).get("message") or {}
|
||||
|
||||
# Shape 1: message.images (Gemini via OpenRouter).
|
||||
for entry in message.get("images") or []:
|
||||
url = _image_url_from_entry(entry)
|
||||
if url:
|
||||
return url
|
||||
|
||||
# Shape 2: message.content as list of parts.
|
||||
content = message.get("content")
|
||||
if isinstance(content, list):
|
||||
for part in content:
|
||||
if not isinstance(part, dict):
|
||||
continue
|
||||
if part.get("type") == "image_url":
|
||||
url = _image_url_from_entry(part)
|
||||
if url:
|
||||
return url
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _image_url_from_entry(entry: Any) -> str | None:
|
||||
"""Handle both ``{image_url: {url: ...}}`` and ``{image_url: "..."}``."""
|
||||
if not isinstance(entry, dict):
|
||||
return None
|
||||
iu = entry.get("image_url")
|
||||
if isinstance(iu, str):
|
||||
return iu
|
||||
if isinstance(iu, dict):
|
||||
u = iu.get("url")
|
||||
return u if isinstance(u, str) else None
|
||||
# Some providers emit a bare "url" key.
|
||||
u = entry.get("url")
|
||||
return u if isinstance(u, str) else None
|
||||
|
||||
|
||||
async def _materialise_image(url: str) -> tuple[bytes, str]:
|
||||
"""Turn an image URL (data: or https) into raw bytes + MIME.
|
||||
|
||||
Returns ``(b"", "")`` on any failure so the caller can emit a
|
||||
uniform ``PROVIDER_BAD_RESPONSE`` error.
|
||||
"""
|
||||
if url.startswith("data:"):
|
||||
try:
|
||||
header, payload = url.split(",", 1)
|
||||
except ValueError:
|
||||
return b"", ""
|
||||
mime = "image/png"
|
||||
rest = header[len("data:") :] if header.startswith("data:") else header
|
||||
if ";" in rest:
|
||||
mime = rest.split(";", 1)[0] or mime
|
||||
elif rest:
|
||||
mime = rest
|
||||
if "base64" in header:
|
||||
try:
|
||||
return base64.b64decode(payload), mime
|
||||
except Exception:
|
||||
return b"", mime
|
||||
return payload.encode("latin-1", errors="ignore"), mime
|
||||
|
||||
if url.startswith("http://") or url.startswith("https://"):
|
||||
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
async with session.get(url) as resp:
|
||||
if resp.status != 200:
|
||||
return b"", ""
|
||||
mime = resp.headers.get("Content-Type", "image/png").split(";", 1)[0].strip()
|
||||
return await resp.read(), mime or "image/png"
|
||||
|
||||
return b"", ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Starlette handler — GET /api/providers/openrouter/models
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def api_openrouter_models(request: Any) -> Any:
|
||||
"""Return the cached OpenRouter image-model catalog.
|
||||
|
||||
Auth: same OAuth user session guard as the other /api/ endpoints.
|
||||
Query params:
|
||||
* ``force=1`` — bypass the 1h module cache.
|
||||
* ``site_id=<uuid>`` — if present, the user's OpenRouter key for
|
||||
that site is used as bearer so availability flags reflect the
|
||||
operator's account; otherwise the endpoint is fetched unauthed.
|
||||
"""
|
||||
from starlette.responses import JSONResponse
|
||||
|
||||
from core.dashboard.routes import _require_user_session
|
||||
|
||||
user_session, redirect = _require_user_session(request)
|
||||
if redirect or user_session is None:
|
||||
return JSONResponse({"ok": False, "error": "unauthorized"}, status_code=401)
|
||||
|
||||
force = request.query_params.get("force") in {"1", "true", "True"}
|
||||
api_key: str | None = None
|
||||
site_id = (request.query_params.get("site_id") or "").strip()
|
||||
if site_id:
|
||||
try:
|
||||
from core.database import get_database
|
||||
from core.site_api import get_site_provider_key
|
||||
|
||||
db = get_database()
|
||||
site = await db.get_site(site_id, user_session["user_id"])
|
||||
if site is not None:
|
||||
api_key = await get_site_provider_key(site_id, "openrouter")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
_logger.debug("openrouter models: site key lookup skipped: %s", exc)
|
||||
|
||||
provider = OpenRouterProvider()
|
||||
models = await provider.list_image_models(api_key=api_key, force=force)
|
||||
return JSONResponse({"ok": True, "provider": "openrouter", "models": models})
|
||||
204
plugins/ai_image/providers/replicate.py
Normal file
204
plugins/ai_image/providers/replicate.py
Normal file
@@ -0,0 +1,204 @@
|
||||
"""Replicate provider (Flux family + other image models).
|
||||
|
||||
Replicate runs predictions asynchronously: POST to ``/v1/predictions``
|
||||
returns a job id, and we poll ``/v1/predictions/{id}`` until status is
|
||||
``succeeded`` or ``failed``. The final output is a URL (or list of URLs)
|
||||
that must be fetched to get bytes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
|
||||
from plugins.ai_image.providers.base import (
|
||||
BaseImageProvider,
|
||||
GenerationRequest,
|
||||
GenerationResult,
|
||||
ProviderError,
|
||||
)
|
||||
|
||||
_logger = logging.getLogger("mcphub.ai_image.replicate")
|
||||
|
||||
_PREDICTIONS_URL = "https://api.replicate.com/v1/predictions"
|
||||
_DEFAULT_MODEL = "black-forest-labs/flux-schnell"
|
||||
_MAX_RETRIES = 3
|
||||
_POLL_INTERVAL = 2.0
|
||||
_POLL_TIMEOUT = 180
|
||||
_REQUEST_TIMEOUT = 60
|
||||
_RETRY_STATUS = {429, 500, 502, 503, 504}
|
||||
|
||||
_COST_TABLE: dict[str, float] = {
|
||||
"black-forest-labs/flux-schnell": 0.003,
|
||||
"black-forest-labs/flux-dev": 0.025,
|
||||
"black-forest-labs/flux-pro": 0.055,
|
||||
}
|
||||
|
||||
|
||||
def _aspect_from_size(size: str) -> str:
|
||||
"""Replicate's Flux models accept an aspect_ratio enum, not WxH."""
|
||||
mapping = {
|
||||
"1024x1024": "1:1",
|
||||
"1344x768": "16:9",
|
||||
"768x1344": "9:16",
|
||||
"1152x896": "3:2",
|
||||
"896x1152": "2:3",
|
||||
}
|
||||
return mapping.get(size, "1:1")
|
||||
|
||||
|
||||
class ReplicateProvider(BaseImageProvider):
|
||||
name = "replicate"
|
||||
|
||||
async def generate(self, api_key: str, request: GenerationRequest) -> GenerationResult:
|
||||
if not api_key:
|
||||
raise ProviderError("PROVIDER_AUTH", "Replicate API token is missing.")
|
||||
|
||||
model = request.model or _DEFAULT_MODEL
|
||||
payload: dict[str, Any] = {
|
||||
"model": model,
|
||||
"input": {
|
||||
"prompt": request.prompt,
|
||||
"aspect_ratio": _aspect_from_size(request.size),
|
||||
**(request.extras or {}),
|
||||
},
|
||||
}
|
||||
if request.negative_prompt:
|
||||
payload["input"]["negative_prompt"] = request.negative_prompt
|
||||
|
||||
prediction = await self._create_prediction(api_key, payload)
|
||||
final = await self._poll_until_done(api_key, prediction)
|
||||
|
||||
output = final.get("output")
|
||||
url = _first_url(output)
|
||||
if not url:
|
||||
raise ProviderError(
|
||||
"PROVIDER_BAD_RESPONSE",
|
||||
"Replicate prediction finished without an image URL.",
|
||||
{"prediction": final},
|
||||
)
|
||||
image_bytes = await _fetch_url(url)
|
||||
|
||||
meta = {
|
||||
"model": model,
|
||||
"prediction_id": final.get("id"),
|
||||
"status": final.get("status"),
|
||||
"metrics": final.get("metrics", {}),
|
||||
}
|
||||
return GenerationResult(
|
||||
data=image_bytes,
|
||||
mime="image/webp" if url.endswith(".webp") else "image/png",
|
||||
filename=f"replicate-{model.split('/')[-1]}.png",
|
||||
meta=meta,
|
||||
cost_usd=_COST_TABLE.get(model),
|
||||
)
|
||||
|
||||
async def _create_prediction(self, api_key: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
"Prefer": "wait=0",
|
||||
}
|
||||
delay = 1.0
|
||||
last_error = ""
|
||||
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
for attempt in range(1, _MAX_RETRIES + 1):
|
||||
async with session.post(_PREDICTIONS_URL, json=payload, headers=headers) as resp:
|
||||
if resp.status in (200, 201):
|
||||
return await resp.json()
|
||||
text = await resp.text()
|
||||
last_error = text[:500]
|
||||
if resp.status in (401, 403):
|
||||
raise ProviderError(
|
||||
"PROVIDER_AUTH",
|
||||
f"Replicate rejected auth ({resp.status}).",
|
||||
{"body": last_error},
|
||||
)
|
||||
if resp.status == 422:
|
||||
raise ProviderError(
|
||||
"PROVIDER_BAD_REQUEST",
|
||||
f"Replicate rejected request: {last_error}",
|
||||
{"body": last_error},
|
||||
)
|
||||
if resp.status in _RETRY_STATUS and attempt < _MAX_RETRIES:
|
||||
_logger.warning("Replicate %d, retry %d", resp.status, attempt)
|
||||
await asyncio.sleep(delay)
|
||||
delay *= 2
|
||||
continue
|
||||
if resp.status == 429:
|
||||
raise ProviderError(
|
||||
"PROVIDER_QUOTA",
|
||||
"Replicate rate-limit hit.",
|
||||
{"status": 429, "body": last_error},
|
||||
)
|
||||
raise ProviderError(
|
||||
"PROVIDER_UNAVAILABLE",
|
||||
f"Replicate HTTP {resp.status}.",
|
||||
{"status": resp.status, "body": last_error},
|
||||
)
|
||||
raise ProviderError(
|
||||
"PROVIDER_UNAVAILABLE",
|
||||
f"Replicate create_prediction failed after retries: {last_error}",
|
||||
)
|
||||
|
||||
async def _poll_until_done(self, api_key: str, prediction: dict[str, Any]) -> dict[str, Any]:
|
||||
headers = {"Authorization": f"Bearer {api_key}"}
|
||||
poll_url = (prediction.get("urls") or {}).get("get") or (
|
||||
f"{_PREDICTIONS_URL}/{prediction.get('id')}"
|
||||
)
|
||||
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)
|
||||
deadline = asyncio.get_event_loop().time() + _POLL_TIMEOUT
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
while True:
|
||||
if asyncio.get_event_loop().time() > deadline:
|
||||
raise ProviderError(
|
||||
"PROVIDER_TIMEOUT",
|
||||
"Replicate prediction did not finish within timeout.",
|
||||
{"prediction_id": prediction.get("id")},
|
||||
)
|
||||
async with session.get(poll_url, headers=headers) as resp:
|
||||
if resp.status != 200:
|
||||
text = await resp.text()
|
||||
raise ProviderError(
|
||||
"PROVIDER_UNAVAILABLE",
|
||||
f"Replicate poll HTTP {resp.status}.",
|
||||
{"body": text[:500]},
|
||||
)
|
||||
body = await resp.json()
|
||||
status = body.get("status")
|
||||
if status in ("succeeded",):
|
||||
return body
|
||||
if status in ("failed", "canceled"):
|
||||
raise ProviderError(
|
||||
"PROVIDER_BAD_RESPONSE",
|
||||
f"Replicate prediction {status}: {body.get('error')}",
|
||||
{"prediction": body},
|
||||
)
|
||||
await asyncio.sleep(_POLL_INTERVAL)
|
||||
|
||||
|
||||
def _first_url(output: Any) -> str | None:
|
||||
if isinstance(output, str):
|
||||
return output
|
||||
if isinstance(output, list) and output:
|
||||
first = output[0]
|
||||
if isinstance(first, str):
|
||||
return first
|
||||
return None
|
||||
|
||||
|
||||
async def _fetch_url(url: str) -> bytes:
|
||||
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
async with session.get(url) as resp:
|
||||
if resp.status != 200:
|
||||
raise ProviderError(
|
||||
"PROVIDER_BAD_RESPONSE",
|
||||
f"Failed to fetch Replicate output (HTTP {resp.status}).",
|
||||
{"status": resp.status},
|
||||
)
|
||||
return await resp.read()
|
||||
154
plugins/ai_image/providers/stability.py
Normal file
154
plugins/ai_image/providers/stability.py
Normal file
@@ -0,0 +1,154 @@
|
||||
"""Stability AI image generation provider (Stable Image Core / Ultra).
|
||||
|
||||
Uses the v2beta generate endpoint with ``Accept: image/*`` to get raw
|
||||
bytes back directly (no JSON indirection).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
|
||||
from plugins.ai_image.providers.base import (
|
||||
BaseImageProvider,
|
||||
GenerationRequest,
|
||||
GenerationResult,
|
||||
ProviderError,
|
||||
)
|
||||
|
||||
_logger = logging.getLogger("mcphub.ai_image.stability")
|
||||
|
||||
_MODEL_ENDPOINTS = {
|
||||
"core": "https://api.stability.ai/v2beta/stable-image/generate/core",
|
||||
"ultra": "https://api.stability.ai/v2beta/stable-image/generate/ultra",
|
||||
"sd3": "https://api.stability.ai/v2beta/stable-image/generate/sd3",
|
||||
}
|
||||
_DEFAULT_MODEL = "core"
|
||||
_MAX_RETRIES = 3
|
||||
_REQUEST_TIMEOUT = 120
|
||||
_RETRY_STATUS = {429, 500, 502, 503, 504}
|
||||
|
||||
_COST_TABLE: dict[str, float] = {
|
||||
"core": 0.03,
|
||||
"ultra": 0.08,
|
||||
"sd3": 0.065,
|
||||
}
|
||||
|
||||
|
||||
def _size_to_aspect(size: str) -> str:
|
||||
"""Map a WxH size string to the Stability aspect_ratio enum."""
|
||||
mapping = {
|
||||
"1024x1024": "1:1",
|
||||
"1152x896": "21:9",
|
||||
"1216x832": "3:2",
|
||||
"1344x768": "16:9",
|
||||
"768x1344": "9:16",
|
||||
"832x1216": "2:3",
|
||||
"1024x576": "16:9",
|
||||
"512x512": "1:1",
|
||||
}
|
||||
return mapping.get(size, "1:1")
|
||||
|
||||
|
||||
class StabilityProvider(BaseImageProvider):
|
||||
name = "stability"
|
||||
|
||||
async def generate(self, api_key: str, request: GenerationRequest) -> GenerationResult:
|
||||
if not api_key:
|
||||
raise ProviderError("PROVIDER_AUTH", "Stability API key is missing.")
|
||||
|
||||
model = (request.model or _DEFAULT_MODEL).lower()
|
||||
endpoint = _MODEL_ENDPOINTS.get(model)
|
||||
if endpoint is None:
|
||||
raise ProviderError(
|
||||
"PROVIDER_BAD_REQUEST",
|
||||
f"Unknown Stability model '{model}'. " f"Allowed: {', '.join(_MODEL_ENDPOINTS)}.",
|
||||
)
|
||||
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("prompt", request.prompt)
|
||||
form.add_field("output_format", "png")
|
||||
form.add_field("aspect_ratio", _size_to_aspect(request.size))
|
||||
if request.negative_prompt:
|
||||
form.add_field("negative_prompt", request.negative_prompt)
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Accept": "image/*",
|
||||
}
|
||||
data, meta = await self._post_with_retry(endpoint, form, headers)
|
||||
|
||||
meta = dict(meta)
|
||||
meta["model"] = model
|
||||
meta["aspect_ratio"] = _size_to_aspect(request.size)
|
||||
return GenerationResult(
|
||||
data=data,
|
||||
mime="image/png",
|
||||
filename=f"stability-{model}.png",
|
||||
meta=meta,
|
||||
cost_usd=_COST_TABLE.get(model),
|
||||
)
|
||||
|
||||
async def _post_with_retry(
|
||||
self,
|
||||
endpoint: str,
|
||||
form: aiohttp.FormData,
|
||||
headers: dict[str, str],
|
||||
) -> tuple[bytes, dict[str, Any]]:
|
||||
last_error: str = ""
|
||||
delay = 1.0
|
||||
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
for attempt in range(1, _MAX_RETRIES + 1):
|
||||
try:
|
||||
async with session.post(endpoint, data=form, headers=headers) as resp:
|
||||
if resp.status == 200:
|
||||
body = await resp.read()
|
||||
return body, {"attempt": attempt, "status": 200}
|
||||
text = await resp.text()
|
||||
last_error = text[:500]
|
||||
if resp.status == 401 or resp.status == 403:
|
||||
raise ProviderError(
|
||||
"PROVIDER_AUTH",
|
||||
f"Stability rejected auth ({resp.status}).",
|
||||
{"body": last_error},
|
||||
)
|
||||
if resp.status == 400:
|
||||
raise ProviderError(
|
||||
"PROVIDER_BAD_REQUEST",
|
||||
f"Stability rejected request: {last_error}",
|
||||
{"body": last_error},
|
||||
)
|
||||
if resp.status in _RETRY_STATUS and attempt < _MAX_RETRIES:
|
||||
_logger.warning("Stability %d, retry %d", resp.status, attempt)
|
||||
await asyncio.sleep(delay)
|
||||
delay *= 2
|
||||
continue
|
||||
if resp.status == 429:
|
||||
raise ProviderError(
|
||||
"PROVIDER_QUOTA",
|
||||
"Stability quota/rate-limit hit.",
|
||||
{"status": 429, "body": last_error},
|
||||
)
|
||||
raise ProviderError(
|
||||
"PROVIDER_UNAVAILABLE",
|
||||
f"Stability upstream error HTTP {resp.status}.",
|
||||
{"status": resp.status, "body": last_error},
|
||||
)
|
||||
except TimeoutError as exc:
|
||||
last_error = f"timeout: {exc}"
|
||||
if attempt < _MAX_RETRIES:
|
||||
await asyncio.sleep(delay)
|
||||
delay *= 2
|
||||
continue
|
||||
raise ProviderError(
|
||||
"PROVIDER_TIMEOUT",
|
||||
"Stability request timed out after retries.",
|
||||
) from exc
|
||||
raise ProviderError(
|
||||
"PROVIDER_UNAVAILABLE",
|
||||
f"Stability call failed after {_MAX_RETRIES} attempts: {last_error}",
|
||||
)
|
||||
36
plugins/ai_image/registry.py
Normal file
36
plugins/ai_image/registry.py
Normal file
@@ -0,0 +1,36 @@
|
||||
"""Provider registry for AI image generation (F.5a.4).
|
||||
|
||||
Lookup by name; raises :class:`ProviderError` with code
|
||||
``PROVIDER_UNKNOWN`` if the caller supplies a non-existent provider.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from plugins.ai_image.providers.base import BaseImageProvider, ProviderError
|
||||
from plugins.ai_image.providers.openai import OpenAIProvider
|
||||
from plugins.ai_image.providers.openrouter import OpenRouterProvider
|
||||
from plugins.ai_image.providers.replicate import ReplicateProvider
|
||||
from plugins.ai_image.providers.stability import StabilityProvider
|
||||
|
||||
_PROVIDERS: dict[str, BaseImageProvider] = {
|
||||
"openai": OpenAIProvider(),
|
||||
"stability": StabilityProvider(),
|
||||
"replicate": ReplicateProvider(),
|
||||
"openrouter": OpenRouterProvider(),
|
||||
}
|
||||
|
||||
|
||||
def get_provider(name: str) -> BaseImageProvider:
|
||||
"""Return the registered provider singleton by name."""
|
||||
try:
|
||||
return _PROVIDERS[name]
|
||||
except KeyError as exc:
|
||||
raise ProviderError(
|
||||
"PROVIDER_UNKNOWN",
|
||||
f"Unknown provider '{name}'. Allowed: {', '.join(_PROVIDERS)}.",
|
||||
) from exc
|
||||
|
||||
|
||||
def list_providers() -> list[str]:
|
||||
"""Return all registered provider names in registration order."""
|
||||
return list(_PROVIDERS.keys())
|
||||
Reference in New Issue
Block a user