generate 联调修改

src/main/java/com/ai/da/common/RabbitMQ/GenerateConsumer.java

@@ -109,6 +109,7 @@ public class GenerateConsumer {
         long start = System.currentTimeMillis();
 
         Map<String, String> generateResult = JSONObject.parseObject(msg.getBody(), Map.class);
+        log.info("tasks_id : {} start ",generateResult.get("tasks_id"));
 //        log.info("tasks_id : {}, message : {}",generateResult.get("tasks_id"), generateResult.get("message") );
         if (generateResult.get("status").equals("SUCCESS")){
             String url = generateResult.get("data");
@@ -118,7 +119,7 @@ public class GenerateConsumer {
             // 修改redis中的数据状态为exception
             String key = generateResultKey + ":" + generateResult.get("tasks_id");
             Long expire = redisUtil.getExpire(key);
-            redisUtil.addToString(key, new Gson().toJson(new GenerateResultVO(null, null, "Fail")), expire);
+            redisUtil.addToString(key, new Gson().toJson(new GenerateResultVO(generateResult.get("tasks_id"), null, null, "Fail")), expire);
             // 将异常信息存到exception中
             HashMap<String, String> exceptionInfo = new HashMap<>();
             exceptionInfo.put(generateResult.get("tasks_id"), generateResult.get("data"));

src/main/java/com/ai/da/controller/GenerateController.java

@@ -33,12 +33,12 @@ public class GenerateController {
         return Response.success(generateService.generateCaption(sketchElementId));
     }
 
-    @ApiOperation("通过文字、图片生成图片")
+    /*@ApiOperation("通过文字、图片生成图片")
     @PostMapping("/sketchAndPrint")
     public void generateThroughImageText(@Valid @RequestBody GenerateThroughImageTextDTO generateThroughImageTextDTO) {
 //        return Response.success(generateService.generateThroughImageText(generateThroughImageTextDTO));
         generateService.generateThroughImageText(generateThroughImageTextDTO);
-    }
+    }*/
 
     @ApiOperation("喜欢生成的图片")
     @PostMapping("/like")
@@ -55,14 +55,14 @@ public class GenerateController {
 
     @ApiOperation(value = "发起生成请求，异步获取结果")
     @PostMapping("/prepare")
-    public Response<List<String>> prepareForGenerate(@Valid @RequestBody GenerateThroughImageTextDTO generateThroughImageTextDTO) {
+    public Response<PrepareForGenerateVO> prepareForGenerate(@Valid @RequestBody GenerateThroughImageTextDTO generateThroughImageTextDTO) {
         return Response.success(generateService.prepareForGenerate(generateThroughImageTextDTO));
     }
 
     @ApiOperation(value = "取消继续生成")
     @GetMapping("/stopWaiting")
     public Response<String> stopWaiting(@RequestParam("userId") Long userId,
-                                        @RequestParam("uniqueId") String uniqueId,
+                                        @RequestParam("uniqueId") List<String> uniqueId,
                                         @RequestParam("timeZone") String timeZone) {
         generateService.cancelGenerate(userId, uniqueId, timeZone);
         return Response.success("stop waiting successfully");

src/main/java/com/ai/da/model/dto/GenerateThroughImageTextDTO.java

@@ -48,6 +48,7 @@ public class GenerateThroughImageTextDTO {
     @ApiModelProperty("唯一id，用于保持消息唯一性")
     String uniqueId;
 
+    @NotNull(message = "Please check if the required fields are empty.(isTestUser)")
     @ApiModelProperty("是否是测试用户")
     Boolean isTestUser;
 }

src/main/java/com/ai/da/model/vo/GenerateResultVO.java

@@ -9,9 +9,17 @@ import lombok.Data;
 @AllArgsConstructor
 public class GenerateResultVO {
 
+    private String taskId;
+
     private Long id;
 
     private String url;
 
     private String status;
+
+    public GenerateResultVO(Long id, String url, String status) {
+        this.id = id;
+        this.url = url;
+        this.status = status;
+    }
 }

src/main/java/com/ai/da/model/vo/PrepareForGenerateVO.java

@@ -4,17 +4,19 @@ import io.swagger.annotations.ApiModel;
 import io.swagger.annotations.ApiModelProperty;
 import lombok.Data;
 
+import java.util.List;
+
 @Data
 @ApiModel("prepare for generate响应vo")
 public class PrepareForGenerateVO {
 
     @ApiModelProperty("uniqueId")
-    private String uniqueId;
+    private List<String> uniqueId;
 
     @ApiModelProperty("剩余使用次数")
     private Integer leftUsageCount;
 
-    public PrepareForGenerateVO(String uniqueId, Integer leftUsageCount) {
+    public PrepareForGenerateVO(List<String> uniqueId, Integer leftUsageCount) {
         this.uniqueId = uniqueId;
         this.leftUsageCount = leftUsageCount;
     }

src/main/java/com/ai/da/service/GenerateService.java

@@ -29,10 +29,10 @@ public interface GenerateService extends IService<Generate> {
 
     List<GenerateResultVO> getGenerateResultList(List<String> taskIdList);
 
-    List<String> prepareForGenerate(GenerateThroughImageTextDTO generateThroughImageTextDTO);
+    PrepareForGenerateVO prepareForGenerate(GenerateThroughImageTextDTO generateThroughImageTextDTO);
 
     Long getRankPosition(String uniqueId);
 
-    void cancelGenerate(Long userId, String uniqueId, String timeZone);
+    void cancelGenerate(Long userId, List<String> uniqueId, String timeZone);
 
 }

src/main/java/com/ai/da/service/impl/GenerateServiceImpl.java

@@ -23,6 +23,7 @@ import com.ai.da.service.RabbitMQService;
 import com.alibaba.fastjson.JSON;
 import com.alibaba.fastjson.JSONObject;
 import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
+import com.baomidou.mybatisplus.core.exceptions.MybatisPlusException;
 import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
 import com.google.gson.Gson;
 import io.minio.errors.MinioException;
@@ -153,7 +154,7 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
         }else {
             status = "Fail";
         }
-        GenerateResultVO generateResultVO = new GenerateResultVO(null, null, status);
+        GenerateResultVO generateResultVO = new GenerateResultVO(generateThroughImageTextDTO.getUniqueId(), null, null, status);
         redisUtil.addToString(key, new Gson().toJson(generateResultVO), CommonConstant.GENERATE_RESULT_EXPIRE_TIME);
 
         // 5、处理模型返回的数据
@@ -193,7 +194,19 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
         // 5.1 将相应的url保存到数据库
         GenerateDetail generateDetail = new GenerateDetail();
         GenerateCollectionItemVO generateCollectionItemVO = new GenerateCollectionItemVO();
-        Generate generate = selectByUniqueId(taskId);
+        Generate generate ;
+        try{
+            generate = selectByUniqueId(taskId);
+        }catch (MybatisPlusException e){
+            log.error(e.getMessage());
+            if (e.getMessage().equals("One record is expected, but the query result is multiple records")){
+                generate = selectListByUniqueId(taskId).get(0);
+            }else {
+                throw new BusinessException("There are some problems with database query, please try again.");
+            }
+
+        }
+//        Generate generate = selectByUniqueId(taskId);
         String md5 = MD5Utils.encryptFile(minioUtil.getPresignedUrl(url, 24 * 60), Boolean.FALSE);
         // 通过MD5值和level1Type,判断不同level1Type下相同的图片是否被like过
         List<Map<String, Long>> libraryIdList = generateDetailMapper.getLibraryIdThroughMD5(md5, generate.getLevel1Type());
@@ -210,7 +223,7 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
 
         String key = generateResultKey + ":" + taskId;
         Long expire = redisUtil.getExpire(key);
-        GenerateResultVO generateResultVO = new GenerateResultVO(generateDetail.getId(), url, "Success");
+        GenerateResultVO generateResultVO = new GenerateResultVO(taskId, generateDetail.getId(), url, "Success");
         redisUtil.addToString(key, new Gson().toJson(generateResultVO), expire);
     }
 
@@ -360,8 +373,8 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
     }
 
     @Override
-//    public PrepareForGenerateVO prepareForGenerate(GenerateThroughImageTextDTO generateThroughImageTextDTO) {
+    public PrepareForGenerateVO prepareForGenerate(GenerateThroughImageTextDTO generateThroughImageTextDTO) {
-    public List<String> prepareForGenerate(GenerateThroughImageTextDTO generateThroughImageTextDTO) {
+//    public List<String> prepareForGenerate(GenerateThroughImageTextDTO generateThroughImageTextDTO) {
         // 1、参数检查，判断必须参数是否为空
         if (Objects.isNull(generateThroughImageTextDTO.getUserId())) {
             throw new BusinessException("userId cannot be empty");
@@ -376,7 +389,7 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
         if (generateThroughImageTextDTO.getIsTestUser()){
             trialsCount = getTrialsCount(generateThroughImageTextDTO.getUserId(), generateThroughImageTextDTO.getLevel1Type());
             if (trialsCount >= 2){
-                return new ArrayList<>();
+                return new PrepareForGenerateVO(0);
             }
         }
 
@@ -422,7 +435,7 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
         }
 
         // 5、返回唯一id
-        return taskIdList;
+        return new PrepareForGenerateVO(taskIdList, 2 - trialsCount);
     }
 
     @Override
@@ -504,9 +517,17 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
         return getOne(qw);
     }
 
+    public List<Generate> selectListByUniqueId(String uniqueId) {
+        QueryWrapper<Generate> qw = new QueryWrapper<>();
+        qw.eq("unique_id", uniqueId).orderByDesc("id");
+
+        return baseMapper.selectList(qw);
+    }
+
     @Override
     @Transactional(rollbackFor = Exception.class)
-    public void cancelGenerate(Long userId, String uniqueId, String timeZone) {
+    public void cancelGenerate(Long userId, List<String> uniqueIdList, String timeZone) {
+        uniqueIdList.forEach(uniqueId -> {
             // 1、确认当前消息是否还在排队中
             Boolean exists = redisUtil.isElementExistsInZSet(consumptionOrderKey, uniqueId);
             Boolean flag = Boolean.FALSE;
@@ -529,21 +550,23 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
                     pythonService.cancelGenerateTask(uniqueId);
                 }
             }
-
             // 3、考虑加一张表，专门用于记录哪些用户在什么时间进行了取消操作,包括已经异常的请求
             GenerateCancel generateCancel = new GenerateCancel(userId, uniqueId, DateUtil.getByTimeZone(timeZone));
             generateCancelMapper.insert(generateCancel);
+        });
+
+
     }
 
-    // 判断试用用户试用generate机会是否使用完毕
+    // 判断试用用户试用generate机会是否使用完毕 每个board 3次机会
     private int getTrialsCount(Long userId, String level1Type){
         List<Generate> getGenerateList = getGenerateByAccountId(userId, level1Type);
         int trialsCount ;
         if (getGenerateList.isEmpty()){
             trialsCount = 0;
-        } else if (getGenerateList.size() == 1) {
+        } else if (getGenerateList.size() == 1 || (getGenerateList.size() >= 4 && getGenerateList.size() / 4 == 1)) {
             trialsCount = 1;
-        } else if (getGenerateList.size() == 2) {
+        } else if (getGenerateList.size() == 2 || (getGenerateList.size() >= 4 && getGenerateList.size() / 4 == 2)) {
             trialsCount = 2;
         }else {
             trialsCount = 2;