"""
Seedream 4.0 MCP工具 - 组图生成工具
实现连续生成多张图像功能，支持自动保存到本地。
"""
from typing import Any, Dict, List, Optional
from pathlib import Path
from mcp.types import Tool, TextContent
from ..client import SeedreamClient
from ..config import SeedreamConfig, get_global_config
from ..utils.logging import get_logger
from ..utils.auto_save import AutoSaveManager, AutoSaveResult
# 工具定义
sequential_generation_tool = Tool(
name="seedream_sequential_generation",
description="使用Seedream 4.0连续生成多张图像（组图生成），支持3种输入类型：文生组图、单图生组图、多图生组图",
inputSchema={
"type": "object",
"properties": {
"prompt": {
"type": "string",
"description": "图像生成的文本提示词，应明确指明生成数量和每一幅图的具体内容描述，建议不超过600个字符",
"maxLength": 600
},
"max_images": {
"type": "integer",
"description": "最大生成图像数量",
"minimum": 1,
"maximum": 15,
"default": 4
},
"size": {
"type": "string",
"description": "生成图像的尺寸，如果不指定则使用配置文件中的默认值",
"enum": ["1K", "2K", "4K"]
},
"watermark": {
"type": "boolean",
"description": "是否在生成的图像上添加水印，如果不指定则使用配置文件中的默认值"
},
"response_format": {
"type": "string",
"description": "响应格式：url返回图像URL，b64_json返回base64编码",
"enum": ["url", "b64_json"],
"default": "url"
},
"image": {
"type": ["string", "array"],
"description": "可选的参考图像。支持单张图片URL/路径（字符串）或多张图片URL/路径（字符串数组）。用于单图生组图或多图生组图",
"items": {
"type": "string"
},
"maxItems": 10
},
"auto_save": {
"type": "boolean",
"description": "是否自动保存生成的图片到本地。如果未指定，将使用全局配置",
"default": None
},
"save_path": {
"type": "string",
"description": "自定义保存目录路径。如果未指定，将使用默认配置路径"
},
"custom_name": {
"type": "string",
"description": "自定义文件名前缀。如果未指定，将根据提示词自动生成"
}
},
"required": ["prompt"]
}
)
async def handle_sequential_generation(arguments: Dict[str, Any]) -> List[TextContent]:
"""处理组图生成请求
Args:
arguments: 工具参数
Returns:
MCP响应内容
"""
logger = get_logger(__name__)
try:
# 获取配置
config = get_global_config()
# 提取参数，按优先级：调用参数 > 配置文件默认值 > 方法默认值
prompt = arguments.get("prompt")
max_images = arguments.get("max_images", 4)
size = arguments.get("size") or config.default_size
watermark = arguments.get("watermark")
if watermark is None:
watermark = config.default_watermark
response_format = arguments.get("response_format", "url")
image = arguments.get("image")
# 提取自动保存参数
auto_save = arguments.get("auto_save")
save_path = arguments.get("save_path")
custom_name = arguments.get("custom_name")
# 验证参数
if not prompt:
return [TextContent(type="text", text="错误：prompt参数是必需的")]
if max_images < 1 or max_images > 15:
return [TextContent(type="text", text="错误：max_images必须在1-15之间")]
if size not in ["1K", "2K", "4K"]:
return [TextContent(type="text", text="错误：size必须是1K、2K或4K")]
if response_format not in ["url", "b64_json"]:
return [TextContent(type="text", text="错误：response_format必须是url或b64_json")]
# 验证image参数
if image is not None:
if isinstance(image, str):
# 单张图片
if not image.strip():
return [TextContent(type="text", text="错误：image参数不能为空字符串")]
elif isinstance(image, list):
# 多张图片
if len(image) == 0:
return [TextContent(type="text", text="错误：image数组不能为空")]
if len(image) > 10:
return [TextContent(type="text", text="错误：最多支持10张参考图片")]
for img in image:
if not isinstance(img, str) or not img.strip():
return [TextContent(type="text", text="错误：image数组中的每个元素都必须是非空字符串")]
else:
return [TextContent(type="text", text="错误：image参数必须是字符串或字符串数组")]
logger.info(f"开始处理组图生成请求: prompt='{prompt[:50]}...', max_images={max_images}, size={size}")
# 确定是否启用自动保存
enable_auto_save = auto_save if auto_save is not None else config.auto_save_enabled
# 创建客户端并调用API
async with SeedreamClient(config) as client:
result = await client.sequential_generation(
prompt=prompt,
max_images=max_images,
size=size,
watermark=watermark,
response_format=response_format,
image=image
)
# 初始化自动保存结果
auto_save_results = []
# 如果启用自动保存且API调用成功，执行自动保存
if enable_auto_save and result.get("success"):
try:
if response_format == "url":
auto_save_results = await _handle_auto_save(
result, prompt, config, save_path, custom_name
)
if auto_save_results:
result = _update_result_with_auto_save(result, auto_save_results)
elif response_format == "b64_json":
auto_save_results = await _handle_auto_save_base64(
result, prompt, config, save_path, custom_name
)
if auto_save_results:
result = _update_result_with_auto_save(result, auto_save_results)
except Exception as e:
logger.warning(f"自动保存失败，但继续返回原始结果: {e}")
# 格式化响应
response_text = _format_sequential_generation_response(
result, prompt, max_images, size, auto_save_results, enable_auto_save
)
logger.info("组图生成请求处理完成")
return [TextContent(type="text", text=response_text)]
except Exception as e:
logger.error(f"组图生成请求处理失败: {str(e)}")
error_msg = f"组图生成失败: {str(e)}"
return [TextContent(type="text", text=error_msg)]
async def _handle_auto_save(
result: Dict[str, Any],
prompt: str,
config: SeedreamConfig,
save_path: Optional[str] = None,
custom_name: Optional[str] = None
) -> List[AutoSaveResult]:
"""处理自动保存逻辑
Args:
result: API响应结果
prompt: 生成提示词
config: 配置对象
save_path: 自定义保存路径
custom_name: 自定义文件名前缀
Returns:
自动保存结果列表
"""
# 初始化自动保存管理器
base_dir = Path(save_path) if save_path else (
Path(config.auto_save_base_dir) if config.auto_save_base_dir else None
)
auto_save_manager = AutoSaveManager(
base_dir=base_dir,
download_timeout=config.auto_save_download_timeout,
max_retries=config.auto_save_max_retries,
max_file_size=config.auto_save_max_file_size,
max_concurrent=config.auto_save_max_concurrent
)
# 提取图片URL
image_urls = []
if result.get("data"):
for item in result["data"]:
if item.get("url"):
image_urls.append(item["url"])
if not image_urls:
return []
# 准备图片数据
image_data = []
for i, url in enumerate(image_urls):
data = {
"url": url,
"prompt": prompt,
"custom_name": custom_name
}
image_data.append(data)
# 执行批量保存
return await auto_save_manager.save_multiple_images(
image_data, "sequential_generation"
)
async def _handle_auto_save_base64(
result: Dict[str, Any],
prompt: str,
config: SeedreamConfig,
save_path: Optional[str] = None,
custom_name: Optional[str] = None
) -> List[AutoSaveResult]:
"""处理 base64 自动保存（组图生成）
当 response_format 为 b64_json 时，从结果中提取 base64 并保存到本地。
"""
logger = get_logger(__name__)
try:
base_dir = Path(save_path) if save_path else (
Path(config.auto_save_base_dir) if config.auto_save_base_dir else None
)
auto_save_manager = AutoSaveManager(
base_dir=base_dir,
download_timeout=config.auto_save_download_timeout,
max_retries=config.auto_save_max_retries,
max_file_size=config.auto_save_max_file_size,
max_concurrent=config.auto_save_max_concurrent
)
data = result.get("data", {})
if isinstance(data, list):
images = data
elif isinstance(data, dict) and "data" in data:
images = data["data"]
else:
images = [data]
image_data = []
for i, image in enumerate(images):
if isinstance(image, dict) and "b64_json" in image:
image_data.append({
'b64_json': image['b64_json'],
'prompt': prompt,
'custom_name': f"{custom_name}_{i+1}" if custom_name else None,
'alt_text': f"Generated image {i+1}: {prompt[:50]}..."
})
if not image_data:
logger.warning("未找到可保存的Base64图片数据")
return []
auto_save_results = await auto_save_manager.save_multiple_base64_images(
image_data, tool_name="sequential_generation"
)
logger.info(f"Base64 自动保存完成: {len(auto_save_results)} 个图片")
return auto_save_results
except Exception as e:
logger.error(f"Base64 自动保存失败: {e}")
return []
def _update_result_with_auto_save(
result: Dict[str, Any],
auto_save_results: List[AutoSaveResult]
) -> Dict[str, Any]:
"""更新结果以包含自动保存信息
Args:
result: 原始API结果
auto_save_results: 自动保存结果列表
Returns:
更新后的结果
"""
# 创建结果副本
updated_result = result.copy()
# 统计保存结果
successful_saves = sum(1 for r in auto_save_results if r.success)
failed_saves = len(auto_save_results) - successful_saves
# 添加自动保存统计信息
updated_result["auto_save_summary"] = {
"total": len(auto_save_results),
"successful": successful_saves,
"failed": failed_saves
}
# 为成功保存的图片添加本地路径信息
if updated_result.get("data") and auto_save_results:
for i, (item, save_result) in enumerate(zip(updated_result["data"], auto_save_results)):
if save_result.success:
item["local_path"] = str(save_result.local_path)
item["markdown_ref"] = save_result.markdown_ref
return updated_result
def _format_sequential_generation_response(
result: Dict[str, Any],
prompt: str,
max_images: int,
size: str,
auto_save_results: Optional[List[AutoSaveResult]] = None,
auto_save_enabled: bool = False
) -> str:
"""格式化组图生成响应
Args:
result: API响应结果
prompt: 原始提示词
max_images: 最大图像数量
size: 图像尺寸
Returns:
格式化的响应文本
"""
if not result.get("success"):
return f"图像生成失败: {result.get('error', '未知错误')}"
data = result.get("data", {})
usage = result.get("usage", {})
# 构建响应文本
response_lines = [
"✅ 组图生成任务完成",
"",
f"📝 提示词: {prompt}",
f"🔢 请求生成数量: {max_images}张",
f"📏 尺寸: {size}",
""
]
# 处理生成的图像
if isinstance(data, list):
images = data
elif isinstance(data, dict) and "data" in data:
images = data["data"]
else:
images = [data]
if images:
actual_count = len(images)
response_lines.append(f"🎨 实际生成图像: {actual_count}张")
response_lines.append("")
for i, image in enumerate(images, 1):
response_lines.append(f"📷 图像 {i}:")
if isinstance(image, dict):
# URL信息（如存在）
if "url" in image:
response_lines.append(f" • URL: {image['url']}")
# Base64信息（如存在）
if "b64_json" in image:
response_lines.append(f" • 数据: [Base64编码，长度: {len(image['b64_json'])}字符]")
# 自动保存后的本地路径与引用（如存在）
if "local_path" in image:
response_lines.append(f" • 💾 本地路径: {image['local_path']}")
if "markdown_ref" in image:
response_lines.append(f" • 📝 Markdown引用: {image['markdown_ref']}")
# 修订提示词（如存在）
if "revised_prompt" in image:
response_lines.append(f" • 修订提示词: {image['revised_prompt']}")
else:
response_lines.append(f" • {str(image)}")
response_lines.append("")
# 生成数量统计
if actual_count != max_images:
response_lines.append(f"ℹ️ 注意: 请求生成{max_images}张，实际生成{actual_count}张")
response_lines.append("")
# 添加使用统计
if usage:
response_lines.extend([
"📊 使用统计:"
])
if "prompt_tokens" in usage:
response_lines.append(f" • 提示词令牌数: {usage['prompt_tokens']}")
if "total_tokens" in usage:
response_lines.append(f" • 总令牌数: {usage['total_tokens']}")
if "cost" in usage:
response_lines.append(f" • 费用: {usage['cost']}")
response_lines.append("")
# 添加自动保存摘要（如果启用）
if auto_save_enabled and auto_save_results:
successful_saves = sum(1 for r in auto_save_results if r.success)
failed_saves = len(auto_save_results) - successful_saves
response_lines.extend([
"",
"💾 自动保存摘要:",
f" • 总计: {len(auto_save_results)}张图片",
f" • 成功: {successful_saves}张",
f" • 失败: {failed_saves}张"
])
# 添加组图生成说明
response_lines.extend([
"",
"💡 组图生成说明:",
" • 组图生成会基于同一个提示词生成多张不同的图像",
" • 每张图像都是独立生成的，会有不同的视觉效果",
" • 适用于需要多个设计方案或创意选择的场景",
" • 可以从生成的多张图像中选择最满意的结果"
])
return "\n".join(response_lines)
