routerai-image-mcp: MCP-сервер генерации картинок через RouterAI
- server.py: MCP-сервер (stdio + streamable HTTP), инструменты generate_image / list_image_models - два бэкенда: chat/completions (Gemini Image, понимает правку по референсу) и images/generations (FLUX, GPT-Image) - ключ RouterAI берётся из env или ~/.hermes/models.yaml автоматически - retry на 429/5xx, картинка сохраняется файлом (~/.hermes/image_cache/mcp-gen) - test_client.py / test_edit.py: проверка по протоколу, включая image-to-image - README: установка, настройка, подключение к Hermes/opencode, питфоллы
This commit is contained in:
@@ -0,0 +1,289 @@
|
||||
#!/usr/bin/env python3
|
||||
"""RouterAI Image MCP server.
|
||||
|
||||
Даёт ЛЮБОЙ модели (в т.ч. той, что не умеет рисовать) инструмент
|
||||
`generate_image` — генерация/редактирование картинок через RouterAI
|
||||
(Nano Banana / Gemini Image, FLUX, GPT-Image).
|
||||
|
||||
Транспорт:
|
||||
MCP_TRANSPORT=stdio (по умолчанию) | http
|
||||
MCP_PORT=8790 (для http)
|
||||
|
||||
Ключ ищем в порядке:
|
||||
1) env ROUTERAI_API_KEY / IMAGE_MCP_API_KEY
|
||||
2) ~/.hermes/models.yaml -> models.deepseek.api_key (base_url = routerai.ru)
|
||||
3) ~/.hermes/models.yaml -> любой api_key с base_url routerai.ru
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
import mimetypes
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
|
||||
BASE_URL = os.environ.get("ROUTERAI_BASE_URL", "https://routerai.ru/api/v1").rstrip("/")
|
||||
DEFAULT_MODEL = os.environ.get("IMAGE_MCP_DEFAULT_MODEL", "google/gemini-3.1-flash-image")
|
||||
DEFAULT_OUT_DIR = os.environ.get(
|
||||
"IMAGE_MCP_OUT_DIR", str(Path.home() / ".hermes" / "image_cache" / "mcp-gen")
|
||||
)
|
||||
REQUEST_TIMEOUT = int(os.environ.get("IMAGE_MCP_TIMEOUT", "300"))
|
||||
|
||||
# модель -> (бэкенд, человеческое имя). backend "chat" = OpenAI chat с modalities,
|
||||
# backend "images" = POST /images/generations
|
||||
MODELS: dict[str, tuple[str, str]] = {
|
||||
"google/gemini-3.1-flash-image": ("chat", "Nano Banana 2 (Gemini 3.1 Flash Image)"),
|
||||
"google/gemini-3.1-flash-lite-image": ("chat", "Nano Banana 2 Lite"),
|
||||
"google/gemini-2.5-flash-image": ("chat", "Nano Banana (Gemini 2.5 Flash Image)"),
|
||||
"google/gemini-3-pro-image": ("chat", "Nano Banana Pro (Gemini 3 Pro Image)"),
|
||||
"google/gemini-3.1-flash-image-preview": ("chat", "Nano Banana 2 (preview)"),
|
||||
"google/gemini-3-pro-image-preview": ("chat", "Nano Banana Pro (preview)"),
|
||||
"openai/gpt-image-1": ("images", "GPT Image 1"),
|
||||
"openai/gpt-image-1-mini": ("images", "GPT Image 1 Mini"),
|
||||
"black-forest-labs/flux.2-pro": ("images", "FLUX.2 Pro"),
|
||||
"black-forest-labs/flux.2-max": ("images", "FLUX.2 Max"),
|
||||
"black-forest-labs/flux.2-flex": ("images", "FLUX.2 Flex"),
|
||||
"black-forest-labs/flux.2-klein-4b": ("images", "FLUX.2 Klein 4B"),
|
||||
}
|
||||
|
||||
mcp = FastMCP("routerai-image")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- key
|
||||
|
||||
|
||||
def _api_key() -> str:
|
||||
for env in ("ROUTERAI_API_KEY", "IMAGE_MCP_API_KEY"):
|
||||
v = os.environ.get(env)
|
||||
if v and len(v) > 20:
|
||||
return v.strip()
|
||||
|
||||
models_yaml = Path.home() / ".hermes" / "models.yaml"
|
||||
if models_yaml.exists():
|
||||
try:
|
||||
import yaml # type: ignore
|
||||
|
||||
data = yaml.safe_load(models_yaml.read_text(encoding="utf-8")) or {}
|
||||
cands = []
|
||||
for name, entry in (data.get("models") or {}).items():
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
key = str(entry.get("api_key") or "")
|
||||
base = str(entry.get("base_url") or "")
|
||||
if "routerai.ru" in base and key.startswith("sk-"):
|
||||
cands.append((name != "deepseek", key))
|
||||
if cands:
|
||||
cands.sort()
|
||||
return cands[0][1]
|
||||
except Exception as exc: # noqa: BLE001
|
||||
print(f"[routerai-image] models.yaml parse failed: {exc}", file=sys.stderr)
|
||||
|
||||
raise RuntimeError(
|
||||
"RouterAI API key not found. Set ROUTERAI_API_KEY, or put a routerai.ru "
|
||||
"entry with api_key into ~/.hermes/models.yaml"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------- http
|
||||
|
||||
|
||||
def _post(path: str, body: dict, timeout: int = REQUEST_TIMEOUT) -> dict:
|
||||
key = _api_key()
|
||||
data = json.dumps(body, ensure_ascii=False).encode("utf-8")
|
||||
last_err = ""
|
||||
|
||||
for attempt in range(4):
|
||||
req = urllib.request.Request(
|
||||
f"{BASE_URL}{path}",
|
||||
data=data,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {key}",
|
||||
"User-Agent": "routerai-image-mcp/1.0",
|
||||
},
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||||
raw = resp.read()
|
||||
# RouterAI иногда отдаёт мусорные префиксы перед JSON
|
||||
start = raw.find(b"{")
|
||||
return json.loads(raw[start:] if start > 0 else raw)
|
||||
except urllib.error.HTTPError as exc:
|
||||
payload = exc.read().decode("utf-8", "replace")[:400]
|
||||
last_err = f"HTTP {exc.code}: {payload}"
|
||||
if exc.code in (429, 500, 502, 503, 504) and attempt < 3:
|
||||
time.sleep(3 * (attempt + 1))
|
||||
continue
|
||||
raise RuntimeError(last_err) from None
|
||||
except Exception as exc: # noqa: BLE001
|
||||
last_err = f"{type(exc).__name__}: {exc}"
|
||||
if attempt < 3:
|
||||
time.sleep(3 * (attempt + 1))
|
||||
continue
|
||||
raise RuntimeError(last_err) from None
|
||||
raise RuntimeError(last_err)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------- helpers
|
||||
|
||||
|
||||
def _data_url(path: str) -> str:
|
||||
p = Path(path).expanduser()
|
||||
if not p.exists():
|
||||
raise RuntimeError(f"reference image not found: {p}")
|
||||
mime = mimetypes.guess_type(p.name)[0] or "image/png"
|
||||
return f"data:{mime};base64,{base64.b64encode(p.read_bytes()).decode()}"
|
||||
|
||||
|
||||
def _decode_image(item: dict) -> bytes:
|
||||
"""Принимает {'b64_json': ...} или {'url': 'data:...'} / http url."""
|
||||
if item.get("b64_json"):
|
||||
return base64.b64decode(item["b64_json"])
|
||||
url = item.get("url") or (item.get("image_url") or {}).get("url") or ""
|
||||
if url.startswith("data:"):
|
||||
return base64.b64decode(url.split(",", 1)[1])
|
||||
if url.startswith("http"):
|
||||
with urllib.request.urlopen(url, timeout=REQUEST_TIMEOUT) as resp: # noqa: S310
|
||||
return resp.read()
|
||||
raise RuntimeError("no image payload in response")
|
||||
|
||||
|
||||
def _ext_from_bytes(blob: bytes) -> str:
|
||||
if blob[:8] == b"\x89PNG\r\n\x1a\n":
|
||||
return ".png"
|
||||
if blob[:3] == b"\xff\xd8\xff":
|
||||
return ".jpg"
|
||||
if blob[:4] == b"RIFF" and blob[8:12] == b"WEBP":
|
||||
return ".webp"
|
||||
return ".png"
|
||||
|
||||
|
||||
def _slug(text: str, limit: int = 40) -> str:
|
||||
txt = re.sub(r"[^\w\s-]", "", text, flags=re.UNICODE).strip()
|
||||
txt = re.sub(r"[\s-]+", "-", txt)
|
||||
return (txt[:limit] or "image").strip("-").lower()
|
||||
|
||||
|
||||
def _run(model: str, prompt: str, refs: list[str] | None, out_dir: str, filename: str | None) -> str:
|
||||
if model not in MODELS:
|
||||
known = ", ".join(sorted(MODELS))
|
||||
raise RuntimeError(f"unknown model '{model}'. Known: {known}")
|
||||
backend, human = MODELS[model]
|
||||
|
||||
refs = refs or []
|
||||
started = time.time()
|
||||
|
||||
if backend == "chat":
|
||||
if refs:
|
||||
content: list[dict] = [{"type": "text", "text": prompt}]
|
||||
content += [
|
||||
{"type": "image_url", "image_url": {"url": _data_url(p)}} for p in refs
|
||||
]
|
||||
else:
|
||||
content = prompt # type: ignore[assignment]
|
||||
body = {
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
"modalities": ["image", "text"],
|
||||
}
|
||||
if not refs:
|
||||
body["response_modalities"] = ["IMAGE"]
|
||||
data = _post("/chat/completions", body)
|
||||
msg = (data.get("choices") or [{}])[0].get("message", {}) or {}
|
||||
images = msg.get("images") or []
|
||||
if not images and isinstance(msg.get("content"), list):
|
||||
images = [c for c in msg["content"] if c.get("type") in ("image_url", "image")]
|
||||
if not images:
|
||||
raise RuntimeError(f"model returned no image. raw: {json.dumps(data)[:400]}")
|
||||
blob = _decode_image(images[0])
|
||||
usage = data.get("usage") or {}
|
||||
else:
|
||||
body = {"model": model, "prompt": prompt, "n": 1}
|
||||
if refs:
|
||||
body["image"] = _data_url(refs[0])
|
||||
data = _post("/images/generations", body)
|
||||
items = data.get("data") or []
|
||||
if not items:
|
||||
raise RuntimeError(f"model returned no image. raw: {json.dumps(data)[:400]}")
|
||||
blob = _decode_image(items[0])
|
||||
usage = data.get("usage") or {}
|
||||
|
||||
out = Path(out_dir).expanduser()
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
name = filename or f"{time.strftime('%Y%m%d-%H%M%S')}-{_slug(prompt)}"
|
||||
if not Path(name).suffix:
|
||||
name += _ext_from_bytes(blob)
|
||||
dest = out / name
|
||||
dest.write_bytes(blob)
|
||||
|
||||
cost = usage.get("cost")
|
||||
secs = time.time() - started
|
||||
cost_line = f"{cost:.2f}₽" if isinstance(cost, (int, float)) else "n/a"
|
||||
kb = len(blob) / 1024
|
||||
return (
|
||||
f"Готово: {human}\n"
|
||||
f"Файл: {dest} ({kb:.0f} KB)\n"
|
||||
f"Модель: {model} | время: {secs:.1f}s | стоимость: {cost_line}\n"
|
||||
f"MEDIA:{dest}\n"
|
||||
f"(покажи файл пользователю: вставь строку MEDIA:<путь> целиком в ответ)"
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- tools
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def generate_image(
|
||||
prompt: str,
|
||||
model: str = DEFAULT_MODEL,
|
||||
reference_images: list[str] | None = None,
|
||||
out_dir: str = DEFAULT_OUT_DIR,
|
||||
filename: str | None = None,
|
||||
) -> str:
|
||||
"""Сгенерировать картинку по текстовому ТЗ (или отредактировать существующую).
|
||||
|
||||
Args:
|
||||
prompt: ТЗ на картинке — что нарисовать. Чем конкретнее (сюжет, стиль,
|
||||
свет, композиция, пропорции), тем лучше результат.
|
||||
model: id модели RouterAI. Дефолт google/gemini-3.1-flash-image
|
||||
(Nano Banana 2). Дешёвые: google/gemini-3.1-flash-lite-image,
|
||||
google/gemini-2.5-flash-image. Дорогие: google/gemini-3-pro-image.
|
||||
Альтернативы: black-forest-labs/flux.2-pro, openai/gpt-image-1.
|
||||
reference_images: список путей к картинкам-референсам/для правки.
|
||||
out_dir: куда сложить результат (по умолчанию ~/.hermes/image_cache/mcp-gen).
|
||||
filename: имя файла без пути (расширение подставится само).
|
||||
|
||||
Returns:
|
||||
Текст с абсолютным путём к файлу и строкой `MEDIA:<путь>` — вставь эту
|
||||
строку в свой ответ, чтобы клиент показал картинку.
|
||||
"""
|
||||
return _run(model, prompt, reference_images, out_dir, filename)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def list_image_models() -> str:
|
||||
"""Список доступных моделей генерации картинок и их бэкендов."""
|
||||
lines = [f"{mid} — {human} (backend: {backend})" for mid, (backend, human) in sorted(MODELS.items())]
|
||||
return "Доступные модели:\n" + "\n".join(lines) + f"\n\nДефолт: {DEFAULT_MODEL}\nБазовый URL: {BASE_URL}"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
transport = os.environ.get("MCP_TRANSPORT", "stdio").lower()
|
||||
if transport in ("http", "streamable-http", "sse"):
|
||||
mcp.settings.host = os.environ.get("MCP_HOST", "0.0.0.0")
|
||||
mcp.settings.port = int(os.environ.get("MCP_PORT", "8790"))
|
||||
mcp.run(transport="streamable-http")
|
||||
else:
|
||||
mcp.run(transport="stdio")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user