From 80787d1dcca7e3cabc5c483e318ef0db3668be37 Mon Sep 17 00:00:00 2001
From: cherishsince <cherishsince@aliyun.com>
Date: Mon, 29 Apr 2024 14:49:09 +0800
Subject: [PATCH] =?UTF-8?q?=E3=80=90=E4=BC=98=E5=8C=96=E3=80=91MidjourneyM?=
 =?UTF-8?q?essageListener=20=E5=A2=9E=E5=8A=A0=20message=20handler=20?=
 =?UTF-8?q?=E5=9B=9E=E8=B0=83=EF=BC=8C=E7=94=A8=E4=BA=8E=E5=A4=84=E7=90=86?=
 =?UTF-8?q?=E4=B8=9A=E5=8A=A1=E6=95=B0=E6=8D=AE?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../webSocket/MidjourneyMessageHandler.java     | 15 +++++++++++++++
 .../listener/MidjourneyMessageListener.java     | 17 ++++++++++++-----
 2 files changed, 27 insertions(+), 5 deletions(-)
 create mode 100644 yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/MidjourneyMessageHandler.java

diff --git a/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/MidjourneyMessageHandler.java b/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/MidjourneyMessageHandler.java
new file mode 100644
index 000000000..d5522fc82
--- /dev/null
+++ b/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/MidjourneyMessageHandler.java
@@ -0,0 +1,15 @@
+package cn.iocoder.yudao.framework.ai.midjourney.webSocket;
+
+import cn.iocoder.yudao.framework.ai.midjourney.MidjourneyMessage;
+
+/**
+ * message handler
+ *
+ * @author fansili
+ * @time 2024/4/29 14:29
+ * @since 1.0
+ */
+public interface MidjourneyMessageHandler {
+
+    void messageHandler(MidjourneyMessage midjourneyMessage);
+}
diff --git a/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/listener/MidjourneyMessageListener.java b/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/listener/MidjourneyMessageListener.java
index a044f3613..e85c4e8f7 100644
--- a/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/listener/MidjourneyMessageListener.java
+++ b/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/listener/MidjourneyMessageListener.java
@@ -4,13 +4,13 @@ package cn.iocoder.yudao.framework.ai.midjourney.webSocket.listener;
 import cn.hutool.core.collection.CollUtil;
 import cn.hutool.core.text.CharSequenceUtil;
 import cn.hutool.core.util.StrUtil;
-import cn.hutool.json.JSONUtil;
 import cn.iocoder.yudao.framework.ai.midjourney.MidjourneyConfig;
 import cn.iocoder.yudao.framework.ai.midjourney.MidjourneyMessage;
 import cn.iocoder.yudao.framework.ai.midjourney.constants.MidjourneyConstants;
 import cn.iocoder.yudao.framework.ai.midjourney.constants.MidjourneyGennerateStatusEnum;
 import cn.iocoder.yudao.framework.ai.midjourney.constants.MidjourneyMessageTypeEnum;
 import cn.iocoder.yudao.framework.ai.midjourney.util.MidjourneyUtil;
+import cn.iocoder.yudao.framework.ai.midjourney.webSocket.MidjourneyMessageHandler;
 import cn.iocoder.yudao.framework.common.util.json.JsonUtils;
 import lombok.extern.slf4j.Slf4j;
 import net.dv8tion.jda.api.utils.data.DataObject;
@@ -21,11 +21,18 @@ import java.util.List;
 public class MidjourneyMessageListener {
 
     private MidjourneyConfig midjourneyConfig;
+    private MidjourneyMessageHandler midjourneyMessageHandler = null;
 
     public MidjourneyMessageListener(MidjourneyConfig midjourneyConfig) {
         this.midjourneyConfig = midjourneyConfig;
     }
 
+    public MidjourneyMessageListener(MidjourneyConfig midjourneyConfig,
+                                     MidjourneyMessageHandler midjourneyMessageHandler) {
+        this.midjourneyConfig = midjourneyConfig;
+        this.midjourneyMessageHandler = midjourneyMessageHandler;
+    }
+
     public void onMessage(DataObject raw) {
         MidjourneyMessageTypeEnum messageType = MidjourneyMessageTypeEnum.of(raw.getString("t"));
         if (messageType == null || MidjourneyMessageTypeEnum.DELETE == messageType) {
@@ -35,7 +42,6 @@ public class MidjourneyMessageListener {
         if (ignoreAndLogMessage(data, messageType)) {
             return;
         }
-
         // 转换几个重要的信息
         MidjourneyMessage mjMessage = new MidjourneyMessage();
         mjMessage.setId(data.getString(MidjourneyConstants.MSG_ID));
@@ -56,9 +62,10 @@ public class MidjourneyMessageListener {
         }
         // 转换状态
         convertGenerateStatus(mjMessage);
-        //
-        log.info("message 信息 {}", JSONUtil.toJsonPrettyStr(mjMessage));
-        System.err.println(JSONUtil.toJsonPrettyStr(mjMessage));
+        // message handler 调用
+        if (midjourneyMessageHandler != null) {
+            midjourneyMessageHandler.messageHandler(mjMessage);
+        }
     }
 
     private void convertGenerateStatus(MidjourneyMessage mjMessage) {