librechat-custom/mcp/image_server.py

#!/usr/bin/env python3
"""Inference Cooperative image-generation MCP server (stdio, stdlib-only).

Exposes a `generate_image` tool that calls the co-op gateway
(`regolo/qwen-image` via POST /v1/images/generations). Each member's own
co-op API key is injected per-user by LibreChat as the IC_API_KEY env var
(`customUserVars` + `env` in librechat.yaml), so usage bills to the member's
own LiteLLM team budget — the same model as the chat endpoint.

Design notes (mirrors mcp/server.py, the docs server):
- stdio transport: newline-delimited JSON-RPC 2.0; logging to stderr only.
- No third-party dependencies. The image is returned inline as MCP image
  content (base64 PNG), which LibreChat renders directly in the chat.
"""

import base64
import json
import os
import sys
import urllib.error
import urllib.request

GATEWAY_IMAGES_URL = "https://gateway.inference.coop/v1/images/generations"
MODEL = "regolo/qwen-image"
REQUEST_TIMEOUT = 90
ALLOWED_SIZES = ("1024x1024", "1536x1024", "1024x1536")
DEFAULT_SIZE = "1024x1024"
DASHBOARD_URL = "https://dashboard.inference.coop"

SERVER_NAME = "inference-cooperative-image"
SERVER_VERSION = "1.0.0"


def log(msg):
    sys.stderr.write(f"[inference-image-mcp] {msg}\n")
    sys.stderr.flush()


def tool_generate_image(args):
    key = (os.environ.get("IC_API_KEY") or "").strip()
    if not key:
        return {
            "content": [{
                "type": "text",
                "text": (
                    "No Inference Co-op API key configured yet. Open the MCP panel "
                    f"(inference-image) and save your key — create one at {DASHBOARD_URL}. "
                    "Your key is only visible to you, and image generation bills to your "
                    "own co-op balance."
                ),
            }],
            "isError": True,
        }

    prompt = str(args.get("prompt") or "").strip()
    if not prompt:
        return {"content": [{"type": "text", "text": "Provide a prompt for the image."}]}

    size = str(args.get("size") or DEFAULT_SIZE).strip()
    if size not in ALLOWED_SIZES:
        return {
            "content": [{
                "type": "text",
                "text": f"Unsupported size '{size}'. Use one of: {', '.join(ALLOWED_SIZES)}.",
            }],
            "isError": True,
        }

    body = json.dumps({"model": MODEL, "prompt": prompt, "n": 1, "size": size}).encode()
    req = urllib.request.Request(
        GATEWAY_IMAGES_URL,
        data=body,
        headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
    )
    try:
        with urllib.request.urlopen(req, timeout=REQUEST_TIMEOUT) as resp:
            data = json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", "replace")[:400]
        log(f"gateway HTTP {exc.code}: {detail}")
        hint = "Check that your co-op API key is valid and has balance."
        if exc.code == 401:
            hint = f"Your key was rejected — re-save it in the MCP panel (create one at {DASHBOARD_URL})."
        elif exc.code == 429:
            hint = "Budget or rate limit reached on your key."
        return {
            "content": [{"type": "text", "text": f"Image generation failed (HTTP {exc.code}). {hint}"}],
            "isError": True,
        }
    except Exception as exc:  # noqa: BLE001 - report transport errors in-band
        log(f"gateway error: {exc}")
        return {
            "content": [{"type": "text", "text": f"Image generation failed: {exc}"}],
            "isError": True,
        }

    items = (data or {}).get("data") or []
    if not items:
        return {"content": [{"type": "text", "text": "The gateway returned no image."}], "isError": True}

    item = items[0]
    content = []
    if item.get("b64_json"):
        content.append({"type": "image", "data": item["b64_json"], "mimeType": "image/png"})
    notes = [f"Generated with {MODEL} ({size}) — billed to your co-op key."]
    if item.get("revised_prompt"):
        notes.append(f"Prompt: {item['revised_prompt']}")
    if not item.get("b64_json") and item.get("url"):
        notes.append(f"Image URL: {item['url']}")
    content.append({"type": "text", "text": " ".join(notes)})
    return {"content": content}


TOOLS = [
    {
        "name": "generate_image",
        "description": (
            "Generate an image from a text prompt (regolo/qwen-image via the Inference "
            "Co-op gateway, EU renewable-powered). Costs roughly 1-2 cents per image, "
            "billed to your own co-op API key."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "prompt": {"type": "string", "description": "What to draw"},
                "size": {
                    "type": "string",
                    "enum": list(ALLOWED_SIZES),
                    "description": "Image size (default 1024x1024)",
                },
            },
            "required": ["prompt"],
        },
    }
]

TOOLS_HANDLERS = {"generate_image": tool_generate_image}


def handle(msg):
    method = msg.get("method")
    msg_id = msg.get("id")

    if method == "initialize":
        return {
            "jsonrpc": "2.0",
            "id": msg_id,
            "result": {
                "protocolVersion": "2024-11-05",
                "capabilities": {"tools": {}},
                "serverInfo": {"name": SERVER_NAME, "version": SERVER_VERSION},
            },
        }
    if method == "ping":
        return {"jsonrpc": "2.0", "id": msg_id, "result": {}}
    if method == "tools/list":
        return {"jsonrpc": "2.0", "id": msg_id, "result": {"tools": TOOLS}}
    if method == "tools/call":
        params = msg.get("params") or {}
        name = params.get("name")
        handler = TOOLS_HANDLERS.get(name)
        if not handler:
            return {
                "jsonrpc": "2.0",
                "id": msg_id,
                "error": {"code": -32602, "message": f"Unknown tool: {name}"},
            }
        try:
            result = handler(params.get("arguments") or {})
            return {"jsonrpc": "2.0", "id": msg_id, "result": result}
        except Exception as exc:  # noqa: BLE001 - report tool errors in-band
            log(f"tool error: {exc}")
            return {
                "jsonrpc": "2.0",
                "id": msg_id,
                "result": {"content": [{"type": "text", "text": f"Tool failed: {exc}"}], "isError": True},
            }
    if method and method.startswith("notifications/"):
        return None
    if msg_id is not None:
        return {"jsonrpc": "2.0", "id": msg_id, "error": {"code": -32601, "message": "Method not found"}}
    return None


def main():
    log("starting (stdio)")
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            msg = json.loads(line)
        except json.JSONDecodeError as exc:
            log(f"bad json: {exc}")
            continue
        try:
            resp = handle(msg)
        except Exception as exc:  # noqa: BLE001 - never crash the protocol loop
            log(f"handler crash: {exc}")
            if msg.get("id") is not None:
                resp = {"jsonrpc": "2.0", "id": msg["id"], "error": {"code": -32603, "message": "Internal error"}}
            else:
                resp = None
        if resp is not None:
            sys.stdout.write(json.dumps(resp) + "\n")
            sys.stdout.flush()


if __name__ == "__main__":
    main()