#!/usr/bin/env python3
"""
智谱AI图像生成MCP服务器
这个MCP服务器提供了调用智谱AI图像生成API的功能。
支持CogView-4系列模型进行图像生成。
使用方式:
python jimeng_image_server.py
环境变量:
ZHIPU_API_KEY: 智谱AI的API密钥
"""
import asyncio
import os
import json
from typing import Optional, Literal
import httpx
from pydantic import BaseModel, Field
from mcp.server.fastmcp import FastMCP
# 创建MCP服务器实例
mcp = FastMCP("智谱AI图像生成服务器")
# API配置
API_BASE_URL = "https://open.bigmodel.cn/api/paas/v4/images/generations"
class ImageGenerationRequest(BaseModel):
"""图像生成请求模型"""
model: Literal["cogview-4-250304", "cogview-4", "cogview-3-flash"] = Field(
default="cogview-4-250304",
description="图像生成模型"
)
prompt: str = Field(
description="图像描述文本",
min_length=1,
max_length=1000
)
quality: Optional[Literal["hd", "standard"]] = Field(
default="standard",
description="图像质量。hd: 高质量(20秒), standard: 标准质量(5-10秒)"
)
size: Optional[str] = Field(
default="1024x1024",
description="图像尺寸，如: 1024x1024, 768x1344, 864x1152等"
)
user_id: Optional[str] = Field(
default=None,
description="用户ID，用于违规行为监控",
min_length=6,
max_length=128
)
class ImageGenerationResponse(BaseModel):
"""图像生成响应模型"""
created: int = Field(description="请求创建时间戳")
data: list[dict] = Field(description="生成的图像数据")
content_filter: Optional[list[dict]] = Field(default=None, description="内容过滤信息")
@mcp.tool()
async def generate_image(
prompt: str,
model: Literal["cogview-4-250304", "cogview-4", "cogview-3-flash"] = "cogview-4-250304",
quality: Literal["hd", "standard"] = "standard",
size: str = "1024x1024",
user_id: Optional[str] = None,
api_key: Optional[str] = None
) -> ImageGenerationResponse:
"""
生成图像
Args:
prompt: 图像描述文本
model: 图像生成模型 (cogview-4-250304, cogview-4, cogview-3-flash)
quality: 图像质量 (hd: 高质量约20秒, standard: 标准质量约5-10秒)
size: 图像尺寸 (如: 1024x1024, 768x1344, 864x1152)
user_id: 用户ID，用于违规行为监控 (可选，6-128字符)
api_key: API密钥 (可选，如果不提供则使用环境变量ZHIPU_API_KEY)
Returns:
ImageGenerationResponse: 包含生成图像URL的响应
Raises:
ValueError: 参数验证失败
RuntimeError: API调用失败
"""
# 获取API密钥
token = api_key or os.getenv("ZHIPU_API_KEY")
if not token:
raise ValueError("API密钥未提供。请设置ZHIPU_API_KEY环境变量或在函数调用中提供api_key参数")
# 验证参数
if not prompt or len(prompt.strip()) == 0:
raise ValueError("图像描述文本不能为空")
if len(prompt) > 1000:
raise ValueError("图像描述文本不能超过1000个字符")
# 验证图像尺寸格式
if size:
try:
width, height = map(int, size.split('x'))
if width < 512 or width > 2048 or height < 512 or height > 2048:
raise ValueError("图像尺寸宽高必须在512px-2048px之间")
if width % 16 != 0 or height % 16 != 0:
raise ValueError("图像尺寸宽高必须是16的倍数")
if width * height > 2**21:
raise ValueError("图像最大像素数不能超过2^21px")
except ValueError as e:
if "invalid literal" in str(e):
raise ValueError("图像尺寸格式错误，应为 宽x高 格式，如: 1024x1024")
raise
# 验证用户ID
if user_id and (len(user_id) < 6 or len(user_id) > 128):
raise ValueError("用户ID长度必须在6-128字符之间")
# 构建请求数据
request_data = {
"model": model,
"prompt": prompt.strip(),
"quality": quality,
"size": size
}
if user_id:
request_data["user_id"] = user_id
# 构建请求头
headers = {
"Authorization": f"Bearer {token}",
"Content-Type": "application/json"
}
try:
# 使用httpx发送异步请求
async with httpx.AsyncClient(timeout=60.0) as client:
response = await client.post(
API_BASE_URL,
headers=headers,
json=request_data
)
response_text = response.text
if response.status_code == 200:
response_data = response.json()
return ImageGenerationResponse(**response_data)
else:
# 处理错误响应
try:
error_data = response.json()
error_message = error_data.get("error", {}).get("message", response_text)
except:
error_message = response_text
raise RuntimeError(f"API请求失败 (状态码: {response.status_code}): {error_message}")
except httpx.RequestError as e:
raise RuntimeError(f"网络请求失败: {str(e)}")
except httpx.TimeoutException:
raise RuntimeError("请求超时，请稍后重试")
except ValueError as e:
# 重新抛出参数验证错误
raise
except Exception as e:
raise RuntimeError(f"请求处理失败: {str(e)}")
@mcp.tool()
async def get_supported_models() -> dict[str, str]:
"""
获取支持的图像生成模型列表
Returns:
dict: 模型名称和描述的字典
"""
return {
"cogview-4-250304": "CogView-4最新版本，支持高质量和标准质量两种模式",
"cogview-4": "CogView-4标准版本",
"cogview-3-flash": "CogView-3快速版本，生成速度更快"
}
@mcp.tool()
async def get_supported_sizes() -> dict[str, str]:
"""
获取推荐的图像尺寸列表
Returns:
dict: 尺寸规格和描述的字典
"""
return {
"1024x1024": "正方形 - 默认尺寸，适合大部分场景",
"768x1344": "竖屏 - 适合手机壁纸、海报等",
"864x1152": "竖屏 - 适合头像、人物肖像等",
"1344x768": "横屏 - 适合桌面壁纸、横幅等",
"1152x864": "横屏 - 适合风景、全景等",
"1440x720": "宽屏 - 适合超宽屏壁纸",
"720x1440": "长竖屏 - 适合手机长图"
}
@mcp.tool()
async def test_api_connection(api_key: Optional[str] = None) -> dict:
"""
测试API连接状态
Args:
api_key: API密钥 (可选，如果不提供则使用环境变量)
Returns:
dict: 连接测试结果
"""
token = api_key or os.getenv("ZHIPU_API_KEY")
if not token:
return {
"status": "error",
"message": "API密钥未提供"
}
# 发送一个简单的测试请求
headers = {
"Authorization": f"Bearer {token}",
"Content-Type": "application/json"
}
test_data = {
"model": "cogview-3-flash",
"prompt": "test",
"size": "512x512"
}
try:
async with httpx.AsyncClient(timeout=10.0) as client:
response = await client.post(
API_BASE_URL,
headers=headers,
json=test_data
)
if response.status_code == 200:
return {
"status": "success",
"message": "API连接正常",
"api_url": API_BASE_URL
}
elif response.status_code == 401:
return {
"status": "error",
"message": "API密钥无效或已过期"
}
else:
return {
"status": "error",
"message": f"API返回错误: {response.status_code}",
"details": response.text
}
except httpx.TimeoutException:
return {
"status": "error",
"message": "连接超时，请检查网络连接"
}
except httpx.RequestError as e:
return {
"status": "error",
"message": f"网络连接失败: {str(e)}"
}
except Exception as e:
return {
"status": "error",
"message": f"连接测试失败: {str(e)}"
}
@mcp.resource("config://api-info")
def get_api_info() -> str:
"""获取API配置信息"""
api_key_set = "已设置" if os.getenv("ZHIPU_API_KEY") else "未设置"
return f"""
智谱AI图像生成API配置信息:
API地址: {API_BASE_URL}
API密钥: {api_key_set}
环境变量配置:
- ZHIPU_API_KEY: 智谱AI的API密钥
支持的功能:
- 图像生成 (generate_image)
- 获取支持的模型 (get_supported_models)
- 获取支持的尺寸 (get_supported_sizes)
- 测试API连接 (test_api_connection)
技术栈:
- HTTP客户端: httpx (异步)
- 协议: Model Context Protocol (MCP)
- 传输方式: stdio
"""
def main():
"""主函数"""
# 检查环境变量
if not os.getenv("ZHIPU_API_KEY"):
print("警告: 未设置ZHIPU_API_KEY环境变量")
print("请设置环境变量或在工具调用时提供api_key参数")
# 运行MCP服务器
mcp.run()
if __name__ == "__main__":
main()
