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

This commit is contained in:
cherishsince 2024-04-29 14:49:09 +08:00
parent 6bd1d943eb
commit 80787d1dcc
2 changed files with 27 additions and 5 deletions

View File

@ -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);
}

View File

@ -4,13 +4,13 @@ package cn.iocoder.yudao.framework.ai.midjourney.webSocket.listener;
import cn.hutool.core.collection.CollUtil; import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.text.CharSequenceUtil; import cn.hutool.core.text.CharSequenceUtil;
import cn.hutool.core.util.StrUtil; 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.MidjourneyConfig;
import cn.iocoder.yudao.framework.ai.midjourney.MidjourneyMessage; 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.MidjourneyConstants;
import cn.iocoder.yudao.framework.ai.midjourney.constants.MidjourneyGennerateStatusEnum; 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.constants.MidjourneyMessageTypeEnum;
import cn.iocoder.yudao.framework.ai.midjourney.util.MidjourneyUtil; 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 cn.iocoder.yudao.framework.common.util.json.JsonUtils;
import lombok.extern.slf4j.Slf4j; import lombok.extern.slf4j.Slf4j;
import net.dv8tion.jda.api.utils.data.DataObject; import net.dv8tion.jda.api.utils.data.DataObject;
@ -21,11 +21,18 @@ import java.util.List;
public class MidjourneyMessageListener { public class MidjourneyMessageListener {
private MidjourneyConfig midjourneyConfig; private MidjourneyConfig midjourneyConfig;
private MidjourneyMessageHandler midjourneyMessageHandler = null;
public MidjourneyMessageListener(MidjourneyConfig midjourneyConfig) { public MidjourneyMessageListener(MidjourneyConfig midjourneyConfig) {
this.midjourneyConfig = midjourneyConfig; this.midjourneyConfig = midjourneyConfig;
} }
public MidjourneyMessageListener(MidjourneyConfig midjourneyConfig,
MidjourneyMessageHandler midjourneyMessageHandler) {
this.midjourneyConfig = midjourneyConfig;
this.midjourneyMessageHandler = midjourneyMessageHandler;
}
public void onMessage(DataObject raw) { public void onMessage(DataObject raw) {
MidjourneyMessageTypeEnum messageType = MidjourneyMessageTypeEnum.of(raw.getString("t")); MidjourneyMessageTypeEnum messageType = MidjourneyMessageTypeEnum.of(raw.getString("t"));
if (messageType == null || MidjourneyMessageTypeEnum.DELETE == messageType) { if (messageType == null || MidjourneyMessageTypeEnum.DELETE == messageType) {
@ -35,7 +42,6 @@ public class MidjourneyMessageListener {
if (ignoreAndLogMessage(data, messageType)) { if (ignoreAndLogMessage(data, messageType)) {
return; return;
} }
// 转换几个重要的信息 // 转换几个重要的信息
MidjourneyMessage mjMessage = new MidjourneyMessage(); MidjourneyMessage mjMessage = new MidjourneyMessage();
mjMessage.setId(data.getString(MidjourneyConstants.MSG_ID)); mjMessage.setId(data.getString(MidjourneyConstants.MSG_ID));
@ -56,9 +62,10 @@ public class MidjourneyMessageListener {
} }
// 转换状态 // 转换状态
convertGenerateStatus(mjMessage); convertGenerateStatus(mjMessage);
// // message handler 调用
log.info("message 信息 {}", JSONUtil.toJsonPrettyStr(mjMessage)); if (midjourneyMessageHandler != null) {
System.err.println(JSONUtil.toJsonPrettyStr(mjMessage)); midjourneyMessageHandler.messageHandler(mjMessage);
}
} }
private void convertGenerateStatus(MidjourneyMessage mjMessage) { private void convertGenerateStatus(MidjourneyMessage mjMessage) {