Merge branch 'mj' into pre

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

@@ -19,6 +19,7 @@ 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.MidjourneyService;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
@@ -52,6 +53,8 @@ public class MidjourneyServiceImpl implements MidjourneyService {
 
     private static final List<String> status = List.of("MODAL","CANCEL","FAILURE");
 
+    private final RedisService redisService;
+
     private static final GlobalThreadPoolTaskExecutor TASK_EXECUTOR = GlobalThreadPoolTaskExecutor.getInstance();
 
 
@@ -121,7 +124,17 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         return saveConversation(user.getId(), user.getMode(), Long.parseLong(result.getResult()),result.getProperties(),"SHORTEN", StringUtils.EMPTY);
     }
 
-
+    /**
+     * 恢复次数
+     */
+    public void recoverUserLimit(Long id,Integer mode){
+        if (mode == 1){
+            redisService.incr(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + id, 1L);
+        }
+        if (mode == 2){
+            redisService.incr(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + id, 1L);
+        }
+    }
     @Override
     public MidjourneyUserConversation submitAction(MidjourneyUser user, Long taskId, String customId) throws Exception {
         MidjourneyUserConversation conversation = conversationMapper.selectOne(Wrappers.lambdaQuery(MidjourneyUserConversation.class).eq(MidjourneyUserConversation::getTaskId, taskId).last("limit 1"));
@@ -137,6 +150,13 @@ public class MidjourneyServiceImpl implements MidjourneyService {
                 if (button.getCustomId().equals(customId)) {
                     log.info("action customId:{}", customId);
                     button.setStyle(3);
+                    if (button.getLabel().contains("Vary") || button.getCustomId().contains("::pan_")
+                            || button.getEmoji().equals("🔄") || button.getCustomId().contains("PromptAnalyzer:")
+                            || button.getCustomId().contains("PicReader::") || button.getCustomId().contains("::variation::")
+                            || button.getCustomId().contains("::CustomZoom::")) {
+                        // 以上操作有弹窗确认，恢复次数
+                        recoverUserLimit(user.getId(), user.getMode());
+                    }
                 }
             });
             conversation.setButtons(Jsons.toJson(messageButtons));

netflix-web/src/main/java/com/cyksj/web/controller/mirror/MidjourneyController.java

@@ -1,5 +1,6 @@
 package com.cyksj.web.controller.mirror;
 
+import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
 import com.baomidou.mybatisplus.core.toolkit.Wrappers;
 import com.cyksj.common.exception.BusinessRuntimeException;
 import com.cyksj.dto.Result;
@@ -49,6 +50,10 @@ public class MidjourneyController {
         if (midjourneyUser == null){
             midjourneyUser = midjourneyUserMapper.selectOne(Wrappers.lambdaQuery(MidjourneyUser.class)
                     .eq(MidjourneyUser::getUserToken, userToken).last("limit 1"));
+            if (midjourneyUser == null){
+                throw new BusinessRuntimeException("账号不存在,请重新登录");
+            }
+            redisService.set(RedisService.key.MIDJOURNEY_USER.getName() + userToken, midjourneyUser, RedisService.key.MIDJOURNEY_USER.getTimeout());
             Object fastNum = redisService.get(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + midjourneyUser.getId());
             if (fastNum == null){
                 redisService.set(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + midjourneyUser.getId(), midjourneyUser.getMjFastNum(), RedisService.key.MIDJOURNEY_FAST_LIMIT.getTimeout());
@@ -82,39 +87,60 @@ public class MidjourneyController {
     /**
      * 检查用户次数
      */
-    public void checkUserLimit(MidjourneyUser user){
-        boolean flag = true;
-        Long fastNum = null;
-        Long relaxNum = null;
+    public Long checkUserLimit(MidjourneyUser user){
+        Long num = 0L;
+        if (user.getExpireTime().before(new Date())){
+            return num;
+        }
        if (user.getMode() == 1){
-           if (user.getExpireTime().before(new Date())){
-               flag = false;
-           }
-           fastNum = redisService.decr(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + user.getId(), 1L);
-           if (fastNum < 0) {
+           num = redisService.decr(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + user.getId(), 1L);
+           if (num < 0) {
                user.setMode(2);
-               fastNum = redisService.incr(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + user.getId(), 1L);
+               num = redisService.incr(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + user.getId(), 1L);
            }
        }
        if (user.getMode() == 2){
            Object relax = redisService.get(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + user.getId());
            if (relax != null){
-               relaxNum = Long.parseLong(relax.toString());
-               if (relaxNum > 0){
-                   relaxNum = redisService.decr(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + user.getId(), 1L);
-               }else {
-                   flag = false;
+               num = redisService.decr(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + user.getId(), 1L);
+               if (num < 0){
+                   num = redisService.incr(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + user.getId(), 1L);
                }
+           }else {
+               num = null;
            }
        }
-       if (!flag){
-           throw BusinessRuntimeException.getInstance("可用次数不足");
-       }else {
-           midjourneyUserMapper.update(null, Wrappers.lambdaUpdate(MidjourneyUser.class)
-                   .eq(MidjourneyUser::getId, user.getId()).set(MidjourneyUser::getMode, user.getMode())
-                   .set(fastNum != null,MidjourneyUser::getMjFastNum, fastNum)
-                   .set(relaxNum != null,MidjourneyUser::getMjRelaxNum, relaxNum));
-       }
+       return num;
+    }
+
+    /**
+     * 恢复次数
+     */
+    public void recoverUserLimit(Long id,Integer mode){
+        if (mode == 1){
+            redisService.incr(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + id, 1L);
+        }
+        if (mode == 2){
+            redisService.incr(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + id, 1L);
+        }
+    }
+
+    /**
+     * 同步数据库
+     */
+    public void syncUser(Long id,Integer mode,Long num){
+        if (num == null){
+            return;
+        }
+        LambdaUpdateWrapper<MidjourneyUser> wrapper = Wrappers.lambdaUpdate(MidjourneyUser.class)
+                .eq(MidjourneyUser::getId, id);
+        if (mode == 1){
+            wrapper.set(MidjourneyUser::getMjFastNum, num);
+        }
+        if (mode == 2){
+            wrapper.set(MidjourneyUser::getMjRelaxNum, num);
+        }
+        midjourneyUserMapper.update(null, wrapper);
     }
 
 
@@ -132,11 +158,21 @@ public class MidjourneyController {
      * 提交Imagine任务
      */
     @PostMapping("/submit/imagine")
-    public Result<MidjourneyUserConversation> submitImagine(@RequestBody SubmitImagineDTO submitImagineDTO)throws Exception{
+    public Result<MidjourneyUserConversation> submitImagine(@RequestBody SubmitImagineDTO submitImagineDTO) {
         log.info("提交Imagine任务，提示：{}，base64数组长度：{}",submitImagineDTO.getPrompt(),submitImagineDTO.getBase64Array());
         MidjourneyUser user = getUser();
-        MidjourneyUserConversation conversation = midjourneyService.submitImagine(user, submitImagineDTO.getPrompt(), submitImagineDTO.getBase64Array());
-        checkUserLimit(user);
+        Long num = checkUserLimit(user);
+        if (num != null && num == 0){
+            throw BusinessRuntimeException.getInstance("次数已用完");
+        }
+        MidjourneyUserConversation conversation;
+        try {
+            conversation = midjourneyService.submitImagine(user, submitImagineDTO.getPrompt(), submitImagineDTO.getBase64Array());
+        } catch (Exception e) {
+            recoverUserLimit(user.getId(), user.getMode());
+            throw BusinessRuntimeException.getInstance(e.getMessage());
+        }
+        syncUser(user.getId(), user.getMode(), num);
         return GatewayResponse.SUCCESS.newBuilder().toResult(conversation);
     }
 
@@ -144,11 +180,21 @@ public class MidjourneyController {
      * 提交Describe任务
      */
     @PostMapping("/submit/describe")
-    public Result<MidjourneyUserConversation> submitDescribe(@RequestBody SubmitDescribeDTO submitDescribeDTO)throws Exception{
+    public Result<MidjourneyUserConversation> submitDescribe(@RequestBody SubmitDescribeDTO submitDescribeDTO) {
         log.info("提交Describe任务");
         MidjourneyUser user = getUser();
-        MidjourneyUserConversation conversation = midjourneyService.submitDescribe(user, submitDescribeDTO.getBase64());
-        checkUserLimit(user);
+        Long num = checkUserLimit(user);
+        if (num != null && num == 0){
+            throw BusinessRuntimeException.getInstance("次数已用完");
+        }
+        MidjourneyUserConversation conversation;
+        try {
+            conversation = midjourneyService.submitDescribe(user, submitDescribeDTO.getBase64());
+        } catch (Exception e) {
+            recoverUserLimit(user.getId(), user.getMode());
+            throw BusinessRuntimeException.getInstance(e.getMessage());
+        }
+        syncUser(user.getId(), user.getMode(), num);
         return GatewayResponse.SUCCESS.newBuilder().toResult(conversation);
     }
 
@@ -156,11 +202,21 @@ public class MidjourneyController {
      * 提交Blend任务
      */
     @PostMapping("/submit/blend")
-    public Result<MidjourneyUserConversation> submitBlend(@RequestBody SubmitBlendDTO submitBlendDTO) throws Exception{
+    public Result<MidjourneyUserConversation> submitBlend(@RequestBody SubmitBlendDTO submitBlendDTO) {
         log.info("提交Blend任务 dimensions:{}，base64数组长度：{}", submitBlendDTO.getDimensions(), submitBlendDTO.getBase64Array().size());
         MidjourneyUser user = getUser();
-        MidjourneyUserConversation conversation = midjourneyService.submitBlend(user, submitBlendDTO.getDimensions(), submitBlendDTO.getBase64Array());
-        checkUserLimit(user);
+        Long num = checkUserLimit(user);
+        if (num != null && num == 0){
+            throw BusinessRuntimeException.getInstance("次数已用完");
+        }
+        MidjourneyUserConversation conversation;
+        try {
+            conversation = midjourneyService.submitBlend(user, submitBlendDTO.getDimensions(), submitBlendDTO.getBase64Array());
+        } catch (Exception e) {
+            recoverUserLimit(user.getId(), user.getMode());
+            throw BusinessRuntimeException.getInstance(e.getMessage());
+        }
+        syncUser(user.getId(), user.getMode(), num);
         return GatewayResponse.SUCCESS.newBuilder().toResult(conversation);
     }
 
@@ -168,11 +224,21 @@ public class MidjourneyController {
      * 提交Modal任务
      */
     @PostMapping("/submit/modal")
-    public Result<MidjourneyUserConversation> submitModal(@RequestBody SubmitModalDTO submitModalDTO) throws Exception{
+    public Result<MidjourneyUserConversation> submitModal(@RequestBody SubmitModalDTO submitModalDTO) {
         log.info("提交Modal任务，taskId:{},提示：{}，base64数组长度：{}",submitModalDTO.getTaskId(),submitModalDTO.getPrompt(),submitModalDTO.getMaskBase64());
         MidjourneyUser user = getUser();
-        MidjourneyUserConversation conversation = midjourneyService.submitModal(user, submitModalDTO.getTaskId(), submitModalDTO.getPrompt(), submitModalDTO.getMaskBase64());
-        checkUserLimit(user);
+        Long num = checkUserLimit(user);
+        if (num != null && num == 0){
+            throw BusinessRuntimeException.getInstance("次数已用完");
+        }
+        MidjourneyUserConversation conversation;
+        try {
+            conversation = midjourneyService.submitModal(user, submitModalDTO.getTaskId(), submitModalDTO.getPrompt(), submitModalDTO.getMaskBase64());
+        } catch (Exception e) {
+            recoverUserLimit(user.getId(), user.getMode());
+            throw BusinessRuntimeException.getInstance(e.getMessage());
+        }
+        syncUser(user.getId(), user.getMode(), num);
         return GatewayResponse.SUCCESS.newBuilder().toResult(conversation);
     }
 
@@ -180,11 +246,21 @@ public class MidjourneyController {
      * 提交Shorten任务
      */
     @PostMapping("/submit/shorten")
-    public Result<MidjourneyUserConversation> submitShorten(@RequestBody SubmitShortenDTO submitShortenDTO) throws Exception {
+    public Result<MidjourneyUserConversation> submitShorten(@RequestBody SubmitShortenDTO submitShortenDTO) {
         log.info("提交Shorten任务 提示词:{}",submitShortenDTO.getPrompt());
         MidjourneyUser user = getUser();
-        MidjourneyUserConversation conversation = midjourneyService.submitShorten(user, submitShortenDTO.getPrompt());
-        checkUserLimit(user);
+        Long num = checkUserLimit(user);
+        if (num != null && num == 0){
+            throw BusinessRuntimeException.getInstance("次数已用完");
+        }
+        MidjourneyUserConversation conversation;
+        try {
+            conversation = midjourneyService.submitShorten(user, submitShortenDTO.getPrompt());
+        } catch (Exception e) {
+            recoverUserLimit(user.getId(), user.getMode());
+            throw BusinessRuntimeException.getInstance(e.getMessage());
+        }
+        syncUser(user.getId(), user.getMode(), num);
         return GatewayResponse.SUCCESS.newBuilder().toResult(conversation);
     }
 
@@ -192,11 +268,21 @@ public class MidjourneyController {
      * 执行动作
      */
     @PostMapping("/submit/action")
-    public Result<MidjourneyUserConversation> action(@RequestBody SubmitActionDTO actionDTO) throws Exception {
+    public Result<MidjourneyUserConversation> action(@RequestBody SubmitActionDTO actionDTO) {
         log.info("任务id:{},执行动作：{}",actionDTO.getTaskId(),actionDTO.getCustomId());
         MidjourneyUser user = getUser();
-        MidjourneyUserConversation conversation = midjourneyService.submitAction(user, actionDTO.getTaskId(), actionDTO.getCustomId());
-        checkUserLimit(user);
+        Long num = checkUserLimit(user);
+        if (num != null && num == 0){
+            throw BusinessRuntimeException.getInstance("次数已用完");
+        }
+        MidjourneyUserConversation conversation;
+        try {
+            conversation = midjourneyService.submitAction(user, actionDTO.getTaskId(), actionDTO.getCustomId());
+        } catch (Exception e) {
+            recoverUserLimit(user.getId(), user.getMode());
+            throw BusinessRuntimeException.getInstance(e.getMessage());
+        }
+        syncUser(user.getId(), user.getMode(), num);
         return GatewayResponse.SUCCESS.newBuilder().toResult(conversation);
     }