mirror of
https://gitee.com/mateos/mateclaw.git
synced 2026-09-13 19:23:42 +08:00
201 lines
9.0 KiB
Java
201 lines
9.0 KiB
Java
package vip.mate.tool.builtin;
|
||
|
||
import lombok.RequiredArgsConstructor;
|
||
import lombok.extern.slf4j.Slf4j;
|
||
import org.springframework.ai.chat.model.ToolContext;
|
||
import org.springframework.ai.tool.annotation.Tool;
|
||
import org.springframework.ai.tool.annotation.ToolParam;
|
||
import org.springframework.lang.Nullable;
|
||
import org.springframework.stereotype.Component;
|
||
import vip.mate.system.model.SystemSettingsDTO;
|
||
import vip.mate.system.service.SystemSettingService;
|
||
import vip.mate.task.AsyncTaskService;
|
||
import vip.mate.task.model.AsyncTaskInfo;
|
||
import vip.mate.tool.image.*;
|
||
|
||
import java.util.ArrayList;
|
||
import java.util.List;
|
||
import java.util.StringJoiner;
|
||
|
||
/**
|
||
* 图片生成工具 — Agent 可调用的 @Tool，提交图片生成任务
|
||
*
|
||
* @author MateClaw Team
|
||
*/
|
||
@Slf4j
|
||
@Component
|
||
@RequiredArgsConstructor
|
||
public class ImageGenerateTool {
|
||
|
||
private final ImageGenerationService imageGenerationService;
|
||
private final ImageProviderRegistry providerRegistry;
|
||
private final SystemSettingService systemSettingService;
|
||
private final AsyncTaskService asyncTaskService;
|
||
private final ImageReferenceLoader imageReferenceLoader;
|
||
|
||
@vip.mate.tool.ConcurrencyUnsafe("creates async tasks and persists generated artifacts; provider rate limits also forbid parallel calls")
|
||
@Tool(description = "Image generation tool. Supports actions: generate (default — text-to-image, OR image-edit when "
|
||
+ "image/images parameters are set), list (show available providers/models), status (check task status). "
|
||
+ "Reference images may be local paths, http(s) URLs, data: URLs, or msg:<messageId>:<idx> for an attachment "
|
||
+ "from an earlier conversation message. Async providers take 30s-2min; results auto-display in the conversation.")
|
||
public String image_generate(
|
||
@ToolParam(description = "Action type: generate, list, status. Default: generate", required = false) String action,
|
||
@ToolParam(description = "Image content description, be detailed (required for generate)", required = false) String prompt,
|
||
@ToolParam(description = "Single reference image for edit mode. Path / http(s) URL / data: URL / msg:<messageId>[:<idx>]", required = false) String image,
|
||
@ToolParam(description = "Multiple reference images for edit mode (provider caps the count). Same formats as 'image'.", required = false) List<String> images,
|
||
@ToolParam(description = "Image size: 1024x1024 / 1024x1792 / 1792x1024", required = false) String size,
|
||
@ToolParam(description = "Aspect ratio: 1:1 / 16:9 / 9:16, default 1:1", required = false) String aspectRatio,
|
||
@ToolParam(description = "Generation count (1-4), default 1", required = false) Integer count,
|
||
@ToolParam(description = "Model name (optional)", required = false) String model,
|
||
@ToolParam(description = "Task ID to check status (for status action)", required = false) String taskId,
|
||
@Nullable ToolContext ctx
|
||
) {
|
||
String normalizedAction = (action == null || action.isBlank()) ? "generate" : action.trim().toLowerCase();
|
||
|
||
return switch (normalizedAction) {
|
||
case "list" -> handleListAction();
|
||
case "status" -> handleStatusAction(taskId, ctx);
|
||
default -> handleGenerateAction(prompt, image, images, size, aspectRatio, count, model, ctx);
|
||
};
|
||
}
|
||
|
||
// ==================== action=list ====================
|
||
|
||
private String handleListAction() {
|
||
SystemSettingsDTO config = systemSettingService.getAllSettings();
|
||
List<ImageGenerationProvider> providers = providerRegistry.allSorted();
|
||
|
||
if (providers.isEmpty()) {
|
||
return "当前没有注册的图片生成 Provider。";
|
||
}
|
||
|
||
StringJoiner sb = new StringJoiner("\n\n");
|
||
sb.add("## 可用的图片生成 Provider\n");
|
||
|
||
for (ImageGenerationProvider p : providers) {
|
||
boolean available = p.isAvailable(config);
|
||
ImageProviderCapabilities caps = p.detailedCapabilities();
|
||
|
||
StringJoiner entry = new StringJoiner("\n");
|
||
entry.add("### " + p.label() + " (" + p.id() + ") " + (available ? "[已配置]" : "[未配置]"));
|
||
|
||
if (caps != null) {
|
||
if (caps.getModels() != null && !caps.getModels().isEmpty()) {
|
||
entry.add("- 模型: " + String.join(", ", caps.getModels()));
|
||
}
|
||
entry.add("- 支持尺寸: " + String.join(", ", caps.getSupportedSizes()));
|
||
entry.add("- 最大数量: " + caps.getMaxCount());
|
||
}
|
||
sb.add(entry.toString());
|
||
}
|
||
return sb.toString();
|
||
}
|
||
|
||
// ==================== action=status ====================
|
||
|
||
private String handleStatusAction(String taskId, @Nullable ToolContext ctx) {
|
||
String conversationId = ToolExecutionContext.conversationId(ctx);
|
||
|
||
if (taskId != null && !taskId.isBlank()) {
|
||
AsyncTaskInfo info = imageGenerationService.checkTaskStatus(taskId);
|
||
if (info == null) {
|
||
return "未找到任务 ID: " + taskId;
|
||
}
|
||
return formatTaskStatus(info);
|
||
}
|
||
|
||
if (conversationId == null) {
|
||
return "无法获取当前会话信息";
|
||
}
|
||
List<AsyncTaskInfo> activeTasks = asyncTaskService.listActiveTasks(conversationId);
|
||
List<AsyncTaskInfo> imageTasks = activeTasks.stream()
|
||
.filter(t -> "image_generation".equals(t.getTaskType()))
|
||
.toList();
|
||
if (imageTasks.isEmpty()) {
|
||
return "当前会话没有进行中的图片生成任务。";
|
||
}
|
||
|
||
StringJoiner sb = new StringJoiner("\n");
|
||
sb.add("当前会话有 " + imageTasks.size() + " 个进行中的任务：");
|
||
for (AsyncTaskInfo task : imageTasks) {
|
||
sb.add("- 任务 " + task.getTaskId() + ": " + formatTaskStatus(task));
|
||
}
|
||
return sb.toString();
|
||
}
|
||
|
||
// ==================== action=generate ====================
|
||
|
||
private String handleGenerateAction(String prompt, String image, List<String> images,
|
||
String size, String aspectRatio,
|
||
Integer count, String model, @Nullable ToolContext ctx) {
|
||
String conversationId = ToolExecutionContext.conversationId(ctx);
|
||
String username = ToolExecutionContext.username(ctx);
|
||
|
||
if (conversationId == null || conversationId.isBlank()) {
|
||
return "错误：无法获取当前会话信息，请重试";
|
||
}
|
||
|
||
if (prompt == null || prompt.isBlank()) {
|
||
return "错误：prompt 为必填参数，请描述你想要生成的图片内容";
|
||
}
|
||
|
||
// Combine the singular and plural forms — the agent picks whichever is
|
||
// ergonomic. Order: image (first) then images[].
|
||
List<String> referenceInputs = new ArrayList<>();
|
||
if (image != null && !image.isBlank()) {
|
||
referenceInputs.add(image);
|
||
}
|
||
if (images != null) {
|
||
for (String s : images) {
|
||
if (s != null && !s.isBlank()) referenceInputs.add(s);
|
||
}
|
||
}
|
||
|
||
List<ImageReference> inputImages;
|
||
try {
|
||
inputImages = imageReferenceLoader.loadAll(referenceInputs, conversationId);
|
||
} catch (Exception e) {
|
||
log.warn("[ImageGenerateTool] Failed to load reference images: {}", e.getMessage());
|
||
return "错误：无法加载参考图片：" + e.getMessage();
|
||
}
|
||
|
||
ImageGenerationRequest request = ImageGenerationRequest.builder()
|
||
.prompt(prompt)
|
||
.size(size)
|
||
.aspectRatio(aspectRatio != null ? aspectRatio : "1:1")
|
||
.count(count != null ? count : 1)
|
||
.model(model)
|
||
.inputImages(inputImages)
|
||
.build();
|
||
|
||
ImageGenerationResult result = imageGenerationService.submitGeneration(
|
||
request, conversationId, username != null ? username : "system");
|
||
|
||
if (result.isCompleted()) {
|
||
// 同步模式：图片已生成
|
||
return result.getMessage();
|
||
} else if (result.isSubmitted()) {
|
||
// 异步模式：已提交
|
||
return result.getMessage();
|
||
} else {
|
||
return "图片生成失败：" + result.getMessage();
|
||
}
|
||
}
|
||
|
||
// ==================== 辅助方法 ====================
|
||
|
||
private String formatTaskStatus(AsyncTaskInfo info) {
|
||
return switch (info.getStatus()) {
|
||
case "pending" -> "排队中，请稍候...";
|
||
case "running" -> {
|
||
String progressStr = info.getProgress() != null && info.getProgress() > 0
|
||
? "（进度: " + info.getProgress() + "%）" : "";
|
||
yield "生成中" + progressStr + "（" + info.getProviderName() + "）";
|
||
}
|
||
case "succeeded" -> "已完成，图片已显示在对话中";
|
||
case "failed" -> "失败：" + (info.getErrorMessage() != null ? info.getErrorMessage() : "未知错误");
|
||
default -> "状态: " + info.getStatus();
|
||
};
|
||
}
|
||
}
|