fix midjourney

netflix-service/src/main/java/com/cyksj/redis/RedisService.java

@@ -840,6 +840,7 @@ public class RedisService {
         MIDJOURNEY_EXPIRE_TIME("midjourney:expire:time:", "midjourney expire time", 60 * 60 * 48L),
         MIDJOURNEY_USER("midjourney:user:", "midjourney user", 60 * 60 * 2L),
         MIDJOURNEY_CONVERSATION("midjourney:conversation:", "midjourney conversation", 60 * 60 * 2L),
+        MIDJOURNEY_ACCOUNT("midjourney:account:", "midjourney account", 60 * 60 * 48L),
         ;
 
         private String name;

netflix-service/src/main/java/com/cyksj/service/midjourney/impl/MidjourneyServiceImpl.java

@@ -20,10 +20,12 @@ import com.cyksj.mapper.MidjourneyUserConversationMapper;
 import com.cyksj.mapper.MidjourneyUserMapper;
 import com.cyksj.model.dto.BlendDimensions;
 import com.cyksj.model.dto.MessageButton;
+import com.cyksj.model.entity.MidjourneyAccount;
 import com.cyksj.model.entity.MidjourneyUser;
 import com.cyksj.model.entity.MidjourneyUserConversation;
 import com.cyksj.model.response.SubmitResult;
 import com.cyksj.redis.RedisService;
+import com.cyksj.service.midjourney.MidjourneyAccountService;
 import com.cyksj.service.midjourney.MidjourneyService;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
@@ -76,6 +78,8 @@ public class MidjourneyServiceImpl implements MidjourneyService {
 
     private static final GlobalThreadPoolTaskExecutor TASK_EXECUTOR = GlobalThreadPoolTaskExecutor.getInstance();
 
+    private final MidjourneyAccountService midjourneyAccountService;
+
 
     @Override
     public MidjourneyUserConversation submitImagine(MidjourneyUser user, String prompt, List<String> base64Array) throws Exception {
@@ -153,15 +157,16 @@ public class MidjourneyServiceImpl implements MidjourneyService {
     /**
      * 恢复次数
      */
-    public void recoverUserLimit(Long id,Integer mode,Integer num){
+    public Long recoverUserLimit(Long id,Integer mode,Long num){
         if (mode == 1){
-            redisService.incr(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + id, 1L);
+            num =  redisService.incr(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + id, 1L);
         }
         if (mode == 2){
             if (num != null) {
-                redisService.incr(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + id, 1L);
+                num = redisService.incr(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + id, 1L);
             }
         }
+        return num;
     }
     /**
      * 同步数据库
@@ -218,7 +223,7 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         SubmitResult result = submit(user.getMode(),"action", param);
         if (flag.get()) {
             // 以上操作有弹窗确认，恢复次数
-            recoverUserLimit(user.getId(), user.getMode(),user.getMjRelaxNum());
+            recoverUserLimit(user.getId(), user.getMode(),num);
         }else {
             syncUser(user.getId(), user.getMode(),num);
         }
@@ -227,10 +232,20 @@ public class MidjourneyServiceImpl implements MidjourneyService {
 
 
     public SubmitResult submit(Integer mode,String action, Map<String, Object> param) throws Exception {
+        String url = "";
         if (mode == 1) {
             param.put("mode", "FAST");
+            url = FAST_HOST;
+        } else if (mode == 2) {
+            url = RELAX_HOST;
+            if (action.equals("action")) {
+                //慢速查询在使用次数最少的账号
+                String accountWithMinUsage = getAccountWithMinUsage();
+                param.put("accountFilter",MapUtil.builder(new HashMap<String,Object>())
+                        .put("instanceId",accountWithMinUsage).build());
+            }
         }
-        String url = (mode == 1 ? FAST_HOST : RELAX_HOST) + getActionUrl(action);
+        url = url + getActionUrl(action);
         param.put("notifyHook",midjourneyHost +(EnvCommonService.active.equals(envCommonService.getEnv()) ? "/8081":"/8082") + "/api/applets/midjourney/notifyHook");
         String body = HttpRequest.post(url).body(Jsons.toJson(param)).header("Authorization", FAST_TOKEN).execute().body();
         log.info("action body:{}", body);
@@ -251,6 +266,38 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         }
         return submitResult;
     }
+
+    public String getAccountWithMinUsage() {
+        String key = RedisService.key.MIDJOURNEY_ACCOUNT.getName();
+        try {
+            // 获取所有账号ID和使用次数
+            Map<Object, Object> accountsUsage = redisService.hmget(key);
+
+            if (accountsUsage == null || accountsUsage.isEmpty()) {
+                Map<String, Object> map = midjourneyAccountService.list(Wrappers.lambdaQuery(MidjourneyAccount.class)).stream().collect(Collectors.toMap(k -> k.getInstanceId().toString(), v -> 0));
+                redisService.hmset(key, map);
+            }
+
+            // 找到使用次数最少的账号ID
+            String minAccountId = null;
+            int minUsage = Integer.MAX_VALUE;
+
+            for (Map.Entry<Object, Object> entry : accountsUsage.entrySet()) {
+                int usage = Integer.parseInt(entry.getValue().toString());
+                if (usage < minUsage) {
+                    minUsage = usage;
+                    minAccountId = entry.getKey().toString();
+                }
+            }
+            // 增加使用次数
+            if (minAccountId != null) {
+                redisService.hincr(key, minAccountId, 1.0);
+            }
+            return minAccountId;
+        } catch (Exception e) {
+            throw BusinessRuntimeException.getInstance("获取账号失败");
+        }
+    }
     private String getActionUrl(String action) {
          switch (action) {
             case "imagine":
@@ -281,8 +328,9 @@ public class MidjourneyServiceImpl implements MidjourneyService {
     public void sync(MidjourneyUserConversation conversation) throws IOException {
         if ((StringUtils.isNotBlank(conversation.getProgress()) && progress.contains(conversation.getProgress())) || status.contains(conversation.getStatus())) {
             log.info("同步任务：{},进度:{}",conversation.getTaskId(),conversation.getProgress());
+            Long userId = conversation.getUserId();
             MidjourneyUserConversation dbConversation = conversationMapper.selectOne(Wrappers.lambdaQuery(MidjourneyUserConversation.class).eq(MidjourneyUserConversation::getTaskId, conversation.getTaskId())
-                    .eq(MidjourneyUserConversation::getUserId, conversation.getUserId()).last("limit 1"));
+                    .eq(MidjourneyUserConversation::getUserId, userId).last("limit 1"));
             if (dbConversation != null) {
                 try {
                     String imageUrl = conversation.getImageUrl();
@@ -294,6 +342,28 @@ public class MidjourneyServiceImpl implements MidjourneyService {
                 }finally {
                     conversation.setId(dbConversation.getId());
                     conversationMapper.updateById(conversation);
+                    if (dbConversation.getMode() == 2){
+                        redisService.hdecr(RedisService.key.MIDJOURNEY_ACCOUNT.getName() , conversation.getInstanceId().toString(), 1.0);
+                    }
+                }
+                //失败返还次数
+
+                if ("FAILURE".equals(conversation.getStatus())){
+                    Object num = redisService.get(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + userId);
+                    LambdaUpdateWrapper<MidjourneyUser> wrapper = Wrappers.lambdaUpdate(MidjourneyUser.class)
+                            .eq(MidjourneyUser::getId, userId);
+                    Integer mode = dbConversation.getMode();
+                    if (mode == 1){
+                        num =  redisService.incr(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + userId, 1L);
+                        wrapper.set(MidjourneyUser::getMjFastNum, num);
+                    }
+                    if (mode == 2){
+                        if (num != null) {
+                            num = redisService.incr(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + userId, 1L);
+                        }
+                        wrapper.set(MidjourneyUser::getMjRelaxNum, num);
+                    }
+                    midjourneyUserMapper.update(null, wrapper);
                 }
             }
         }