Просмотр исходного кода

fix 发送顺序以及返回值剔除无用字段

chenbiao 2 лет назад
Родитель
Сommit
96b0a31ac5

+ 21 - 0
midjourney/src/main/java/com/yhlxj/dao/model/entity/MidjourneyUserConversation.java

@@ -1,9 +1,14 @@
 package com.yhlxj.dao.model.entity;
 
+import com.baomidou.mybatisplus.annotation.TableField;
+import com.ejlchina.searcher.bean.DbIgnore;
 import com.ejlchina.searcher.bean.SearchBean;
+import com.fasterxml.jackson.annotation.JsonInclude;
 import com.yhlxj.dao.model.BaseEntity;
 import lombok.Data;
+import lombok.NoArgsConstructor;
 import lombok.experimental.Accessors;
+import org.apache.commons.lang3.StringUtils;
 
 import java.io.Serializable;
 
@@ -14,6 +19,7 @@ import java.io.Serializable;
 @Data
 @Accessors(chain = true)
 @SearchBean(tables = "midjourney_user_conversation")
+@JsonInclude(JsonInclude.Include.NON_NULL)
 public class MidjourneyUserConversation extends BaseEntity implements Serializable {
 
     private Long userId;
@@ -104,4 +110,19 @@ public class MidjourneyUserConversation extends BaseEntity implements Serializab
      */
     private Boolean bookmark;
 
+    @DbIgnore
+    @TableField(exist = false)
+    private Integer progressNum;
+
+    public MidjourneyUserConversation setProgress(String progress) {
+        this.progress = progress;
+        try {
+            if(StringUtils.isNotBlank(progress)){
+                this.progressNum = Integer.parseInt(this.progress.replace("%",""));
+            }
+        }catch (Exception e) {
+
+        }
+        return this;
+    }
 }

+ 1 - 1
midjourney/src/main/java/com/yhlxj/redis/RedisService.java

@@ -753,7 +753,7 @@ public class RedisService {
         MIDJOURNEY_CONVERSATION("midjourney:conversation:", "midjourney conversation", 60 * 60 * 2L),
         MIDJOURNEY_ACCOUNT("midjourney:account:", "midjourney account", 60 * 60 * 48L),
         MIDJOURNEY_QUERY("midjourney:query:", "midjourney query", 60 * 60 * 2L),
-        MIDJOURNEY_WS_SESSION("midjourney:ws:session", "midjourney ws session", 60 * 60 * 30L),
+        MIDJOURNEY_PROGRESS("midjourney:task:progress:","task progress 任务进度", 60 * 60 * 1L),
         MIDJOURNEY_USER_TIME("midjourney:user:time:", "midjourney user time", 60 * 60 * 24L),
         MIDJOURNEY_USER_SUBMIT("midjourney:user:submit:", "midjourney user submit", 30L),
         ;

+ 60 - 6
midjourney/src/main/java/com/yhlxj/service/midjourney/impl/MidjourneyServiceImpl.java

@@ -46,13 +46,14 @@ import javax.imageio.stream.ImageInputStream;
 import java.awt.image.BufferedImage;
 import java.io.File;
 import java.io.IOException;
-import java.net.MalformedURLException;
 import java.net.URL;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
 import static com.yhlxj.web.wss.MidjourneyServerEndpoint.WSS_SESSION_MAP;
@@ -727,11 +728,7 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         redisService.set(RedisService.key.MIDJOURNEY_CONVERSATION.getName() + conversation.getTaskId(), conversation,RedisService.key.MIDJOURNEY_CONVERSATION.getTimeout());
 
         TASK_EXECUTOR.execute(()->{
-            //WebSocket 直接发送 任务状态
-            WssSession wssSession = WSS_SESSION_MAP.get(conversation.getUserId());
-            if (wssSession != null) {
-                wssSession.sendMessage(WsMessageTypeEnum.SERVER_TASK_INFO, new JSONObject(conversation));
-            }
+            handleCallback(conversation);
             conversationMapper.update(null, QueryWrapperUtils.buildUpdateWrapper((wrapper)->{
                 wrapper.eq(MidjourneyUserConversation::getAction,"MODEL").eq(MidjourneyUserConversation::getTaskId, conversation.getTaskId()).set(MidjourneyUserConversation::getAction, conversation.getAction());
             }));
@@ -746,6 +743,63 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         });
     }
 
+    public void handleCallback(MidjourneyUserConversation conversation) {
+       String key = RedisService.key.MIDJOURNEY_PROGRESS.getName() + conversation.getTaskId();
+
+        synchronized (key.intern()) {
+            String currentProgressStr = redisService.getStr(key);
+            int currentProgress = currentProgressStr == null ? -1 : Integer.parseInt(currentProgressStr);
+            if (conversation.getProgressNum() > currentProgress) {
+                redisService.set(key, String.valueOf(progress), RedisService.key.MIDJOURNEY_PROGRESS.getTimeout());
+                sendProgressToClient(conversation);
+            }
+        }
+    }
+
+    private void sendProgressToClient( MidjourneyUserConversation conversation) {
+        //WebSocket 直接发送 任务状态
+        WssSession wssSession = WSS_SESSION_MAP.get(conversation.getUserId());
+        if (wssSession != null) {
+            conversation.setUserId(null);
+            conversation.setChannelId(null);
+            conversation.setInstanceId(null);
+            if(JSONUtil.isJson(conversation.getProperties())){
+                //cancelComponent
+                //discordChannelId
+                //discordInstanceId
+                //messageContent 这个需要替换中间字符
+                //messageHash
+                //notifyHook
+                JSONObject jsonObject = new JSONObject(conversation.getProperties());
+                jsonObject.remove("cancelComponent");
+                jsonObject.remove("discordChannelId");
+                jsonObject.remove("discordInstanceId");
+                jsonObject.remove("messageHash");
+                jsonObject.remove("notifyHook");
+                jsonObject.remove("nonce");
+                String content = jsonObject.getStr("messageContent");
+                jsonObject.set("messageContent", replaceUserId(content,"mj"));
+                conversation.setProperties(jsonObject.toString());
+            }
+
+            wssSession.sendMessage(WsMessageTypeEnum.SERVER_TASK_INFO, new JSONObject(conversation));
+        }
+    }
+
+    public static String replaceUserId(String messageContent, String newUserId) {
+        // 定义匹配用户ID的正则表达式模式
+        String regex = "<@[^>]+>";
+        Pattern pattern = Pattern.compile(regex);
+        Matcher matcher = pattern.matcher(messageContent);
+
+        // 使用新的用户ID替换第一个匹配的用户ID
+        if (matcher.find()) {
+            messageContent = matcher.replaceFirst("<@" + newUserId + ">");
+        }
+
+        return messageContent;
+    }
+
     @Data
     public class ImgSize{
         /**

+ 1 - 1
midjourney/src/main/java/com/yhlxj/web/wss/MidjourneyServerEndpoint.java

@@ -44,7 +44,7 @@ public class MidjourneyServerEndpoint {
     public void onOpen(Session session, @PathParam("userToken") String userToken, @PathParam("taskId") String taskId) throws IOException {
         log.info("[WS建立连接],token:{}, taskId:{}", userToken, taskId);
 
-        session.setMaxIdleTimeout(TimeUnit.MINUTES.toMillis(5));
+        session.setMaxIdleTimeout(TimeUnit.MINUTES.toMillis(30));
         //检查uniodId是否合法
         if (StringUtils.isBlank(userToken) || StringUtils.isBlank(taskId)) {