207 lines
7.3 KiB
Python
207 lines
7.3 KiB
Python
#!/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 the user has not set a personal API key, fall back to an unauthenticated request.
|
|
# The gateway will bill the request to the member's account based on the session
|
|
# (the LibreChat backend already knows the member identity). This mirrors the
|
|
# behaviour of the chat endpoint, which does not require the user to manually
|
|
# provide a key.
|
|
if not key:
|
|
auth_header = {}
|
|
else:
|
|
auth_header = {"Authorization": f"Bearer {key}"}
|
|
|
|
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={**auth_header, "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()
|