From 2b3593dfe019080602402014cdfd32ae03cdc1fa Mon Sep 17 00:00:00 2001
From: cherishsince <cherishsince@aliyun.com>
Date: Wed, 5 Jun 2024 15:22:10 +0800
Subject: [PATCH] =?UTF-8?q?=E3=80=90=E4=BC=98=E5=8C=96=E3=80=91=E6=8A=BD?=
 =?UTF-8?q?=E7=A6=BB=20midjourneyNotify=20=E9=9C=80=E8=A6=81=E6=9B=B4?=
 =?UTF-8?q?=E6=96=B0=E7=9A=84=20image=20=E5=AF=B9=E8=B1=A1=EF=BC=8C?=
 =?UTF-8?q?=E6=96=B9=E4=BE=BF=E5=9C=A8job=E4=B8=AD=E5=85=AC=E7=94=A8?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../ai/service/image/AiImageService.java      |  9 ++++
 .../ai/service/image/AiImageServiceImpl.java  | 44 ++++++++++++++-----
 2 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/image/AiImageService.java b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/image/AiImageService.java
index d7395f0e3..f84bfcb21 100644
--- a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/image/AiImageService.java
+++ b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/image/AiImageService.java
@@ -64,6 +64,15 @@ public interface AiImageService {
      */
     Boolean midjourneyNotify(MidjourneyNotifyReqVO notifyReqVO);
 
+    /**
+     * 构建 midjourney - 更新对象
+     *
+     * @param imageId
+     * @param notifyReqVO
+     * @return
+     */
+    AiImageDO buildUpdateImage(Long imageId, MidjourneyNotifyReqVO notifyReqVO);
+
     /**
      * midjourney - action(放大、缩小、U1、U2...)
      *
diff --git a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/image/AiImageServiceImpl.java b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/image/AiImageServiceImpl.java
index 2a0dfbd1f..e4909f988 100644
--- a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/image/AiImageServiceImpl.java
+++ b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/image/AiImageServiceImpl.java
@@ -193,6 +193,13 @@ public class AiImageServiceImpl implements AiImageService {
             log.warn("midjourneyNotify 回调的 jobId 不存在! jobId: {}", notifyReqVO.getId());
             return false;
         }
+        // 2、转换状态
+        AiImageDO updateImage = buildUpdateImage(image.getId(), notifyReqVO);
+        // 3、更新 image 状态
+        return imageMapper.updateById(updateImage) > 0;
+    }
+
+    public AiImageDO buildUpdateImage(Long imageId, MidjourneyNotifyReqVO notifyReqVO) {
         // 2、转换状态
         String imageStatus = null;
         MidjourneyTaskStatusEnum taskStatusEnum = MidjourneyTaskStatusEnum.valueOf(notifyReqVO.getStatus());
@@ -211,17 +218,14 @@ public class AiImageServiceImpl implements AiImageService {
             }
         }
         // 4、更新 image 状态
-        imageMapper.updateById(
-                new AiImageDO()
-                        .setId(image.getId())
-                        .setStatus(imageStatus)
-                        .setPicUrl(filePath)
-                        .setProgress(notifyReqVO.getProgress())
-                        .setResponse(notifyReqVO)
-                        .setButtons(notifyReqVO.getButtons())
-                        .setErrorMessage(notifyReqVO.getFailReason())
-        );
-        return true;
+        return new AiImageDO()
+                .setId(imageId)
+                .setStatus(imageStatus)
+                .setPicUrl(filePath)
+                .setProgress(notifyReqVO.getProgress())
+                .setResponse(notifyReqVO)
+                .setButtons(notifyReqVO.getButtons())
+                .setErrorMessage(notifyReqVO.getFailReason());
     }
 
     @Override
@@ -234,12 +238,28 @@ public class AiImageServiceImpl implements AiImageService {
             throw exception(AI_IMAGE_CUSTOM_ID_NOT_EXISTS);
         }
         // 3、调用 midjourney proxy
-        midjourneyProxyClient.action(
+        MidjourneySubmitRespVO submitRespVO = midjourneyProxyClient.action(
                 new MidjourneyActionReqVO()
                         .setCustomId(customId)
                         .setTaskId(aiImageDO.getTaskId())
                         .setNotifyHook(midjourneyNotifyUrl)
         );
+        // 6、保存任务 id (状态码: 1(提交成功), 21(已存在), 22(排队中), other(错误))
+        if (!MidjourneySubmitCodeEnum.SUCCESS_CODES.contains(submitRespVO.getCode())) {
+            throw exception(AI_IMAGE_MIDJOURNEY_SUBMIT_FAIL, submitRespVO.getDescription());
+        }
+        // 4、新增 image 记录
+        AiImageDO newImage = BeanUtils.toBean(aiImageDO, AiImageDO.class);
+        // 4.1、重置参数
+        newImage.setId(null);
+        newImage.setStatus(AiImageStatusEnum.IN_PROGRESS.getStatus());
+        newImage.setPicUrl(null);
+        newImage.setResponse(null);
+        newImage.setProgress(null);
+        newImage.setTaskId(submitRespVO.getResult());
+        newImage.setErrorMessage(null);
+        // 4.2、保存数据库
+        imageMapper.insert(newImage);
         return Boolean.TRUE;
     }