generate 添加试用用户判断

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

@@ -6,6 +6,7 @@ import com.ai.da.model.dto.GenerateThroughImageTextDTO;
 import com.ai.da.model.vo.GenerateCaptionVO;
 import com.ai.da.model.vo.GenerateCollectionVO;
 import com.ai.da.model.vo.GenerateLikeVO;
+import com.ai.da.model.vo.PrepareForGenerateVO;
 import com.ai.da.service.GenerateService;
 import io.swagger.annotations.Api;
 import io.swagger.annotations.ApiOperation;
@@ -55,7 +56,7 @@ public class GenerateController {
 
     @ApiOperation(value = "发起生成请求，异步获取结果")
     @PostMapping("/prepare")
-    public Response<String> prepareForGenerate(@Valid @RequestBody GenerateThroughImageTextDTO generateThroughImageTextDTO) {
+    public Response<PrepareForGenerateVO> prepareForGenerate(@Valid @RequestBody GenerateThroughImageTextDTO generateThroughImageTextDTO) {
         return Response.success(generateService.prepareForGenerate(generateThroughImageTextDTO));
     }
 

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

@@ -47,4 +47,7 @@ public class GenerateThroughImageTextDTO {
 
     @ApiModelProperty("唯一id，用于保持消息唯一性")
     String uniqueId;
+
+    @ApiModelProperty("是否是测试用户")
+    Boolean isTestUser;
 }

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

@@ -0,0 +1,25 @@
+package com.ai.da.model.vo;
+
+import io.swagger.annotations.ApiModel;
+import io.swagger.annotations.ApiModelProperty;
+import lombok.Data;
+
+@Data
+@ApiModel("prepare for generate响应vo")
+public class PrepareForGenerateVO {
+
+    @ApiModelProperty("uniqueId")
+    private String uniqueId;
+
+    @ApiModelProperty("剩余使用次数")
+    private Integer leftUsageCount;
+
+    public PrepareForGenerateVO(String uniqueId, Integer leftUsageCount) {
+        this.uniqueId = uniqueId;
+        this.leftUsageCount = leftUsageCount;
+    }
+
+    public PrepareForGenerateVO(Integer leftUsageCount) {
+        this.leftUsageCount = leftUsageCount;
+    }
+}

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

@@ -7,6 +7,7 @@ import com.ai.da.model.dto.GenerateThroughImageTextDTO;
 import com.ai.da.model.vo.GenerateCaptionVO;
 import com.ai.da.model.vo.GenerateCollectionVO;
 import com.ai.da.model.vo.GenerateLikeVO;
+import com.ai.da.model.vo.PrepareForGenerateVO;
 import com.baomidou.mybatisplus.extension.service.IService;
 
 import java.util.List;
@@ -27,7 +28,7 @@ public interface GenerateService extends IService<Generate> {
 
     GenerateCollectionVO getGenerateResult(String uniqueId);
 
-    String prepareForGenerate(GenerateThroughImageTextDTO generateThroughImageTextDTO);
+    PrepareForGenerateVO prepareForGenerate(GenerateThroughImageTextDTO generateThroughImageTextDTO);
 
     Long getRankPosition(String uniqueId);
 

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

@@ -314,7 +314,7 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
     }
 
     @Override
-    public String prepareForGenerate(GenerateThroughImageTextDTO generateThroughImageTextDTO) {
+    public PrepareForGenerateVO prepareForGenerate(GenerateThroughImageTextDTO generateThroughImageTextDTO) {
         // 1、参数检查，判断必须参数是否为空
         if (Objects.isNull(generateThroughImageTextDTO.getUserId())) {
             throw new BusinessException("userId cannot be empty");
@@ -323,6 +323,16 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
         if (!GenerateModeEnum.getGenerateModeList().contains(generateType)) {
             throw new BusinessException("unknown.generate.type");
         }
+
+        // 判断试用用户是否还有剩余试用机会
+        int trialsCount = 0;
+        if (generateThroughImageTextDTO.getIsTestUser()){
+            trialsCount = getTrialsCount(generateThroughImageTextDTO.getUserId(), generateThroughImageTextDTO.getLevel1Type());
+            if (trialsCount >= 2){
+                return new PrepareForGenerateVO(0);
+            }
+        }
+
         String text = generateThroughImageTextDTO.getText();
         Long elementId = generateThroughImageTextDTO.getCollectionElementId();
         validateGeneraType(new Generate(), text, elementId, generateType);
@@ -361,7 +371,7 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
         rabbitMQService.publishMessage(jsonString);
 
         // 5、返回唯一id
-        return uuid;
+        return new PrepareForGenerateVO(uuid, 2 - trialsCount);
     }
 
     @Override
@@ -458,4 +468,28 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
         GenerateCancel generateCancel = new GenerateCancel(userId, uniqueId, DateUtil.getByTimeZone(timeZone));
         generateCancelMapper.insert(generateCancel);
     }
+
+    // 判断试用用户试用generate机会是否使用完毕
+    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) {
+            trialsCount = 1;
+        } else if (getGenerateList.size() == 2) {
+            trialsCount = 2;
+        }else {
+            trialsCount = 2;
+        }
+        return trialsCount;
+    }
+
+    public List<Generate> getGenerateByAccountId(Long accountId, String level1Type){
+        QueryWrapper<Generate> qw = new QueryWrapper<>();
+        qw.eq("account_id",accountId);
+        qw.eq("level1_type", level1Type);
+
+        return baseMapper.selectList(qw);
+    }
 }