TASK: 对话 扣除积分

src/main/java/com/ai/da/common/enums/CreditsEventsEnum.java

@@ -57,6 +57,8 @@ public enum CreditsEventsEnum {
     LOCAL_TEXT2IMG_HIGH("Local_text2img_high","5"),
     LOCAL_IMG2IMG_HIGH("Local_img2img_high","5"),
     LOCAL_ANIMATION("Local_Animation","15"),
+
+    LLM_CONVERSATION("LLM_Conversation", "0")
     ;
 
     private final String name;

src/main/java/com/ai/da/mapper/primary/entity/ChatMessage.java

@@ -39,6 +39,22 @@ public class ChatMessage implements Serializable {
 
     @ApiModelProperty("0对话内容1颜色2图片")
     private Integer isImage;
+    /**
+     * 输入
+     */
+    private Long inputTokens;
+    /**
+     * 输出
+     */
+    private Long outputTokens;
+    /**
+     * 思考
+     */
+    private Long reasoningTokens;
+    /**
+     * 本次输出消耗的总金额
+     */
+    private String totalCost;
 
     @ApiModelProperty("创建时间")
     private LocalDateTime createTime;

src/main/java/com/ai/da/mapper/primary/entity/WorkspaceRelStyle.java

@@ -9,10 +9,10 @@ import java.io.Serializable;
 
 @Data
 @TableName("workspace_rel_style")
-public class WorkspaceRelStyle implements Serializable {
+public class WorkspaceRelStyle extends BaseEntity implements Serializable  {
     private static final long serialVersionUID = 1L;
-    @TableId(value = "id", type = IdType.AUTO)
+//    @TableId(value = "id", type = IdType.AUTO)
-    private Long id;
+//    private Long id;
     private Long workspaceId;
     private Long styleId;
 }

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

@@ -7,6 +7,8 @@ import com.ai.da.mapper.primary.entity.CreditsDetail;
 import com.ai.da.model.dto.QueryIncomeOrExpenditureDTO;
 import com.baomidou.mybatisplus.extension.service.IService;
 
+import java.math.BigDecimal;
+
 public interface CreditsService extends IService<CreditsDetail> {
 
     Boolean buyCredits(Long accountId, Float quantity);
@@ -40,4 +42,6 @@ public interface CreditsService extends IService<CreditsDetail> {
     void updateChangedCredits(String accountId, String taskId);
 
     CreditsDetail queryDetailByTaskId(String taskId);
+
+    void tokenUsage(Long accountId, BigDecimal totalTokenUsage);
 }

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

@@ -26,12 +26,10 @@ import org.springframework.util.CollectionUtils;
 
 import javax.annotation.Resource;
 import java.math.BigDecimal;
+import java.math.RoundingMode;
 import java.time.LocalDateTime;
 import java.time.format.DateTimeFormatter;
-import java.util.List;
+import java.util.*;
-import java.util.Objects;
-import java.util.Optional;
-import java.util.Set;
 
 @Service
 @Slf4j
@@ -368,4 +366,23 @@ public class CreditsServiceImpl extends ServiceImpl<CreditsDetailMapper, Credits
         return baseMapper.selectOne(queryWrapper);
     }
 
+    public void tokenUsage(Long accountId, BigDecimal totalTokenUsage){
+        Account account = accountService.getById(accountId);
+        BigDecimal changeCredits = totalTokenUsage.setScale(0, RoundingMode.CEILING);
+        BigDecimal subtracted = account.getCredits().subtract(changeCredits);
+        account.setCredits(subtracted);
+        account.setUpdateDate(new Date());
+        accountService.updateById(account);
+
+        CreditsDetail creditsDetail = new CreditsDetail();
+        creditsDetail.setAccountId(accountId);
+        creditsDetail.setChangeEvent(CreditsEventsEnum.LLM_CONVERSATION.getName());
+        creditsDetail.setChangedCredits("-" + changeCredits);
+        creditsDetail.setCredits(subtracted);
+        creditsDetail.setCreateTime(LocalDateTime.now());
+        baseMapper.insert(creditsDetail);
+        log.info("对话扣除积分，原始积分消耗量：{}， 取整：{}， 剩余积分：{}", totalTokenUsage, changeCredits, subtracted);
+    }
+
+
 }

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

@@ -18,10 +18,7 @@ import com.ai.da.model.enums.*;
 import com.ai.da.model.vo.AuthPrincipalVo;
 import com.ai.da.model.vo.UserLikeGroupVO;
 import com.ai.da.python.PythonService;
-import com.ai.da.service.DesignService;
+import com.ai.da.service.*;
-import com.ai.da.service.LLMService;
-import com.ai.da.service.SysFileService;
-import com.ai.da.service.WorkspaceService;
 import com.alibaba.fastjson.JSON;
 import com.alibaba.fastjson.JSONArray;
 import com.alibaba.fastjson.JSONObject;
@@ -42,6 +39,8 @@ import javax.annotation.Resource;
 import java.io.BufferedReader;
 import java.io.InputStreamReader;
 import java.io.OutputStream;
+import java.math.BigDecimal;
+import java.math.RoundingMode;
 import java.net.HttpURLConnection;
 import java.net.URL;
 import java.nio.charset.StandardCharsets;
@@ -81,6 +80,8 @@ public class LLMServiceImpl implements LLMService {
     private CollectionElementRelModelMapper collectionElementRelModelMapper;
     @Value("${minio.bucketName.sysImage}")
     private String sysImage;
+    @Resource
+    private CreditsService creditsService;
 
     @Override
     public SseEmitter stream(String prompt, Long projectParamId, String fileUrl, List<String> imageUrlList, String token, Boolean enableThinking, String process) {
@@ -295,6 +296,14 @@ public class LLMServiceImpl implements LLMService {
                     systemMessage.setAccountId(accountId);
                     StringBuilder responseContentBuilder = new StringBuilder();
                     String contentType = "";
+                    BigDecimal totalTokenUsage = BigDecimal.ZERO;
+                    // 需要存input_token  output_token reasoning_tokens total_cost
+                    Long inputToken = 0L;
+                    Long outputToken = 0L;
+                    Long reasoningToken = 0L;
+                    String totalCost = "0";
+                    boolean tokenUsageFlag = false;
+
                     try (BufferedReader reader = new BufferedReader(new InputStreamReader(conn.getInputStream(), StandardCharsets.UTF_8))) {
                         String line;
                         while ((line = reader.readLine()) != null) {
@@ -334,6 +343,13 @@ public class LLMServiceImpl implements LLMService {
                                             systemMessage.setSeq(getNextSeq(projectId));
                                             systemMessage.setCreateTime(LocalDateTime.now());
                                             systemMessage.setContent(responseContentBuilder.toString());
+                                            if (tokenUsageFlag){
+                                                systemMessage.setInputTokens(inputToken);
+                                                systemMessage.setOutputTokens(outputToken);
+                                                systemMessage.setReasoningTokens(reasoningToken);
+                                                systemMessage.setTotalCost(totalCost);
+                                                tokenUsageFlag = false;
+                                            }
                                             chatMessageMapper.insert(systemMessage);
                                             systemMessage.setId(null);
                                             responseContentBuilder = new StringBuilder();
@@ -368,6 +384,13 @@ public class LLMServiceImpl implements LLMService {
                                             systemImage.setCreateTime(LocalDateTime.now());
                                             systemImage.setContent(contentSave);
                                             systemImage.setAccountId(accountId);
+                                            if (tokenUsageFlag){
+                                                systemImage.setInputTokens(inputToken);
+                                                systemImage.setOutputTokens(outputToken);
+                                                systemImage.setReasoningTokens(reasoningToken);
+                                                systemImage.setTotalCost(totalCost);
+                                                tokenUsageFlag = false;
+                                            }
                                             chatMessageMapper.insert(systemImage);
                                         } else if (Objects.nonNull(toolsName) && toolsName.equals("search_sketch_img")) {
                                             json.put("content", processSearchSketchToolCon(projectId, toolsData));
@@ -376,6 +399,16 @@ public class LLMServiceImpl implements LLMService {
                                             updateProjectParams(projectId, toolsData);
                                         }
                                         emitter.send(json.toJSONString());
+                                    }else if ("cost".equals(type)) {
+                                        JSONObject contentObj = json.getJSONObject("content");
+                                        log.info("token usage: {}", contentObj);
+                                        inputToken = contentObj.getLong("input_tokens");
+                                        outputToken = contentObj.getLong("output_tokens");
+                                        reasoningToken = contentObj.getLong("reasoning_tokens");
+                                        totalCost = contentObj.getString("total_cost");
+                                        totalTokenUsage = totalTokenUsage.add(new BigDecimal(totalCost).divide(new BigDecimal("0.03"), 6, RoundingMode.CEILING));
+                                        log.info("totalTokenUsage: {}", totalTokenUsage);
+                                        tokenUsageFlag = true;
                                     }
                                 }
                             }
@@ -384,8 +417,16 @@ public class LLMServiceImpl implements LLMService {
                             systemMessage.setSeq(getNextSeq(projectId));
                             systemMessage.setCreateTime(LocalDateTime.now());
                             systemMessage.setContent(responseContentBuilder.toString());
+                            if (tokenUsageFlag){
+                                systemMessage.setInputTokens(inputToken);
+                                systemMessage.setOutputTokens(outputToken);
+                                systemMessage.setReasoningTokens(reasoningToken);
+                                systemMessage.setTotalCost(totalCost);
+                            }
                             chatMessageMapper.insert(systemMessage);
                         }
+                        // 扣积分
+                        creditsService.tokenUsage(accountId, totalTokenUsage);
                     }
                 }
                 emitter.complete();