Commit 80787d1dcca7e3cabc5c483e318ef0db3668be37

Authored by cherishsince
1 parent 6bd1d943

【优化】MidjourneyMessageListener 增加 message handler 回调,用于处理业务数据

yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/MidjourneyMessageHandler.java 0 → 100644
  1 +package cn.iocoder.yudao.framework.ai.midjourney.webSocket;
  2 +
  3 +import cn.iocoder.yudao.framework.ai.midjourney.MidjourneyMessage;
  4 +
  5 +/**
  6 + * message handler
  7 + *
  8 + * @author fansili
  9 + * @time 2024/4/29 14:29
  10 + * @since 1.0
  11 + */
  12 +public interface MidjourneyMessageHandler {
  13 +
  14 + void messageHandler(MidjourneyMessage midjourneyMessage);
  15 +}
... ...
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;
4 4 import cn.hutool.core.collection.CollUtil;
5 5 import cn.hutool.core.text.CharSequenceUtil;
6 6 import cn.hutool.core.util.StrUtil;
7   -import cn.hutool.json.JSONUtil;
8 7 import cn.iocoder.yudao.framework.ai.midjourney.MidjourneyConfig;
9 8 import cn.iocoder.yudao.framework.ai.midjourney.MidjourneyMessage;
10 9 import cn.iocoder.yudao.framework.ai.midjourney.constants.MidjourneyConstants;
11 10 import cn.iocoder.yudao.framework.ai.midjourney.constants.MidjourneyGennerateStatusEnum;
12 11 import cn.iocoder.yudao.framework.ai.midjourney.constants.MidjourneyMessageTypeEnum;
13 12 import cn.iocoder.yudao.framework.ai.midjourney.util.MidjourneyUtil;
  13 +import cn.iocoder.yudao.framework.ai.midjourney.webSocket.MidjourneyMessageHandler;
14 14 import cn.iocoder.yudao.framework.common.util.json.JsonUtils;
15 15 import lombok.extern.slf4j.Slf4j;
16 16 import net.dv8tion.jda.api.utils.data.DataObject;
... ... @@ -21,11 +21,18 @@ import java.util.List;
21 21 public class MidjourneyMessageListener {
22 22  
23 23 private MidjourneyConfig midjourneyConfig;
  24 + private MidjourneyMessageHandler midjourneyMessageHandler = null;
24 25  
25 26 public MidjourneyMessageListener(MidjourneyConfig midjourneyConfig) {
26 27 this.midjourneyConfig = midjourneyConfig;
27 28 }
28 29  
  30 + public MidjourneyMessageListener(MidjourneyConfig midjourneyConfig,
  31 + MidjourneyMessageHandler midjourneyMessageHandler) {
  32 + this.midjourneyConfig = midjourneyConfig;
  33 + this.midjourneyMessageHandler = midjourneyMessageHandler;
  34 + }
  35 +
29 36 public void onMessage(DataObject raw) {
30 37 MidjourneyMessageTypeEnum messageType = MidjourneyMessageTypeEnum.of(raw.getString("t"));
31 38 if (messageType == null || MidjourneyMessageTypeEnum.DELETE == messageType) {
... ... @@ -35,7 +42,6 @@ public class MidjourneyMessageListener {
35 42 if (ignoreAndLogMessage(data, messageType)) {
36 43 return;
37 44 }
38   -
39 45 // 转换几个重要的信息
40 46 MidjourneyMessage mjMessage = new MidjourneyMessage();
41 47 mjMessage.setId(data.getString(MidjourneyConstants.MSG_ID));
... ... @@ -56,9 +62,10 @@ public class MidjourneyMessageListener {
56 62 }
57 63 // 转换状态
58 64 convertGenerateStatus(mjMessage);
59   - //
60   - log.info("message 信息 {}", JSONUtil.toJsonPrettyStr(mjMessage));
61   - System.err.println(JSONUtil.toJsonPrettyStr(mjMessage));
  65 + // message handler 调用
  66 + if (midjourneyMessageHandler != null) {
  67 + midjourneyMessageHandler.messageHandler(mjMessage);
  68 + }
62 69 }
63 70  
64 71 private void convertGenerateStatus(MidjourneyMessage mjMessage) {
... ...