102 lines
3.1 KiB
Python
102 lines
3.1 KiB
Python
"""MCP Server：向 Agent 暴露工具。
|
|
|
|
强制流程：list_providers() → discover_models(provider_id) → generate_image(...)
|
|
matched 为空时，Agent 必须直接回复「没找到有效的模型」。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from fastmcp import FastMCP
|
|
|
|
from .providers import ProviderError
|
|
from .services import generation as generation_service
|
|
from .services import providers as provider_service
|
|
from .services import skill_doc
|
|
|
|
mcp = FastMCP("ai-image")
|
|
|
|
|
|
@mcp.tool
|
|
def list_providers() -> list[dict[str, Any]]:
|
|
"""第 1 步：列出可用供应商及其已启用模型数。返回后必须调用 discover_models 拉取模型。"""
|
|
out = []
|
|
for p in provider_service.list_providers():
|
|
if not p["enabled"]:
|
|
continue
|
|
out.append(
|
|
{
|
|
"provider": p["id"],
|
|
"name": p["name"],
|
|
"adapter": p["adapter"],
|
|
"base_url": p["base_url"],
|
|
"default_model": p["default_model"],
|
|
"model_count": p["model_count"],
|
|
"enabled_model_count": p["enabled_model_count"],
|
|
"note": p["note"],
|
|
}
|
|
)
|
|
return out
|
|
|
|
|
|
@mcp.tool
|
|
def discover_models(provider_id: str) -> dict[str, Any]:
|
|
"""第 2 步（必做）：拉取指定供应商的上游模型清单。
|
|
|
|
返回 {provider_id, total, matched}。matched 为模型 ID 转小写后包含 "image" 的集合，
|
|
是唯一可用的模型来源；matched 为空时必须停止并回复「没找到有效的模型」。
|
|
"""
|
|
try:
|
|
return provider_service.discover(provider_id, require_enabled=True)
|
|
except ProviderError as exc:
|
|
return {"provider_id": provider_id, "total": 0, "matched": [], "error": exc.message}
|
|
|
|
|
|
@mcp.tool
|
|
def generate_image(
|
|
provider: str,
|
|
model: str,
|
|
prompt: str,
|
|
size: str = "1024x1024",
|
|
ratio: str = "",
|
|
n: int = 1,
|
|
images: list[str] | None = None,
|
|
profile_name: str = "agent",
|
|
) -> dict[str, Any]:
|
|
"""第 3 步：出图。model 必须是 discover_models 返回的 matched 之一。"""
|
|
try:
|
|
result = generation_service.generate(
|
|
provider_id=provider,
|
|
model=model,
|
|
prompt=prompt,
|
|
size=size,
|
|
ratio=ratio,
|
|
n=n,
|
|
images=images or [],
|
|
profile_name=profile_name,
|
|
)
|
|
except ProviderError as exc:
|
|
return {"status": "failed", "error": exc.message}
|
|
return {"status": "success", **result}
|
|
|
|
|
|
@mcp.tool
|
|
def list_generations(limit: int = 20) -> list[dict[str, Any]]:
|
|
"""查询最近的生成记录。"""
|
|
return generation_service.list_generations()[:limit]
|
|
|
|
|
|
@mcp.tool
|
|
def get_skill_doc() -> str:
|
|
"""获取最新版 SKILL.md 原文（不含固定模型名）。"""
|
|
return skill_doc.get_skill()["body"]
|
|
|
|
|
|
def build_http_app(path: str = "/"):
|
|
"""把 MCP 挂到已有 FastAPI 应用上（同一端口）。
|
|
|
|
返回的 Starlette 应用自带 path 路径，由入口分发器转发请求。
|
|
"""
|
|
return mcp.http_app(path=path)
|