subochev
585c6a4ca4
- 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, питфоллы
290 lines
11 KiB
Python
290 lines
11 KiB
Python
#!/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()
|