Add async AI generation tasks
This commit is contained in:
王鹏
@@ -0,0 +1,89 @@
|
||||
package com.ruoyi.generator.domain.front;
|
||||
|
||||
import java.util.Date;
|
||||
import com.ruoyi.common.core.domain.BaseEntity;
|
||||
|
||||
public class AiGenerationTask extends BaseEntity
|
||||
{
|
||||
private static final long serialVersionUID = 1L;
|
||||
|
||||
private Long taskId;
|
||||
private Long generationId;
|
||||
private Long projectId;
|
||||
private Long userId;
|
||||
private String generateType;
|
||||
private String provider;
|
||||
private String model;
|
||||
private String status;
|
||||
private String requestHash;
|
||||
private String requestPayload;
|
||||
private String resultPayload;
|
||||
private String errorCode;
|
||||
private String errorMessage;
|
||||
private Integer attempts;
|
||||
private Integer maxAttempts;
|
||||
private Date nextRetryTime;
|
||||
private String lockedBy;
|
||||
private Date lockedUntil;
|
||||
private Integer progress;
|
||||
private String currentStep;
|
||||
private Integer inputTokens;
|
||||
private Integer outputTokens;
|
||||
private Integer totalTokens;
|
||||
private Integer costCents;
|
||||
private Date startedAt;
|
||||
private Date finishedAt;
|
||||
|
||||
public Long getTaskId() { return taskId; }
|
||||
public void setTaskId(Long taskId) { this.taskId = taskId; }
|
||||
public Long getGenerationId() { return generationId; }
|
||||
public void setGenerationId(Long generationId) { this.generationId = generationId; }
|
||||
public Long getProjectId() { return projectId; }
|
||||
public void setProjectId(Long projectId) { this.projectId = projectId; }
|
||||
public Long getUserId() { return userId; }
|
||||
public void setUserId(Long userId) { this.userId = userId; }
|
||||
public String getGenerateType() { return generateType; }
|
||||
public void setGenerateType(String generateType) { this.generateType = generateType; }
|
||||
public String getProvider() { return provider; }
|
||||
public void setProvider(String provider) { this.provider = provider; }
|
||||
public String getModel() { return model; }
|
||||
public void setModel(String model) { this.model = model; }
|
||||
public String getStatus() { return status; }
|
||||
public void setStatus(String status) { this.status = status; }
|
||||
public String getRequestHash() { return requestHash; }
|
||||
public void setRequestHash(String requestHash) { this.requestHash = requestHash; }
|
||||
public String getRequestPayload() { return requestPayload; }
|
||||
public void setRequestPayload(String requestPayload) { this.requestPayload = requestPayload; }
|
||||
public String getResultPayload() { return resultPayload; }
|
||||
public void setResultPayload(String resultPayload) { this.resultPayload = resultPayload; }
|
||||
public String getErrorCode() { return errorCode; }
|
||||
public void setErrorCode(String errorCode) { this.errorCode = errorCode; }
|
||||
public String getErrorMessage() { return errorMessage; }
|
||||
public void setErrorMessage(String errorMessage) { this.errorMessage = errorMessage; }
|
||||
public Integer getAttempts() { return attempts; }
|
||||
public void setAttempts(Integer attempts) { this.attempts = attempts; }
|
||||
public Integer getMaxAttempts() { return maxAttempts; }
|
||||
public void setMaxAttempts(Integer maxAttempts) { this.maxAttempts = maxAttempts; }
|
||||
public Date getNextRetryTime() { return nextRetryTime; }
|
||||
public void setNextRetryTime(Date nextRetryTime) { this.nextRetryTime = nextRetryTime; }
|
||||
public String getLockedBy() { return lockedBy; }
|
||||
public void setLockedBy(String lockedBy) { this.lockedBy = lockedBy; }
|
||||
public Date getLockedUntil() { return lockedUntil; }
|
||||
public void setLockedUntil(Date lockedUntil) { this.lockedUntil = lockedUntil; }
|
||||
public Integer getProgress() { return progress; }
|
||||
public void setProgress(Integer progress) { this.progress = progress; }
|
||||
public String getCurrentStep() { return currentStep; }
|
||||
public void setCurrentStep(String currentStep) { this.currentStep = currentStep; }
|
||||
public Integer getInputTokens() { return inputTokens; }
|
||||
public void setInputTokens(Integer inputTokens) { this.inputTokens = inputTokens; }
|
||||
public Integer getOutputTokens() { return outputTokens; }
|
||||
public void setOutputTokens(Integer outputTokens) { this.outputTokens = outputTokens; }
|
||||
public Integer getTotalTokens() { return totalTokens; }
|
||||
public void setTotalTokens(Integer totalTokens) { this.totalTokens = totalTokens; }
|
||||
public Integer getCostCents() { return costCents; }
|
||||
public void setCostCents(Integer costCents) { this.costCents = costCents; }
|
||||
public Date getStartedAt() { return startedAt; }
|
||||
public void setStartedAt(Date startedAt) { this.startedAt = startedAt; }
|
||||
public Date getFinishedAt() { return finishedAt; }
|
||||
public void setFinishedAt(Date finishedAt) { this.finishedAt = finishedAt; }
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
package com.ruoyi.generator.domain.front;
|
||||
|
||||
import com.ruoyi.common.core.domain.BaseEntity;
|
||||
|
||||
public class AiQuotaBucket extends BaseEntity
|
||||
{
|
||||
private static final long serialVersionUID = 1L;
|
||||
|
||||
private Long bucketId;
|
||||
private Long userId;
|
||||
private String periodType;
|
||||
private String periodKey;
|
||||
private Integer taskLimit;
|
||||
private Integer taskUsed;
|
||||
private Integer tokenLimit;
|
||||
private Integer tokenUsed;
|
||||
private Integer costLimitCents;
|
||||
private Integer costUsedCents;
|
||||
private Integer runningLimit;
|
||||
private Integer runningCount;
|
||||
|
||||
public Long getBucketId() { return bucketId; }
|
||||
public void setBucketId(Long bucketId) { this.bucketId = bucketId; }
|
||||
public Long getUserId() { return userId; }
|
||||
public void setUserId(Long userId) { this.userId = userId; }
|
||||
public String getPeriodType() { return periodType; }
|
||||
public void setPeriodType(String periodType) { this.periodType = periodType; }
|
||||
public String getPeriodKey() { return periodKey; }
|
||||
public void setPeriodKey(String periodKey) { this.periodKey = periodKey; }
|
||||
public Integer getTaskLimit() { return taskLimit; }
|
||||
public void setTaskLimit(Integer taskLimit) { this.taskLimit = taskLimit; }
|
||||
public Integer getTaskUsed() { return taskUsed; }
|
||||
public void setTaskUsed(Integer taskUsed) { this.taskUsed = taskUsed; }
|
||||
public Integer getTokenLimit() { return tokenLimit; }
|
||||
public void setTokenLimit(Integer tokenLimit) { this.tokenLimit = tokenLimit; }
|
||||
public Integer getTokenUsed() { return tokenUsed; }
|
||||
public void setTokenUsed(Integer tokenUsed) { this.tokenUsed = tokenUsed; }
|
||||
public Integer getCostLimitCents() { return costLimitCents; }
|
||||
public void setCostLimitCents(Integer costLimitCents) { this.costLimitCents = costLimitCents; }
|
||||
public Integer getCostUsedCents() { return costUsedCents; }
|
||||
public void setCostUsedCents(Integer costUsedCents) { this.costUsedCents = costUsedCents; }
|
||||
public Integer getRunningLimit() { return runningLimit; }
|
||||
public void setRunningLimit(Integer runningLimit) { this.runningLimit = runningLimit; }
|
||||
public Integer getRunningCount() { return runningCount; }
|
||||
public void setRunningCount(Integer runningCount) { this.runningCount = runningCount; }
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package com.ruoyi.generator.domain.front;
|
||||
|
||||
import com.ruoyi.common.core.domain.BaseEntity;
|
||||
|
||||
public class AiUsageLedger extends BaseEntity
|
||||
{
|
||||
private static final long serialVersionUID = 1L;
|
||||
|
||||
private Long ledgerId;
|
||||
private Long taskId;
|
||||
private Long generationId;
|
||||
private Long projectId;
|
||||
private Long userId;
|
||||
private String generateType;
|
||||
private String provider;
|
||||
private String model;
|
||||
private Integer inputTokens;
|
||||
private Integer outputTokens;
|
||||
private Integer totalTokens;
|
||||
private Integer costCents;
|
||||
private String success;
|
||||
private String errorMessage;
|
||||
|
||||
public Long getLedgerId() { return ledgerId; }
|
||||
public void setLedgerId(Long ledgerId) { this.ledgerId = ledgerId; }
|
||||
public Long getTaskId() { return taskId; }
|
||||
public void setTaskId(Long taskId) { this.taskId = taskId; }
|
||||
public Long getGenerationId() { return generationId; }
|
||||
public void setGenerationId(Long generationId) { this.generationId = generationId; }
|
||||
public Long getProjectId() { return projectId; }
|
||||
public void setProjectId(Long projectId) { this.projectId = projectId; }
|
||||
public Long getUserId() { return userId; }
|
||||
public void setUserId(Long userId) { this.userId = userId; }
|
||||
public String getGenerateType() { return generateType; }
|
||||
public void setGenerateType(String generateType) { this.generateType = generateType; }
|
||||
public String getProvider() { return provider; }
|
||||
public void setProvider(String provider) { this.provider = provider; }
|
||||
public String getModel() { return model; }
|
||||
public void setModel(String model) { this.model = model; }
|
||||
public Integer getInputTokens() { return inputTokens; }
|
||||
public void setInputTokens(Integer inputTokens) { this.inputTokens = inputTokens; }
|
||||
public Integer getOutputTokens() { return outputTokens; }
|
||||
public void setOutputTokens(Integer outputTokens) { this.outputTokens = outputTokens; }
|
||||
public Integer getTotalTokens() { return totalTokens; }
|
||||
public void setTotalTokens(Integer totalTokens) { this.totalTokens = totalTokens; }
|
||||
public Integer getCostCents() { return costCents; }
|
||||
public void setCostCents(Integer costCents) { this.costCents = costCents; }
|
||||
public String getSuccess() { return success; }
|
||||
public void setSuccess(String success) { this.success = success; }
|
||||
public String getErrorMessage() { return errorMessage; }
|
||||
public void setErrorMessage(String errorMessage) { this.errorMessage = errorMessage; }
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package com.ruoyi.generator.domain.front.dto;
|
||||
|
||||
public class AiGenerationTaskCreateRequest
|
||||
{
|
||||
private String generateType;
|
||||
private String projectName;
|
||||
private String projectDesc;
|
||||
private String industryTemplate;
|
||||
private String extraRequirements;
|
||||
private String appBlueprint;
|
||||
|
||||
public String getGenerateType() { return generateType; }
|
||||
public void setGenerateType(String generateType) { this.generateType = generateType; }
|
||||
public String getProjectName() { return projectName; }
|
||||
public void setProjectName(String projectName) { this.projectName = projectName; }
|
||||
public String getProjectDesc() { return projectDesc; }
|
||||
public void setProjectDesc(String projectDesc) { this.projectDesc = projectDesc; }
|
||||
public String getIndustryTemplate() { return industryTemplate; }
|
||||
public void setIndustryTemplate(String industryTemplate) { this.industryTemplate = industryTemplate; }
|
||||
public String getExtraRequirements() { return extraRequirements; }
|
||||
public void setExtraRequirements(String extraRequirements) { this.extraRequirements = extraRequirements; }
|
||||
public String getAppBlueprint() { return appBlueprint; }
|
||||
public void setAppBlueprint(String appBlueprint) { this.appBlueprint = appBlueprint; }
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
package com.ruoyi.generator.domain.front.dto;
|
||||
|
||||
import java.util.Date;
|
||||
|
||||
public class AiGenerationTaskStatusResponse
|
||||
{
|
||||
private Long taskId;
|
||||
private Long projectId;
|
||||
private Long generationId;
|
||||
private String generateType;
|
||||
private String status;
|
||||
private Integer progress;
|
||||
private String currentStep;
|
||||
private String resultPayload;
|
||||
private String errorMessage;
|
||||
private Integer attempts;
|
||||
private Integer maxAttempts;
|
||||
private Date nextRetryTime;
|
||||
private Date createTime;
|
||||
private Date updateTime;
|
||||
|
||||
public Long getTaskId() { return taskId; }
|
||||
public void setTaskId(Long taskId) { this.taskId = taskId; }
|
||||
public Long getProjectId() { return projectId; }
|
||||
public void setProjectId(Long projectId) { this.projectId = projectId; }
|
||||
public Long getGenerationId() { return generationId; }
|
||||
public void setGenerationId(Long generationId) { this.generationId = generationId; }
|
||||
public String getGenerateType() { return generateType; }
|
||||
public void setGenerateType(String generateType) { this.generateType = generateType; }
|
||||
public String getStatus() { return status; }
|
||||
public void setStatus(String status) { this.status = status; }
|
||||
public Integer getProgress() { return progress; }
|
||||
public void setProgress(Integer progress) { this.progress = progress; }
|
||||
public String getCurrentStep() { return currentStep; }
|
||||
public void setCurrentStep(String currentStep) { this.currentStep = currentStep; }
|
||||
public String getResultPayload() { return resultPayload; }
|
||||
public void setResultPayload(String resultPayload) { this.resultPayload = resultPayload; }
|
||||
public String getErrorMessage() { return errorMessage; }
|
||||
public void setErrorMessage(String errorMessage) { this.errorMessage = errorMessage; }
|
||||
public Integer getAttempts() { return attempts; }
|
||||
public void setAttempts(Integer attempts) { this.attempts = attempts; }
|
||||
public Integer getMaxAttempts() { return maxAttempts; }
|
||||
public void setMaxAttempts(Integer maxAttempts) { this.maxAttempts = maxAttempts; }
|
||||
public Date getNextRetryTime() { return nextRetryTime; }
|
||||
public void setNextRetryTime(Date nextRetryTime) { this.nextRetryTime = nextRetryTime; }
|
||||
public Date getCreateTime() { return createTime; }
|
||||
public void setCreateTime(Date createTime) { this.createTime = createTime; }
|
||||
public Date getUpdateTime() { return updateTime; }
|
||||
public void setUpdateTime(Date updateTime) { this.updateTime = updateTime; }
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package com.ruoyi.generator.domain.front.dto;
|
||||
|
||||
public class AiQuotaResponse
|
||||
{
|
||||
private Integer dailyTaskLimit;
|
||||
private Integer dailyTaskUsed;
|
||||
private Integer monthlyCostLimitCents;
|
||||
private Integer monthlyCostUsedCents;
|
||||
private Integer runningLimit;
|
||||
private Integer runningCount;
|
||||
|
||||
public Integer getDailyTaskLimit() { return dailyTaskLimit; }
|
||||
public void setDailyTaskLimit(Integer dailyTaskLimit) { this.dailyTaskLimit = dailyTaskLimit; }
|
||||
public Integer getDailyTaskUsed() { return dailyTaskUsed; }
|
||||
public void setDailyTaskUsed(Integer dailyTaskUsed) { this.dailyTaskUsed = dailyTaskUsed; }
|
||||
public Integer getMonthlyCostLimitCents() { return monthlyCostLimitCents; }
|
||||
public void setMonthlyCostLimitCents(Integer monthlyCostLimitCents) { this.monthlyCostLimitCents = monthlyCostLimitCents; }
|
||||
public Integer getMonthlyCostUsedCents() { return monthlyCostUsedCents; }
|
||||
public void setMonthlyCostUsedCents(Integer monthlyCostUsedCents) { this.monthlyCostUsedCents = monthlyCostUsedCents; }
|
||||
public Integer getRunningLimit() { return runningLimit; }
|
||||
public void setRunningLimit(Integer runningLimit) { this.runningLimit = runningLimit; }
|
||||
public Integer getRunningCount() { return runningCount; }
|
||||
public void setRunningCount(Integer runningCount) { this.runningCount = runningCount; }
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
package com.ruoyi.generator.mapper.front;
|
||||
|
||||
import java.util.List;
|
||||
import org.apache.ibatis.annotations.Param;
|
||||
import com.ruoyi.generator.domain.front.AiGenerationTask;
|
||||
|
||||
public interface AiGenerationTaskMapper
|
||||
{
|
||||
public AiGenerationTask selectAiGenerationTaskById(Long taskId);
|
||||
public AiGenerationTask selectTaskForUser(@Param("userId") Long userId, @Param("projectId") Long projectId, @Param("taskId") Long taskId);
|
||||
public AiGenerationTask selectActiveTaskByRequestHash(@Param("userId") Long userId, @Param("projectId") Long projectId, @Param("requestHash") String requestHash);
|
||||
public List<AiGenerationTask> selectTasksByProjectId(@Param("userId") Long userId, @Param("projectId") Long projectId);
|
||||
public List<AiGenerationTask> selectClaimableTasks(@Param("limit") Integer limit);
|
||||
public int insertAiGenerationTask(AiGenerationTask aiGenerationTask);
|
||||
public int updateAiGenerationTask(AiGenerationTask aiGenerationTask);
|
||||
public int claimTask(@Param("taskId") Long taskId, @Param("lockedBy") String lockedBy);
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
package com.ruoyi.generator.mapper.front;
|
||||
|
||||
import org.apache.ibatis.annotations.Param;
|
||||
import com.ruoyi.generator.domain.front.AiQuotaBucket;
|
||||
|
||||
public interface AiQuotaBucketMapper
|
||||
{
|
||||
public AiQuotaBucket selectQuotaBucket(@Param("userId") Long userId, @Param("periodType") String periodType, @Param("periodKey") String periodKey);
|
||||
public int insertQuotaBucket(AiQuotaBucket aiQuotaBucket);
|
||||
public int updateQuotaBucket(AiQuotaBucket aiQuotaBucket);
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
package com.ruoyi.generator.mapper.front;
|
||||
|
||||
import com.ruoyi.generator.domain.front.AiUsageLedger;
|
||||
|
||||
public interface AiUsageLedgerMapper
|
||||
{
|
||||
public int insertAiUsageLedger(AiUsageLedger aiUsageLedger);
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
package com.ruoyi.generator.service.front;
|
||||
|
||||
import org.springframework.stereotype.Component;
|
||||
|
||||
@Component
|
||||
public class AiCostCalculator
|
||||
{
|
||||
private static final double DEEPSEEK_CHAT_INPUT_CENTS_PER_1K = 0.1D;
|
||||
private static final double DEEPSEEK_CHAT_OUTPUT_CENTS_PER_1K = 0.6D;
|
||||
|
||||
public int estimateCostCents(String provider, String model, Integer inputTokens, Integer outputTokens)
|
||||
{
|
||||
if (inputTokens == null && outputTokens == null)
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
int input = inputTokens == null ? 0 : Math.max(inputTokens, 0);
|
||||
int output = outputTokens == null ? 0 : Math.max(outputTokens, 0);
|
||||
if (input == 0 && output == 0)
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
double cents = input / 1000D * DEEPSEEK_CHAT_INPUT_CENTS_PER_1K
|
||||
+ output / 1000D * DEEPSEEK_CHAT_OUTPUT_CENTS_PER_1K;
|
||||
return Math.max(1, (int) Math.ceil(cents));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,222 @@
|
||||
package com.ruoyi.generator.service.front;
|
||||
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.security.MessageDigest;
|
||||
import java.security.NoSuchAlgorithmException;
|
||||
import java.util.Date;
|
||||
import java.util.List;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
|
||||
import org.springframework.stereotype.Service;
|
||||
import org.springframework.transaction.annotation.Transactional;
|
||||
import com.alibaba.fastjson2.JSON;
|
||||
import com.ruoyi.common.exception.ServiceException;
|
||||
import com.ruoyi.common.utils.StringUtils;
|
||||
import com.ruoyi.generator.domain.front.AiGenerationTask;
|
||||
import com.ruoyi.generator.domain.front.FrontProject;
|
||||
import com.ruoyi.generator.domain.front.FrontProjectGeneration;
|
||||
import com.ruoyi.generator.domain.front.dto.AiGenerationTaskCreateRequest;
|
||||
import com.ruoyi.generator.domain.front.dto.AiGenerationTaskStatusResponse;
|
||||
import com.ruoyi.generator.domain.front.dto.AiQuotaResponse;
|
||||
import com.ruoyi.generator.mapper.front.AiGenerationTaskMapper;
|
||||
import com.ruoyi.generator.mapper.front.FrontProjectGenerationMapper;
|
||||
import com.ruoyi.generator.mapper.front.FrontProjectMapper;
|
||||
|
||||
@Service
|
||||
public class AiGenerationTaskServiceImpl implements IAiGenerationTaskService
|
||||
{
|
||||
private static final int MAX_ATTEMPTS = 3;
|
||||
|
||||
@Autowired
|
||||
private FrontProjectMapper frontProjectMapper;
|
||||
@Autowired
|
||||
private AiGenerationTaskMapper aiGenerationTaskMapper;
|
||||
@Autowired(required = false)
|
||||
private FrontProjectGenerationMapper frontProjectGenerationMapper;
|
||||
@Autowired
|
||||
private AiQuotaService aiQuotaService;
|
||||
@Autowired(required = false)
|
||||
private AiGenerationTaskWorker aiGenerationTaskWorker;
|
||||
@Autowired(required = false)
|
||||
private ThreadPoolTaskExecutor threadPoolTaskExecutor;
|
||||
|
||||
@Override
|
||||
@Transactional
|
||||
public AiGenerationTaskStatusResponse createTask(Long userId, Long projectId, AiGenerationTaskCreateRequest request)
|
||||
{
|
||||
assertOwnedProject(userId, projectId);
|
||||
validateRequest(request);
|
||||
String requestPayload = JSON.toJSONString(request);
|
||||
String requestHash = sha256(request.getGenerateType() + ":" + requestPayload);
|
||||
AiGenerationTask existing = aiGenerationTaskMapper.selectActiveTaskByRequestHash(userId, projectId, requestHash);
|
||||
if (existing != null)
|
||||
{
|
||||
return toResponse(existing);
|
||||
}
|
||||
aiQuotaService.reserve(userId);
|
||||
AiGenerationTask task = new AiGenerationTask();
|
||||
task.setProjectId(projectId);
|
||||
task.setUserId(userId);
|
||||
task.setGenerateType(request.getGenerateType());
|
||||
task.setProvider("deepseek");
|
||||
task.setModel("deepseek-chat");
|
||||
task.setStatus("QUEUED");
|
||||
task.setRequestHash(requestHash);
|
||||
task.setRequestPayload(requestPayload);
|
||||
task.setAttempts(0);
|
||||
task.setMaxAttempts(MAX_ATTEMPTS);
|
||||
task.setProgress(0);
|
||||
task.setCurrentStep("等待生成");
|
||||
task.setNextRetryTime(new Date());
|
||||
aiGenerationTaskMapper.insertAiGenerationTask(task);
|
||||
dispatch(task.getTaskId());
|
||||
return toResponse(task);
|
||||
}
|
||||
|
||||
@Override
|
||||
public AiGenerationTaskStatusResponse getTask(Long userId, Long projectId, Long taskId)
|
||||
{
|
||||
AiGenerationTask task = selectOwnedTask(userId, projectId, taskId);
|
||||
return toResponse(task);
|
||||
}
|
||||
|
||||
@Override
|
||||
@Transactional
|
||||
public AiGenerationTaskStatusResponse retryTask(Long userId, Long projectId, Long taskId)
|
||||
{
|
||||
AiGenerationTask task = selectOwnedTask(userId, projectId, taskId);
|
||||
if (!"FAILED".equals(task.getStatus()) && !"RETRY_WAITING".equals(task.getStatus()))
|
||||
{
|
||||
throw new ServiceException("当前任务状态不允许重试");
|
||||
}
|
||||
aiQuotaService.reserve(userId);
|
||||
task.setStatus("QUEUED");
|
||||
task.setProgress(0);
|
||||
task.setCurrentStep("等待重试");
|
||||
task.setErrorCode("");
|
||||
task.setErrorMessage("");
|
||||
task.setNextRetryTime(new Date());
|
||||
aiGenerationTaskMapper.updateAiGenerationTask(task);
|
||||
dispatch(taskId);
|
||||
return toResponse(task);
|
||||
}
|
||||
|
||||
@Override
|
||||
@Transactional
|
||||
public AiGenerationTaskStatusResponse cancelTask(Long userId, Long projectId, Long taskId)
|
||||
{
|
||||
AiGenerationTask task = selectOwnedTask(userId, projectId, taskId);
|
||||
if (!"QUEUED".equals(task.getStatus()) && !"RETRY_WAITING".equals(task.getStatus()))
|
||||
{
|
||||
throw new ServiceException("当前任务状态不允许取消");
|
||||
}
|
||||
task.setStatus("CANCELED");
|
||||
task.setProgress(0);
|
||||
task.setCurrentStep("已取消");
|
||||
task.setFinishedAt(new Date());
|
||||
aiGenerationTaskMapper.updateAiGenerationTask(task);
|
||||
aiQuotaService.releaseRunning(userId);
|
||||
return toResponse(task);
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<FrontProjectGeneration> listGenerations(Long userId, Long projectId)
|
||||
{
|
||||
assertOwnedProject(userId, projectId);
|
||||
return frontProjectGenerationMapper.selectGenerationsByProjectId(projectId);
|
||||
}
|
||||
|
||||
@Override
|
||||
public AiQuotaResponse getQuota(Long userId)
|
||||
{
|
||||
return aiQuotaService.getQuota(userId);
|
||||
}
|
||||
|
||||
private FrontProject assertOwnedProject(Long userId, Long projectId)
|
||||
{
|
||||
FrontProject project = frontProjectMapper.selectFrontProjectByUserAndId(userId, projectId);
|
||||
if (project == null)
|
||||
{
|
||||
throw new ServiceException("项目不存在或无权限");
|
||||
}
|
||||
return project;
|
||||
}
|
||||
|
||||
private AiGenerationTask selectOwnedTask(Long userId, Long projectId, Long taskId)
|
||||
{
|
||||
AiGenerationTask task = aiGenerationTaskMapper.selectTaskForUser(userId, projectId, taskId);
|
||||
if (task == null)
|
||||
{
|
||||
throw new ServiceException("任务不存在或无权限");
|
||||
}
|
||||
return task;
|
||||
}
|
||||
|
||||
private void validateRequest(AiGenerationTaskCreateRequest request)
|
||||
{
|
||||
if (request == null || StringUtils.isEmpty(request.getGenerateType()))
|
||||
{
|
||||
throw new ServiceException("生成类型不能为空");
|
||||
}
|
||||
if (!"app_blueprint".equals(request.getGenerateType()) && !"database".equals(request.getGenerateType())
|
||||
&& !"business_blueprint".equals(request.getGenerateType()))
|
||||
{
|
||||
throw new ServiceException("不支持的生成类型");
|
||||
}
|
||||
}
|
||||
|
||||
private AiGenerationTaskStatusResponse toResponse(AiGenerationTask task)
|
||||
{
|
||||
AiGenerationTaskStatusResponse response = new AiGenerationTaskStatusResponse();
|
||||
response.setTaskId(task.getTaskId());
|
||||
response.setProjectId(task.getProjectId());
|
||||
response.setGenerationId(task.getGenerationId());
|
||||
response.setGenerateType(task.getGenerateType());
|
||||
response.setStatus(task.getStatus());
|
||||
response.setProgress(task.getProgress());
|
||||
response.setCurrentStep(task.getCurrentStep());
|
||||
response.setResultPayload(task.getResultPayload());
|
||||
response.setErrorMessage(task.getErrorMessage());
|
||||
response.setAttempts(task.getAttempts());
|
||||
response.setMaxAttempts(task.getMaxAttempts());
|
||||
response.setNextRetryTime(task.getNextRetryTime());
|
||||
response.setCreateTime(task.getCreateTime());
|
||||
response.setUpdateTime(task.getUpdateTime());
|
||||
return response;
|
||||
}
|
||||
|
||||
private String sha256(String value)
|
||||
{
|
||||
try
|
||||
{
|
||||
MessageDigest digest = MessageDigest.getInstance("SHA-256");
|
||||
byte[] hash = digest.digest(value.getBytes(StandardCharsets.UTF_8));
|
||||
StringBuilder builder = new StringBuilder();
|
||||
for (byte b : hash)
|
||||
{
|
||||
builder.append(String.format("%02x", b));
|
||||
}
|
||||
return builder.toString();
|
||||
}
|
||||
catch (NoSuchAlgorithmException e)
|
||||
{
|
||||
throw new ServiceException("生成请求指纹失败").setDetailMessage(e.getMessage());
|
||||
}
|
||||
}
|
||||
|
||||
private void dispatch(final Long taskId)
|
||||
{
|
||||
if (taskId == null || aiGenerationTaskWorker == null || threadPoolTaskExecutor == null)
|
||||
{
|
||||
return;
|
||||
}
|
||||
threadPoolTaskExecutor.execute(new Runnable()
|
||||
{
|
||||
@Override
|
||||
public void run()
|
||||
{
|
||||
aiGenerationTaskWorker.processTask(taskId);
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,243 @@
|
||||
package com.ruoyi.generator.service.front;
|
||||
|
||||
import java.util.Date;
|
||||
import java.util.List;
|
||||
import java.util.UUID;
|
||||
import java.util.concurrent.ScheduledExecutorService;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import javax.annotation.PostConstruct;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.stereotype.Service;
|
||||
import com.alibaba.fastjson2.JSON;
|
||||
import com.ruoyi.common.utils.StringUtils;
|
||||
import com.ruoyi.generator.domain.front.AiGenerationTask;
|
||||
import com.ruoyi.generator.domain.front.AiUsageLedger;
|
||||
import com.ruoyi.generator.domain.front.dto.AiGenerationTaskCreateRequest;
|
||||
import com.ruoyi.generator.domain.front.dto.GenerateAppBlueprintRequest;
|
||||
import com.ruoyi.generator.domain.front.dto.GenerateBusinessBlueprintRequest;
|
||||
import com.ruoyi.generator.domain.front.dto.GenerateDatabaseRequest;
|
||||
import com.ruoyi.generator.mapper.front.AiGenerationTaskMapper;
|
||||
import com.ruoyi.generator.mapper.front.AiUsageLedgerMapper;
|
||||
|
||||
@Service
|
||||
public class AiGenerationTaskWorker
|
||||
{
|
||||
@Autowired
|
||||
private AiGenerationTaskMapper aiGenerationTaskMapper;
|
||||
@Autowired(required = false)
|
||||
private AiUsageLedgerMapper aiUsageLedgerMapper;
|
||||
@Autowired
|
||||
private IAiGenerateService aiGenerateService;
|
||||
@Autowired
|
||||
private AiQuotaService aiQuotaService;
|
||||
@Autowired
|
||||
private AiTaskRetryPolicy retryPolicy;
|
||||
@Autowired
|
||||
private AiCostCalculator costCalculator;
|
||||
@Autowired(required = false)
|
||||
private ScheduledExecutorService scheduledExecutorService;
|
||||
|
||||
@PostConstruct
|
||||
public void startPolling()
|
||||
{
|
||||
if (scheduledExecutorService == null)
|
||||
{
|
||||
return;
|
||||
}
|
||||
scheduledExecutorService.scheduleWithFixedDelay(new Runnable()
|
||||
{
|
||||
@Override
|
||||
public void run()
|
||||
{
|
||||
processDueTasks(5);
|
||||
}
|
||||
}, 10, 10, TimeUnit.SECONDS);
|
||||
}
|
||||
|
||||
public void processTask(Long taskId)
|
||||
{
|
||||
AiGenerationTask task = aiGenerationTaskMapper.selectAiGenerationTaskById(taskId);
|
||||
if (task == null)
|
||||
{
|
||||
return;
|
||||
}
|
||||
if (aiGenerationTaskMapper.claimTask(taskId, workerId()) == 0)
|
||||
{
|
||||
return;
|
||||
}
|
||||
int attempts = safe(task.getAttempts()) + 1;
|
||||
task.setAttempts(attempts);
|
||||
task.setStatus("RUNNING");
|
||||
task.setProgress(10);
|
||||
task.setCurrentStep("正在调用 AI");
|
||||
task.setStartedAt(new Date());
|
||||
aiGenerationTaskMapper.updateAiGenerationTask(task);
|
||||
|
||||
try
|
||||
{
|
||||
Object result = execute(task);
|
||||
task.setStatus("SUCCEEDED");
|
||||
task.setProgress(100);
|
||||
task.setCurrentStep("生成完成");
|
||||
task.setResultPayload(JSON.toJSONString(result));
|
||||
task.setErrorCode("");
|
||||
task.setErrorMessage("");
|
||||
task.setFinishedAt(new Date());
|
||||
task.setInputTokens(estimateTokens(task.getRequestPayload()));
|
||||
task.setOutputTokens(estimateTokens(task.getResultPayload()));
|
||||
task.setTotalTokens(safe(task.getInputTokens()) + safe(task.getOutputTokens()));
|
||||
int totalTokens = safe(task.getTotalTokens());
|
||||
int costCents = costCalculator.estimateCostCents(task.getProvider(), task.getModel(), task.getInputTokens(), task.getOutputTokens());
|
||||
task.setCostCents(costCents);
|
||||
aiGenerationTaskMapper.updateAiGenerationTask(task);
|
||||
insertLedger(task, "1");
|
||||
aiQuotaService.settleCost(task.getUserId(), totalTokens, costCents);
|
||||
}
|
||||
catch (RuntimeException e)
|
||||
{
|
||||
String message = sanitize(e.getMessage());
|
||||
task.setErrorMessage(message);
|
||||
task.setProgress(0);
|
||||
if (retryPolicy.canRetry(message, attempts, safeMax(task.getMaxAttempts())))
|
||||
{
|
||||
task.setStatus("RETRY_WAITING");
|
||||
task.setCurrentStep("等待重试");
|
||||
task.setNextRetryTime(retryPolicy.nextRetryTime(attempts, new Date()));
|
||||
}
|
||||
else
|
||||
{
|
||||
task.setStatus("FAILED");
|
||||
task.setCurrentStep("生成失败");
|
||||
task.setFinishedAt(new Date());
|
||||
}
|
||||
aiGenerationTaskMapper.updateAiGenerationTask(task);
|
||||
insertLedger(task, "0");
|
||||
}
|
||||
finally
|
||||
{
|
||||
aiQuotaService.releaseRunning(task.getUserId());
|
||||
}
|
||||
}
|
||||
|
||||
public int processDueTasks(int limit)
|
||||
{
|
||||
List<AiGenerationTask> tasks = aiGenerationTaskMapper.selectClaimableTasks(limit);
|
||||
if (tasks == null || tasks.isEmpty())
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
int count = 0;
|
||||
for (AiGenerationTask task : tasks)
|
||||
{
|
||||
processTask(task.getTaskId());
|
||||
count++;
|
||||
}
|
||||
return count;
|
||||
}
|
||||
|
||||
private Object execute(AiGenerationTask task)
|
||||
{
|
||||
AiGenerationTaskCreateRequest request = JSON.parseObject(task.getRequestPayload(), AiGenerationTaskCreateRequest.class);
|
||||
if ("app_blueprint".equals(task.getGenerateType()))
|
||||
{
|
||||
return aiGenerateService.generateAppBlueprint(task.getUserId(), task.getProjectId(), toAppBlueprintRequest(request));
|
||||
}
|
||||
if ("database".equals(task.getGenerateType()))
|
||||
{
|
||||
return aiGenerateService.generateDatabase(task.getUserId(), task.getProjectId(), toDatabaseRequest(request));
|
||||
}
|
||||
if ("business_blueprint".equals(task.getGenerateType()))
|
||||
{
|
||||
return aiGenerateService.generateBusinessBlueprint(task.getUserId(), task.getProjectId(), toBusinessBlueprintRequest(request));
|
||||
}
|
||||
throw new IllegalArgumentException("不支持的生成类型");
|
||||
}
|
||||
|
||||
private GenerateAppBlueprintRequest toAppBlueprintRequest(AiGenerationTaskCreateRequest source)
|
||||
{
|
||||
GenerateAppBlueprintRequest request = new GenerateAppBlueprintRequest();
|
||||
request.setProjectName(source.getProjectName());
|
||||
request.setProjectDesc(source.getProjectDesc());
|
||||
request.setIndustryTemplate(source.getIndustryTemplate());
|
||||
request.setExtraRequirements(source.getExtraRequirements());
|
||||
return request;
|
||||
}
|
||||
|
||||
private GenerateDatabaseRequest toDatabaseRequest(AiGenerationTaskCreateRequest source)
|
||||
{
|
||||
GenerateDatabaseRequest request = new GenerateDatabaseRequest();
|
||||
request.setProjectName(source.getProjectName());
|
||||
request.setProjectDesc(source.getProjectDesc());
|
||||
request.setIndustryTemplate(source.getIndustryTemplate());
|
||||
request.setExtraRequirements(source.getExtraRequirements());
|
||||
request.setAppBlueprint(source.getAppBlueprint());
|
||||
return request;
|
||||
}
|
||||
|
||||
private GenerateBusinessBlueprintRequest toBusinessBlueprintRequest(AiGenerationTaskCreateRequest source)
|
||||
{
|
||||
GenerateBusinessBlueprintRequest request = new GenerateBusinessBlueprintRequest();
|
||||
request.setProjectName(source.getProjectName());
|
||||
request.setProjectDesc(source.getProjectDesc());
|
||||
request.setIndustryTemplate(source.getIndustryTemplate());
|
||||
request.setExtraRequirements(source.getExtraRequirements());
|
||||
return request;
|
||||
}
|
||||
|
||||
private void insertLedger(AiGenerationTask task, String success)
|
||||
{
|
||||
if (aiUsageLedgerMapper == null)
|
||||
{
|
||||
return;
|
||||
}
|
||||
AiUsageLedger ledger = new AiUsageLedger();
|
||||
ledger.setTaskId(task.getTaskId());
|
||||
ledger.setGenerationId(task.getGenerationId());
|
||||
ledger.setProjectId(task.getProjectId());
|
||||
ledger.setUserId(task.getUserId());
|
||||
ledger.setGenerateType(task.getGenerateType());
|
||||
ledger.setProvider(task.getProvider());
|
||||
ledger.setModel(task.getModel());
|
||||
ledger.setInputTokens(task.getInputTokens());
|
||||
ledger.setOutputTokens(task.getOutputTokens());
|
||||
ledger.setTotalTokens(task.getTotalTokens());
|
||||
ledger.setCostCents(task.getCostCents());
|
||||
ledger.setSuccess(success);
|
||||
ledger.setErrorMessage(task.getErrorMessage());
|
||||
aiUsageLedgerMapper.insertAiUsageLedger(ledger);
|
||||
}
|
||||
|
||||
private String workerId()
|
||||
{
|
||||
return "ai-task-worker-" + UUID.randomUUID().toString();
|
||||
}
|
||||
|
||||
private int safe(Integer value)
|
||||
{
|
||||
return value == null ? 0 : value;
|
||||
}
|
||||
|
||||
private int estimateTokens(String text)
|
||||
{
|
||||
if (StringUtils.isEmpty(text))
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
return Math.max(1, (text.length() + 3) / 4);
|
||||
}
|
||||
|
||||
private int safeMax(Integer value)
|
||||
{
|
||||
return value == null || value <= 0 ? 3 : value;
|
||||
}
|
||||
|
||||
private String sanitize(String message)
|
||||
{
|
||||
if (StringUtils.isEmpty(message))
|
||||
{
|
||||
return "";
|
||||
}
|
||||
String sanitized = message.replaceAll("[\\p{Cntrl}]", " ");
|
||||
return sanitized.length() > 1000 ? sanitized.substring(0, 1000) : sanitized;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
package com.ruoyi.generator.service.front;
|
||||
|
||||
import java.text.SimpleDateFormat;
|
||||
import java.util.Date;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.stereotype.Service;
|
||||
import org.springframework.transaction.annotation.Transactional;
|
||||
import com.ruoyi.common.exception.ServiceException;
|
||||
import com.ruoyi.generator.domain.front.AiQuotaBucket;
|
||||
import com.ruoyi.generator.domain.front.dto.AiQuotaResponse;
|
||||
import com.ruoyi.generator.mapper.front.AiQuotaBucketMapper;
|
||||
|
||||
@Service
|
||||
public class AiQuotaService
|
||||
{
|
||||
private static final int DAILY_TASK_LIMIT = 20;
|
||||
private static final int MONTHLY_COST_LIMIT_CENTS = 5000;
|
||||
private static final int RUNNING_LIMIT = 1;
|
||||
|
||||
@Autowired(required = false)
|
||||
private AiQuotaBucketMapper aiQuotaBucketMapper;
|
||||
|
||||
@Transactional
|
||||
public void reserve(Long userId)
|
||||
{
|
||||
if (aiQuotaBucketMapper == null)
|
||||
{
|
||||
return;
|
||||
}
|
||||
AiQuotaBucket day = ensureBucket(userId, "DAY", new SimpleDateFormat("yyyyMMdd").format(new Date()));
|
||||
AiQuotaBucket month = ensureBucket(userId, "MONTH", new SimpleDateFormat("yyyyMM").format(new Date()));
|
||||
if (safe(day.getTaskUsed()) >= safe(day.getTaskLimit()))
|
||||
{
|
||||
throw new ServiceException("今日 AI 生成次数已用完");
|
||||
}
|
||||
if (safe(month.getCostUsedCents()) >= safe(month.getCostLimitCents()))
|
||||
{
|
||||
throw new ServiceException("本月 AI 生成成本配额已用完");
|
||||
}
|
||||
if (safe(day.getRunningCount()) >= safe(day.getRunningLimit()))
|
||||
{
|
||||
throw new ServiceException("已有 AI 生成任务正在执行，请稍后再试");
|
||||
}
|
||||
day.setTaskUsed(safe(day.getTaskUsed()) + 1);
|
||||
day.setRunningCount(safe(day.getRunningCount()) + 1);
|
||||
aiQuotaBucketMapper.updateQuotaBucket(day);
|
||||
}
|
||||
|
||||
@Transactional
|
||||
public void releaseRunning(Long userId)
|
||||
{
|
||||
if (aiQuotaBucketMapper == null)
|
||||
{
|
||||
return;
|
||||
}
|
||||
AiQuotaBucket day = ensureBucket(userId, "DAY", new SimpleDateFormat("yyyyMMdd").format(new Date()));
|
||||
day.setRunningCount(Math.max(0, safe(day.getRunningCount()) - 1));
|
||||
aiQuotaBucketMapper.updateQuotaBucket(day);
|
||||
}
|
||||
|
||||
@Transactional
|
||||
public void settleCost(Long userId, int totalTokens, int costCents)
|
||||
{
|
||||
if (aiQuotaBucketMapper == null)
|
||||
{
|
||||
return;
|
||||
}
|
||||
AiQuotaBucket month = ensureBucket(userId, "MONTH", new SimpleDateFormat("yyyyMM").format(new Date()));
|
||||
month.setTokenUsed(safe(month.getTokenUsed()) + Math.max(totalTokens, 0));
|
||||
month.setCostUsedCents(safe(month.getCostUsedCents()) + Math.max(costCents, 0));
|
||||
aiQuotaBucketMapper.updateQuotaBucket(month);
|
||||
}
|
||||
|
||||
public AiQuotaResponse getQuota(Long userId)
|
||||
{
|
||||
AiQuotaResponse response = new AiQuotaResponse();
|
||||
response.setDailyTaskLimit(DAILY_TASK_LIMIT);
|
||||
response.setDailyTaskUsed(0);
|
||||
response.setMonthlyCostLimitCents(MONTHLY_COST_LIMIT_CENTS);
|
||||
response.setMonthlyCostUsedCents(0);
|
||||
response.setRunningLimit(RUNNING_LIMIT);
|
||||
response.setRunningCount(0);
|
||||
if (aiQuotaBucketMapper != null)
|
||||
{
|
||||
AiQuotaBucket day = aiQuotaBucketMapper.selectQuotaBucket(userId, "DAY", new SimpleDateFormat("yyyyMMdd").format(new Date()));
|
||||
AiQuotaBucket month = aiQuotaBucketMapper.selectQuotaBucket(userId, "MONTH", new SimpleDateFormat("yyyyMM").format(new Date()));
|
||||
if (day != null)
|
||||
{
|
||||
response.setDailyTaskUsed(safe(day.getTaskUsed()));
|
||||
response.setRunningCount(safe(day.getRunningCount()));
|
||||
}
|
||||
if (month != null)
|
||||
{
|
||||
response.setMonthlyCostUsedCents(safe(month.getCostUsedCents()));
|
||||
}
|
||||
}
|
||||
return response;
|
||||
}
|
||||
|
||||
private AiQuotaBucket ensureBucket(Long userId, String periodType, String periodKey)
|
||||
{
|
||||
AiQuotaBucket bucket = aiQuotaBucketMapper.selectQuotaBucket(userId, periodType, periodKey);
|
||||
if (bucket != null)
|
||||
{
|
||||
return bucket;
|
||||
}
|
||||
bucket = new AiQuotaBucket();
|
||||
bucket.setUserId(userId);
|
||||
bucket.setPeriodType(periodType);
|
||||
bucket.setPeriodKey(periodKey);
|
||||
bucket.setTaskLimit("DAY".equals(periodType) ? DAILY_TASK_LIMIT : 0);
|
||||
bucket.setTaskUsed(0);
|
||||
bucket.setTokenLimit(0);
|
||||
bucket.setTokenUsed(0);
|
||||
bucket.setCostLimitCents("MONTH".equals(periodType) ? MONTHLY_COST_LIMIT_CENTS : 0);
|
||||
bucket.setCostUsedCents(0);
|
||||
bucket.setRunningLimit("DAY".equals(periodType) ? RUNNING_LIMIT : 0);
|
||||
bucket.setRunningCount(0);
|
||||
aiQuotaBucketMapper.insertQuotaBucket(bucket);
|
||||
return bucket;
|
||||
}
|
||||
|
||||
private int safe(Integer value)
|
||||
{
|
||||
return value == null ? 0 : value;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
package com.ruoyi.generator.service.front;
|
||||
|
||||
import java.util.Date;
|
||||
import java.util.Locale;
|
||||
import org.springframework.stereotype.Component;
|
||||
import com.ruoyi.common.utils.StringUtils;
|
||||
|
||||
@Component
|
||||
public class AiTaskRetryPolicy
|
||||
{
|
||||
private static final long[] RETRY_DELAYS_MS = new long[] { 30_000L, 120_000L, 300_000L };
|
||||
|
||||
public boolean canRetry(String errorMessage, int attempts, int maxAttempts)
|
||||
{
|
||||
if (attempts >= maxAttempts)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
if (StringUtils.isEmpty(errorMessage))
|
||||
{
|
||||
return true;
|
||||
}
|
||||
String message = errorMessage.toLowerCase(Locale.ROOT);
|
||||
if (message.contains("配额") || message.contains("quota") || message.contains("无权限")
|
||||
|| message.contains("permission") || message.contains("missing field")
|
||||
|| message.contains("validation"))
|
||||
{
|
||||
return false;
|
||||
}
|
||||
return message.contains("timeout") || message.contains("timed out") || message.contains("429")
|
||||
|| message.contains("too many requests") || message.contains("unavailable")
|
||||
|| message.contains("502") || message.contains("503") || message.contains("504")
|
||||
|| message.contains("5xx");
|
||||
}
|
||||
|
||||
public Date nextRetryTime(int attempts, Date now)
|
||||
{
|
||||
long delay = RETRY_DELAYS_MS[Math.min(Math.max(attempts, 1), RETRY_DELAYS_MS.length) - 1];
|
||||
return new Date(now.getTime() + delay);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
package com.ruoyi.generator.service.front;
|
||||
|
||||
import java.util.List;
|
||||
import com.ruoyi.generator.domain.front.FrontProjectGeneration;
|
||||
import com.ruoyi.generator.domain.front.dto.AiGenerationTaskCreateRequest;
|
||||
import com.ruoyi.generator.domain.front.dto.AiGenerationTaskStatusResponse;
|
||||
import com.ruoyi.generator.domain.front.dto.AiQuotaResponse;
|
||||
|
||||
public interface IAiGenerationTaskService
|
||||
{
|
||||
AiGenerationTaskStatusResponse createTask(Long userId, Long projectId, AiGenerationTaskCreateRequest request);
|
||||
AiGenerationTaskStatusResponse getTask(Long userId, Long projectId, Long taskId);
|
||||
AiGenerationTaskStatusResponse retryTask(Long userId, Long projectId, Long taskId);
|
||||
AiGenerationTaskStatusResponse cancelTask(Long userId, Long projectId, Long taskId);
|
||||
List<FrontProjectGeneration> listGenerations(Long userId, Long projectId);
|
||||
AiQuotaResponse getQuota(Long userId);
|
||||
}
|
||||
@@ -0,0 +1,136 @@
|
||||
<?xml version="1.0" encoding="UTF-8" ?>
|
||||
<!DOCTYPE mapper
|
||||
PUBLIC "-//mybatis.org//DTD Mapper 3.0//EN"
|
||||
"http://mybatis.org/dtd/mybatis-3-mapper.dtd">
|
||||
<mapper namespace="com.ruoyi.generator.mapper.front.AiGenerationTaskMapper">
|
||||
|
||||
<resultMap type="AiGenerationTask" id="AiGenerationTaskResult">
|
||||
<id property="taskId" column="task_id" />
|
||||
<result property="generationId" column="generation_id" />
|
||||
<result property="projectId" column="project_id" />
|
||||
<result property="userId" column="user_id" />
|
||||
<result property="generateType" column="generate_type" />
|
||||
<result property="provider" column="provider" />
|
||||
<result property="model" column="model" />
|
||||
<result property="status" column="status" />
|
||||
<result property="requestHash" column="request_hash" />
|
||||
<result property="requestPayload" column="request_payload" />
|
||||
<result property="resultPayload" column="result_payload" />
|
||||
<result property="errorCode" column="error_code" />
|
||||
<result property="errorMessage" column="error_message" />
|
||||
<result property="attempts" column="attempts" />
|
||||
<result property="maxAttempts" column="max_attempts" />
|
||||
<result property="nextRetryTime" column="next_retry_time" />
|
||||
<result property="lockedBy" column="locked_by" />
|
||||
<result property="lockedUntil" column="locked_until" />
|
||||
<result property="progress" column="progress" />
|
||||
<result property="currentStep" column="current_step" />
|
||||
<result property="inputTokens" column="input_tokens" />
|
||||
<result property="outputTokens" column="output_tokens" />
|
||||
<result property="totalTokens" column="total_tokens" />
|
||||
<result property="costCents" column="cost_cents" />
|
||||
<result property="startedAt" column="started_at" />
|
||||
<result property="finishedAt" column="finished_at" />
|
||||
<result property="createTime" column="create_time" />
|
||||
<result property="updateTime" column="update_time" />
|
||||
</resultMap>
|
||||
|
||||
<sql id="selectAiGenerationTaskVo">
|
||||
select task_id, generation_id, project_id, user_id, generate_type, provider, model, status,
|
||||
request_hash, request_payload, result_payload, error_code, error_message, attempts, max_attempts,
|
||||
next_retry_time, locked_by, locked_until, progress, current_step, input_tokens, output_tokens,
|
||||
total_tokens, cost_cents, started_at, finished_at, create_time, update_time
|
||||
from front_ai_generation_task
|
||||
</sql>
|
||||
|
||||
<select id="selectAiGenerationTaskById" parameterType="Long" resultMap="AiGenerationTaskResult">
|
||||
<include refid="selectAiGenerationTaskVo"/>
|
||||
where task_id = #{taskId}
|
||||
</select>
|
||||
|
||||
<select id="selectTaskForUser" resultMap="AiGenerationTaskResult">
|
||||
<include refid="selectAiGenerationTaskVo"/>
|
||||
where user_id = #{userId} and project_id = #{projectId} and task_id = #{taskId}
|
||||
</select>
|
||||
|
||||
<select id="selectActiveTaskByRequestHash" resultMap="AiGenerationTaskResult">
|
||||
<include refid="selectAiGenerationTaskVo"/>
|
||||
where user_id = #{userId}
|
||||
and project_id = #{projectId}
|
||||
and request_hash = #{requestHash}
|
||||
and status in ('QUEUED', 'RUNNING', 'RETRY_WAITING')
|
||||
order by task_id desc
|
||||
limit 1
|
||||
</select>
|
||||
|
||||
<select id="selectTasksByProjectId" resultMap="AiGenerationTaskResult">
|
||||
<include refid="selectAiGenerationTaskVo"/>
|
||||
where user_id = #{userId} and project_id = #{projectId}
|
||||
order by task_id desc
|
||||
</select>
|
||||
|
||||
<select id="selectClaimableTasks" resultMap="AiGenerationTaskResult">
|
||||
<include refid="selectAiGenerationTaskVo"/>
|
||||
where status in ('QUEUED', 'RETRY_WAITING')
|
||||
and (next_retry_time is null or next_retry_time <= sysdate())
|
||||
and (locked_until is null or locked_until < sysdate())
|
||||
order by task_id asc
|
||||
limit #{limit}
|
||||
</select>
|
||||
|
||||
<insert id="insertAiGenerationTask" parameterType="AiGenerationTask" useGeneratedKeys="true" keyProperty="taskId">
|
||||
insert into front_ai_generation_task
|
||||
<trim prefix="(" suffix=")" suffixOverrides=",">
|
||||
generation_id, project_id, user_id, generate_type, provider, model, status, request_hash,
|
||||
request_payload, result_payload, error_code, error_message, attempts, max_attempts,
|
||||
next_retry_time, locked_by, locked_until, progress, current_step, input_tokens, output_tokens,
|
||||
total_tokens, cost_cents, started_at, finished_at, create_time, update_time,
|
||||
</trim>
|
||||
<trim prefix="values (" suffix=")" suffixOverrides=",">
|
||||
#{generationId}, #{projectId}, #{userId}, #{generateType}, #{provider}, #{model}, #{status}, #{requestHash},
|
||||
#{requestPayload}, #{resultPayload}, #{errorCode}, #{errorMessage}, #{attempts}, #{maxAttempts},
|
||||
#{nextRetryTime}, #{lockedBy}, #{lockedUntil}, #{progress}, #{currentStep}, #{inputTokens}, #{outputTokens},
|
||||
#{totalTokens}, #{costCents}, #{startedAt}, #{finishedAt}, sysdate(), sysdate(),
|
||||
</trim>
|
||||
</insert>
|
||||
|
||||
<update id="updateAiGenerationTask" parameterType="AiGenerationTask">
|
||||
update front_ai_generation_task
|
||||
<trim prefix="SET" suffixOverrides=",">
|
||||
<if test="generationId != null">generation_id = #{generationId},</if>
|
||||
<if test="status != null">status = #{status},</if>
|
||||
<if test="resultPayload != null">result_payload = #{resultPayload},</if>
|
||||
<if test="errorCode != null">error_code = #{errorCode},</if>
|
||||
<if test="errorMessage != null">error_message = #{errorMessage},</if>
|
||||
<if test="attempts != null">attempts = #{attempts},</if>
|
||||
<if test="maxAttempts != null">max_attempts = #{maxAttempts},</if>
|
||||
<if test="nextRetryTime != null">next_retry_time = #{nextRetryTime},</if>
|
||||
<if test="lockedBy != null">locked_by = #{lockedBy},</if>
|
||||
<if test="lockedUntil != null">locked_until = #{lockedUntil},</if>
|
||||
<if test="progress != null">progress = #{progress},</if>
|
||||
<if test="currentStep != null">current_step = #{currentStep},</if>
|
||||
<if test="inputTokens != null">input_tokens = #{inputTokens},</if>
|
||||
<if test="outputTokens != null">output_tokens = #{outputTokens},</if>
|
||||
<if test="totalTokens != null">total_tokens = #{totalTokens},</if>
|
||||
<if test="costCents != null">cost_cents = #{costCents},</if>
|
||||
<if test="startedAt != null">started_at = #{startedAt},</if>
|
||||
<if test="finishedAt != null">finished_at = #{finishedAt},</if>
|
||||
update_time = sysdate(),
|
||||
</trim>
|
||||
where task_id = #{taskId}
|
||||
</update>
|
||||
|
||||
<update id="claimTask">
|
||||
update front_ai_generation_task
|
||||
set status = 'RUNNING',
|
||||
locked_by = #{lockedBy},
|
||||
locked_until = date_add(sysdate(), interval 10 minute),
|
||||
current_step = '正在启动',
|
||||
progress = 5,
|
||||
update_time = sysdate()
|
||||
where task_id = #{taskId}
|
||||
and status in ('QUEUED', 'RETRY_WAITING')
|
||||
and (next_retry_time is null or next_retry_time <= sysdate())
|
||||
and (locked_until is null or locked_until < sysdate())
|
||||
</update>
|
||||
</mapper>
|
||||
@@ -0,0 +1,53 @@
|
||||
<?xml version="1.0" encoding="UTF-8" ?>
|
||||
<!DOCTYPE mapper
|
||||
PUBLIC "-//mybatis.org//DTD Mapper 3.0//EN"
|
||||
"http://mybatis.org/dtd/mybatis-3-mapper.dtd">
|
||||
<mapper namespace="com.ruoyi.generator.mapper.front.AiQuotaBucketMapper">
|
||||
|
||||
<resultMap type="AiQuotaBucket" id="AiQuotaBucketResult">
|
||||
<id property="bucketId" column="bucket_id" />
|
||||
<result property="userId" column="user_id" />
|
||||
<result property="periodType" column="period_type" />
|
||||
<result property="periodKey" column="period_key" />
|
||||
<result property="taskLimit" column="task_limit" />
|
||||
<result property="taskUsed" column="task_used" />
|
||||
<result property="tokenLimit" column="token_limit" />
|
||||
<result property="tokenUsed" column="token_used" />
|
||||
<result property="costLimitCents" column="cost_limit_cents" />
|
||||
<result property="costUsedCents" column="cost_used_cents" />
|
||||
<result property="runningLimit" column="running_limit" />
|
||||
<result property="runningCount" column="running_count" />
|
||||
<result property="createTime" column="create_time" />
|
||||
<result property="updateTime" column="update_time" />
|
||||
</resultMap>
|
||||
|
||||
<select id="selectQuotaBucket" resultMap="AiQuotaBucketResult">
|
||||
select bucket_id, user_id, period_type, period_key, task_limit, task_used, token_limit, token_used,
|
||||
cost_limit_cents, cost_used_cents, running_limit, running_count, create_time, update_time
|
||||
from front_ai_quota_bucket
|
||||
where user_id = #{userId} and period_type = #{periodType} and period_key = #{periodKey}
|
||||
</select>
|
||||
|
||||
<insert id="insertQuotaBucket" parameterType="AiQuotaBucket" useGeneratedKeys="true" keyProperty="bucketId">
|
||||
insert into front_ai_quota_bucket
|
||||
(user_id, period_type, period_key, task_limit, task_used, token_limit, token_used,
|
||||
cost_limit_cents, cost_used_cents, running_limit, running_count, create_time, update_time)
|
||||
values
|
||||
(#{userId}, #{periodType}, #{periodKey}, #{taskLimit}, #{taskUsed}, #{tokenLimit}, #{tokenUsed},
|
||||
#{costLimitCents}, #{costUsedCents}, #{runningLimit}, #{runningCount}, sysdate(), sysdate())
|
||||
</insert>
|
||||
|
||||
<update id="updateQuotaBucket" parameterType="AiQuotaBucket">
|
||||
update front_ai_quota_bucket
|
||||
set task_limit = #{taskLimit},
|
||||
task_used = #{taskUsed},
|
||||
token_limit = #{tokenLimit},
|
||||
token_used = #{tokenUsed},
|
||||
cost_limit_cents = #{costLimitCents},
|
||||
cost_used_cents = #{costUsedCents},
|
||||
running_limit = #{runningLimit},
|
||||
running_count = #{runningCount},
|
||||
update_time = sysdate()
|
||||
where bucket_id = #{bucketId}
|
||||
</update>
|
||||
</mapper>
|
||||
@@ -0,0 +1,15 @@
|
||||
<?xml version="1.0" encoding="UTF-8" ?>
|
||||
<!DOCTYPE mapper
|
||||
PUBLIC "-//mybatis.org//DTD Mapper 3.0//EN"
|
||||
"http://mybatis.org/dtd/mybatis-3-mapper.dtd">
|
||||
<mapper namespace="com.ruoyi.generator.mapper.front.AiUsageLedgerMapper">
|
||||
|
||||
<insert id="insertAiUsageLedger" parameterType="AiUsageLedger" useGeneratedKeys="true" keyProperty="ledgerId">
|
||||
insert into front_ai_usage_ledger
|
||||
(task_id, generation_id, project_id, user_id, generate_type, provider, model, input_tokens,
|
||||
output_tokens, total_tokens, cost_cents, success, error_message, create_time)
|
||||
values
|
||||
(#{taskId}, #{generationId}, #{projectId}, #{userId}, #{generateType}, #{provider}, #{model}, #{inputTokens},
|
||||
#{outputTokens}, #{totalTokens}, #{costCents}, #{success}, #{errorMessage}, sysdate())
|
||||
</insert>
|
||||
</mapper>
|
||||
@@ -0,0 +1,26 @@
|
||||
package com.ruoyi.generator.service.front;
|
||||
|
||||
import static org.junit.Assert.assertEquals;
|
||||
|
||||
import org.junit.Test;
|
||||
|
||||
public class AiCostCalculatorTest
|
||||
{
|
||||
@Test
|
||||
public void missingTokensCostNothing()
|
||||
{
|
||||
AiCostCalculator calculator = new AiCostCalculator();
|
||||
|
||||
assertEquals(0, calculator.estimateCostCents("deepseek", "deepseek-chat", null, null));
|
||||
assertEquals(0, calculator.estimateCostCents("deepseek", "deepseek-chat", 0, 0));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void deepSeekChatUsesInputAndOutputRates()
|
||||
{
|
||||
AiCostCalculator calculator = new AiCostCalculator();
|
||||
|
||||
assertEquals(1, calculator.estimateCostCents("deepseek", "deepseek-chat", 1_000, 1_000));
|
||||
assertEquals(3, calculator.estimateCostCents("deepseek", "deepseek-chat", 2_000, 4_000));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,152 @@
|
||||
package com.ruoyi.generator.service.front;
|
||||
|
||||
import static org.junit.Assert.assertEquals;
|
||||
import static org.junit.Assert.assertNotNull;
|
||||
import static org.junit.Assert.assertSame;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.ArgumentMatchers.anyLong;
|
||||
import static org.mockito.Mockito.never;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
import java.lang.reflect.Field;
|
||||
import org.junit.Before;
|
||||
import org.junit.Test;
|
||||
import org.mockito.ArgumentCaptor;
|
||||
import org.mockito.Mock;
|
||||
import org.mockito.MockitoAnnotations;
|
||||
import com.ruoyi.common.exception.ServiceException;
|
||||
import com.ruoyi.generator.domain.front.AiGenerationTask;
|
||||
import com.ruoyi.generator.domain.front.FrontProject;
|
||||
import com.ruoyi.generator.domain.front.dto.AiGenerationTaskCreateRequest;
|
||||
import com.ruoyi.generator.domain.front.dto.AiGenerationTaskStatusResponse;
|
||||
import com.ruoyi.generator.mapper.front.AiGenerationTaskMapper;
|
||||
import com.ruoyi.generator.mapper.front.FrontProjectMapper;
|
||||
|
||||
public class AiGenerationTaskServiceImplTest
|
||||
{
|
||||
private AiGenerationTaskServiceImpl service;
|
||||
|
||||
@Mock
|
||||
private FrontProjectMapper frontProjectMapper;
|
||||
@Mock
|
||||
private AiGenerationTaskMapper aiGenerationTaskMapper;
|
||||
@Mock
|
||||
private AiQuotaService aiQuotaService;
|
||||
|
||||
@Before
|
||||
public void setUp() throws Exception
|
||||
{
|
||||
MockitoAnnotations.initMocks(this);
|
||||
service = new AiGenerationTaskServiceImpl();
|
||||
setField("frontProjectMapper", frontProjectMapper);
|
||||
setField("aiGenerationTaskMapper", aiGenerationTaskMapper);
|
||||
setField("aiQuotaService", aiQuotaService);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void createTaskQueuesDatabaseGenerationAndReservesQuota()
|
||||
{
|
||||
when(frontProjectMapper.selectFrontProjectByUserAndId(7L, 10L)).thenReturn(project());
|
||||
org.mockito.Mockito.doAnswer(invocation -> {
|
||||
AiGenerationTask task = invocation.getArgument(0);
|
||||
task.setTaskId(88L);
|
||||
return 1;
|
||||
}).when(aiGenerationTaskMapper).insertAiGenerationTask(any(AiGenerationTask.class));
|
||||
|
||||
AiGenerationTaskStatusResponse response = service.createTask(7L, 10L, request("database"));
|
||||
|
||||
ArgumentCaptor<AiGenerationTask> taskCaptor = ArgumentCaptor.forClass(AiGenerationTask.class);
|
||||
verify(aiQuotaService).reserve(7L);
|
||||
verify(aiGenerationTaskMapper).insertAiGenerationTask(taskCaptor.capture());
|
||||
AiGenerationTask task = taskCaptor.getValue();
|
||||
assertEquals(Long.valueOf(10L), task.getProjectId());
|
||||
assertEquals(Long.valueOf(7L), task.getUserId());
|
||||
assertEquals("database", task.getGenerateType());
|
||||
assertEquals("QUEUED", task.getStatus());
|
||||
assertEquals(Integer.valueOf(0), task.getProgress());
|
||||
assertEquals(Integer.valueOf(0), task.getAttempts());
|
||||
assertEquals(Integer.valueOf(3), task.getMaxAttempts());
|
||||
assertNotNull(task.getRequestHash());
|
||||
assertEquals(Long.valueOf(88L), response.getTaskId());
|
||||
assertEquals("QUEUED", response.getStatus());
|
||||
}
|
||||
|
||||
@Test
|
||||
public void createTaskReusesRunningTaskWithSameRequestHash()
|
||||
{
|
||||
when(frontProjectMapper.selectFrontProjectByUserAndId(7L, 10L)).thenReturn(project());
|
||||
AiGenerationTask existing = new AiGenerationTask();
|
||||
existing.setTaskId(99L);
|
||||
existing.setProjectId(10L);
|
||||
existing.setUserId(7L);
|
||||
existing.setGenerateType("database");
|
||||
existing.setStatus("RUNNING");
|
||||
existing.setProgress(35);
|
||||
when(aiGenerationTaskMapper.selectActiveTaskByRequestHash(anyLong(), anyLong(), org.mockito.ArgumentMatchers.anyString()))
|
||||
.thenReturn(existing);
|
||||
|
||||
AiGenerationTaskStatusResponse response = service.createTask(7L, 10L, request("database"));
|
||||
|
||||
verify(aiQuotaService, never()).reserve(anyLong());
|
||||
verify(aiGenerationTaskMapper, never()).insertAiGenerationTask(any(AiGenerationTask.class));
|
||||
assertEquals(Long.valueOf(99L), response.getTaskId());
|
||||
assertEquals("RUNNING", response.getStatus());
|
||||
assertEquals(Integer.valueOf(35), response.getProgress());
|
||||
}
|
||||
|
||||
@Test(expected = ServiceException.class)
|
||||
public void createTaskRejectsMissingProjectOwnership()
|
||||
{
|
||||
when(frontProjectMapper.selectFrontProjectByUserAndId(7L, 10L)).thenReturn(null);
|
||||
|
||||
service.createTask(7L, 10L, request("database"));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void getTaskReturnsOnlyOwnedTask()
|
||||
{
|
||||
AiGenerationTask task = new AiGenerationTask();
|
||||
task.setTaskId(99L);
|
||||
task.setProjectId(10L);
|
||||
task.setUserId(7L);
|
||||
task.setGenerateType("database");
|
||||
task.setStatus("SUCCEEDED");
|
||||
task.setResultPayload("{\"tables\":[]}");
|
||||
when(aiGenerationTaskMapper.selectTaskForUser(7L, 10L, 99L)).thenReturn(task);
|
||||
|
||||
AiGenerationTaskStatusResponse response = service.getTask(7L, 10L, 99L);
|
||||
|
||||
assertSame(task.getTaskId(), response.getTaskId());
|
||||
assertEquals("SUCCEEDED", response.getStatus());
|
||||
assertEquals("{\"tables\":[]}", response.getResultPayload());
|
||||
}
|
||||
|
||||
private AiGenerationTaskCreateRequest request(String generateType)
|
||||
{
|
||||
AiGenerationTaskCreateRequest request = new AiGenerationTaskCreateRequest();
|
||||
request.setGenerateType(generateType);
|
||||
request.setProjectName("客户中心");
|
||||
request.setProjectDesc("管理客户资料");
|
||||
request.setIndustryTemplate("crm");
|
||||
request.setExtraRequirements("客户分层");
|
||||
request.setAppBlueprint("{\"roles\":[]}");
|
||||
return request;
|
||||
}
|
||||
|
||||
private FrontProject project()
|
||||
{
|
||||
FrontProject project = new FrontProject();
|
||||
project.setProjectId(10L);
|
||||
project.setUserId(7L);
|
||||
project.setProjectName("客户中心");
|
||||
return project;
|
||||
}
|
||||
|
||||
private void setField(String name, Object value) throws Exception
|
||||
{
|
||||
Field field = AiGenerationTaskServiceImpl.class.getDeclaredField(name);
|
||||
field.setAccessible(true);
|
||||
field.set(service, value);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
package com.ruoyi.generator.service.front;
|
||||
|
||||
import static org.junit.Assert.assertEquals;
|
||||
import static org.junit.Assert.assertTrue;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
import java.lang.reflect.Field;
|
||||
import java.util.Arrays;
|
||||
import java.util.Collections;
|
||||
import org.junit.Before;
|
||||
import org.junit.Test;
|
||||
import org.mockito.ArgumentCaptor;
|
||||
import org.mockito.Mock;
|
||||
import org.mockito.MockitoAnnotations;
|
||||
import com.alibaba.fastjson2.JSON;
|
||||
import com.ruoyi.generator.domain.front.AiGenerationTask;
|
||||
import com.ruoyi.generator.domain.front.AiUsageLedger;
|
||||
import com.ruoyi.generator.domain.front.dto.AiGenerationTaskCreateRequest;
|
||||
import com.ruoyi.generator.domain.front.dto.AppBlueprintDesign;
|
||||
import com.ruoyi.generator.domain.front.dto.GenerateAppBlueprintRequest;
|
||||
import com.ruoyi.generator.mapper.front.AiGenerationTaskMapper;
|
||||
import com.ruoyi.generator.mapper.front.AiUsageLedgerMapper;
|
||||
|
||||
public class AiGenerationTaskWorkerTest
|
||||
{
|
||||
private AiGenerationTaskWorker worker;
|
||||
|
||||
@Mock
|
||||
private AiGenerationTaskMapper aiGenerationTaskMapper;
|
||||
@Mock
|
||||
private AiUsageLedgerMapper aiUsageLedgerMapper;
|
||||
@Mock
|
||||
private IAiGenerateService aiGenerateService;
|
||||
@Mock
|
||||
private AiQuotaService aiQuotaService;
|
||||
|
||||
@Before
|
||||
public void setUp() throws Exception
|
||||
{
|
||||
MockitoAnnotations.initMocks(this);
|
||||
worker = new AiGenerationTaskWorker();
|
||||
setField("aiGenerationTaskMapper", aiGenerationTaskMapper);
|
||||
setField("aiUsageLedgerMapper", aiUsageLedgerMapper);
|
||||
setField("aiGenerateService", aiGenerateService);
|
||||
setField("aiQuotaService", aiQuotaService);
|
||||
setField("retryPolicy", new AiTaskRetryPolicy());
|
||||
setField("costCalculator", new AiCostCalculator());
|
||||
}
|
||||
|
||||
@Test
|
||||
public void processTaskMarksSuccessAndStoresResult()
|
||||
{
|
||||
AiGenerationTask task = task("app_blueprint");
|
||||
when(aiGenerationTaskMapper.selectAiGenerationTaskById(88L)).thenReturn(task);
|
||||
when(aiGenerationTaskMapper.claimTask(org.mockito.ArgumentMatchers.eq(88L), any(String.class))).thenReturn(1);
|
||||
AppBlueprintDesign blueprint = new AppBlueprintDesign();
|
||||
blueprint.setRoles(Collections.emptyList());
|
||||
when(aiGenerateService.generateAppBlueprint(any(Long.class), any(Long.class), any(GenerateAppBlueprintRequest.class)))
|
||||
.thenReturn(blueprint);
|
||||
|
||||
worker.processTask(88L);
|
||||
|
||||
ArgumentCaptor<AiGenerationTask> taskCaptor = ArgumentCaptor.forClass(AiGenerationTask.class);
|
||||
verify(aiGenerationTaskMapper, org.mockito.Mockito.atLeastOnce()).updateAiGenerationTask(taskCaptor.capture());
|
||||
AiGenerationTask saved = taskCaptor.getAllValues().get(taskCaptor.getAllValues().size() - 1);
|
||||
assertEquals("SUCCEEDED", saved.getStatus());
|
||||
assertEquals(Integer.valueOf(100), saved.getProgress());
|
||||
assertTrue(saved.getTotalTokens() > 0);
|
||||
assertTrue(saved.getCostCents() > 0);
|
||||
assertTrue(saved.getResultPayload().contains("\"roles\":[]"));
|
||||
verify(aiQuotaService).releaseRunning(7L);
|
||||
verify(aiUsageLedgerMapper).insertAiUsageLedger(any(AiUsageLedger.class));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void processTaskSchedulesRetryForRetryableFailure()
|
||||
{
|
||||
AiGenerationTask task = task("app_blueprint");
|
||||
task.setAttempts(1);
|
||||
when(aiGenerationTaskMapper.selectAiGenerationTaskById(88L)).thenReturn(task);
|
||||
when(aiGenerationTaskMapper.claimTask(org.mockito.ArgumentMatchers.eq(88L), any(String.class))).thenReturn(1);
|
||||
when(aiGenerateService.generateAppBlueprint(any(Long.class), any(Long.class), any(GenerateAppBlueprintRequest.class)))
|
||||
.thenThrow(new RuntimeException("DeepSeek request timeout"));
|
||||
|
||||
worker.processTask(88L);
|
||||
|
||||
ArgumentCaptor<AiGenerationTask> taskCaptor = ArgumentCaptor.forClass(AiGenerationTask.class);
|
||||
verify(aiGenerationTaskMapper, org.mockito.Mockito.atLeastOnce()).updateAiGenerationTask(taskCaptor.capture());
|
||||
AiGenerationTask saved = taskCaptor.getAllValues().get(taskCaptor.getAllValues().size() - 1);
|
||||
assertEquals("RETRY_WAITING", saved.getStatus());
|
||||
assertEquals(Integer.valueOf(2), saved.getAttempts());
|
||||
verify(aiQuotaService).releaseRunning(7L);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void processTaskMarksFailedForNonRetryableFailure()
|
||||
{
|
||||
AiGenerationTask task = task("app_blueprint");
|
||||
when(aiGenerationTaskMapper.selectAiGenerationTaskById(88L)).thenReturn(task);
|
||||
when(aiGenerationTaskMapper.claimTask(org.mockito.ArgumentMatchers.eq(88L), any(String.class))).thenReturn(1);
|
||||
when(aiGenerateService.generateAppBlueprint(any(Long.class), any(Long.class), any(GenerateAppBlueprintRequest.class)))
|
||||
.thenThrow(new RuntimeException("配额不足"));
|
||||
|
||||
worker.processTask(88L);
|
||||
|
||||
ArgumentCaptor<AiGenerationTask> taskCaptor = ArgumentCaptor.forClass(AiGenerationTask.class);
|
||||
verify(aiGenerationTaskMapper, org.mockito.Mockito.atLeastOnce()).updateAiGenerationTask(taskCaptor.capture());
|
||||
AiGenerationTask saved = taskCaptor.getAllValues().get(taskCaptor.getAllValues().size() - 1);
|
||||
assertEquals("FAILED", saved.getStatus());
|
||||
assertEquals(Integer.valueOf(1), saved.getAttempts());
|
||||
verify(aiQuotaService).releaseRunning(7L);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void processDueTasksRunsClaimableTasks()
|
||||
{
|
||||
AiGenerationTask first = task("app_blueprint");
|
||||
first.setTaskId(88L);
|
||||
AiGenerationTask second = task("app_blueprint");
|
||||
second.setTaskId(89L);
|
||||
when(aiGenerationTaskMapper.selectClaimableTasks(2)).thenReturn(Arrays.asList(first, second));
|
||||
when(aiGenerationTaskMapper.selectAiGenerationTaskById(88L)).thenReturn(first);
|
||||
when(aiGenerationTaskMapper.selectAiGenerationTaskById(89L)).thenReturn(second);
|
||||
when(aiGenerationTaskMapper.claimTask(any(Long.class), any(String.class))).thenReturn(1);
|
||||
AppBlueprintDesign blueprint = new AppBlueprintDesign();
|
||||
blueprint.setRoles(Collections.emptyList());
|
||||
when(aiGenerateService.generateAppBlueprint(any(Long.class), any(Long.class), any(GenerateAppBlueprintRequest.class)))
|
||||
.thenReturn(blueprint);
|
||||
|
||||
int count = worker.processDueTasks(2);
|
||||
|
||||
assertEquals(2, count);
|
||||
verify(aiGenerationTaskMapper).selectClaimableTasks(2);
|
||||
verify(aiGenerationTaskMapper).selectAiGenerationTaskById(88L);
|
||||
verify(aiGenerationTaskMapper).selectAiGenerationTaskById(89L);
|
||||
}
|
||||
|
||||
private AiGenerationTask task(String generateType)
|
||||
{
|
||||
AiGenerationTaskCreateRequest request = new AiGenerationTaskCreateRequest();
|
||||
request.setGenerateType(generateType);
|
||||
request.setProjectName("客户中心");
|
||||
request.setProjectDesc("管理客户资料");
|
||||
request.setIndustryTemplate("crm");
|
||||
request.setExtraRequirements("客户分层");
|
||||
AiGenerationTask task = new AiGenerationTask();
|
||||
task.setTaskId(88L);
|
||||
task.setProjectId(10L);
|
||||
task.setUserId(7L);
|
||||
task.setGenerateType(generateType);
|
||||
task.setProvider("deepseek");
|
||||
task.setModel("deepseek-chat");
|
||||
task.setStatus("QUEUED");
|
||||
task.setAttempts(0);
|
||||
task.setMaxAttempts(3);
|
||||
task.setRequestPayload(JSON.toJSONString(request));
|
||||
return task;
|
||||
}
|
||||
|
||||
private void setField(String name, Object value) throws Exception
|
||||
{
|
||||
Field field = AiGenerationTaskWorker.class.getDeclaredField(name);
|
||||
field.setAccessible(true);
|
||||
field.set(worker, value);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package com.ruoyi.generator.service.front;
|
||||
|
||||
import static org.junit.Assert.assertFalse;
|
||||
import static org.junit.Assert.assertTrue;
|
||||
import static org.junit.Assert.assertEquals;
|
||||
|
||||
import java.util.Date;
|
||||
import org.junit.Test;
|
||||
|
||||
public class AiTaskRetryPolicyTest
|
||||
{
|
||||
@Test
|
||||
public void retryableFailuresCanRetryBeforeMaxAttempts()
|
||||
{
|
||||
AiTaskRetryPolicy policy = new AiTaskRetryPolicy();
|
||||
|
||||
assertTrue(policy.canRetry("DeepSeek request timeout", 1, 3));
|
||||
assertTrue(policy.canRetry("429 too many requests", 1, 3));
|
||||
assertTrue(policy.canRetry("DeepSeek unavailable", 2, 3));
|
||||
assertTrue(policy.canRetry("HTTP 502 bad gateway", 2, 3));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void nonRetryableFailuresDoNotRetry()
|
||||
{
|
||||
AiTaskRetryPolicy policy = new AiTaskRetryPolicy();
|
||||
|
||||
assertFalse(policy.canRetry("配额不足", 1, 3));
|
||||
assertFalse(policy.canRetry("项目不存在或无权限", 1, 3));
|
||||
assertFalse(policy.canRetry("Business action references missing field", 1, 3));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void maxAttemptsStopsRetry()
|
||||
{
|
||||
AiTaskRetryPolicy policy = new AiTaskRetryPolicy();
|
||||
|
||||
assertFalse(policy.canRetry("DeepSeek request timeout", 3, 3));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void nextRetryTimeUsesBackoffSequence()
|
||||
{
|
||||
AiTaskRetryPolicy policy = new AiTaskRetryPolicy();
|
||||
Date now = new Date(1_000L);
|
||||
|
||||
assertEquals(new Date(31_000L), policy.nextRetryTime(1, now));
|
||||
assertEquals(new Date(121_000L), policy.nextRetryTime(2, now));
|
||||
assertEquals(new Date(301_000L), policy.nextRetryTime(3, now));
|
||||
assertEquals(new Date(301_000L), policy.nextRetryTime(9, now));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user