xieyinlu/ruoyi-ai
1
0
Fork 0
mirror of https://gitcode.com/ageerle/ruoyi-ai.git synced 2026-03-13 20:53:42 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(billing): 修复Token计费逻辑和消息更新机制

Browse Source 
* 修复Token计费算法：按批次计费而非Token数量计费
* 添加ChatRequest.messageId字段支持消息关联更新
* 优化消息保存流程：分离基础信息保存和计费信息更新
* 修复预检查逻辑：统一预检查和实际扣费计算方式
* 调整Token阈值：100 → 1000，减少扣费频次
* 完善事件传递：ChatMessageCreatedEvent增加messageId

Fixes: 余额预检查误判、消息计费信息缺失、Token计费不准确
This commit is contained in:
Administrator
2025-08-29 15:19:37 +08:00
parent 1e3b49c9b8
commit 1e4af3d01b
5 changed files with 126 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
ruoyi-common/ruoyi-common-chat/src/main/java/org/ruoyi/common/chat/request/ChatRequest.java
View File

@@ -82,4 +82,9 @@ public class ChatRequest {
*/
private String token;
/**
* 消息ID（保存消息成功后设置，用于后续扣费更新）
*/
private Long messageId;
}

5
ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/chat/event/ChatMessageCreatedEvent.java
View File

@@ -12,14 +12,16 @@ public class ChatMessageCreatedEvent extends ApplicationEvent {
private final String modelName;
private final String role;
private final String content;
private final Long messageId;
public ChatMessageCreatedEvent(Long userId, Long sessionId, String modelName, String role, String content) {
public ChatMessageCreatedEvent(Long userId, Long sessionId, String modelName, String role, String content, Long messageId) {
super(userId);
this.userId = userId;
this.sessionId = sessionId;
this.modelName = modelName;
this.role = role;
this.content = content;
this.messageId = messageId;
}
public Long getUserId() { return userId; }
@@ -27,5 +29,6 @@ public class ChatMessageCreatedEvent extends ApplicationEvent {
public String getModelName() { return modelName; }
public String getRole() { return role; }
public String getContent() { return content; }
public Long getMessageId() { return messageId; }
}

1
ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/chat/listener/BillingEventListener.java
View File

@@ -30,6 +30,7 @@ public class BillingEventListener {
chatRequest.setModel(event.getModelName());
chatRequest.setRole(event.getRole());
chatRequest.setPrompt(event.getContent());
chatRequest.setMessageId(event.getMessageId()); // 设置消息ID
// 异步执行计费累计与扣费
log.debug("BillingEventListener->开始执行计费逻辑");
chatCostService.deductToken(chatRequest);

198
ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/chat/service/chat/impl/ChatCostServiceImpl.java
View File

@@ -70,7 +70,7 @@ public class ChatCostServiceImpl implements IChatCostService {
BigDecimal numberCost = unitPrice.setScale(2, RoundingMode.HALF_UP);
deductUserBalance(chatRequest.getUserId(), numberCost.doubleValue());
log.debug("deductToken->按次数扣费，费用: {}，模型: {}", numberCost, modelName);
// 清理可能存在的历史累计token（模型计费方式可能发生过变更）
ChatUsageToken existingToken = chatTokenService.queryByUserId(chatRequest.getUserId(), modelName);
if (existingToken != null && existingToken.getToken() > 0) {
@@ -78,14 +78,14 @@ public class ChatCostServiceImpl implements IChatCostService {
chatTokenService.editToken(existingToken);
log.debug("deductToken->按次计费，清理历史累计token: {}", existingToken.getToken());
}
// 记录账单消息
saveBillingRecord(chatRequest, tokens, numberCost.doubleValue(), chatModelVo.getModelType());
// 更新消息的计费信息到备注
updateMessageBilling(chatRequest, tokens, numberCost.doubleValue(), chatModelVo.getModelType());
return;
}
// 按token计费：累加并按阈值批量扣费，保留余数
final int threshold = 100;
final int threshold = 1000;
// 获得记录的累计token数
// TODO: 这里存在并发竞态条件，需要在chatTokenService层面添加乐观锁或分布式锁
@@ -105,11 +105,14 @@ public class ChatCostServiceImpl implements IChatCostService {
int remainder = totalTokens - billable; // 结算后保留的余数
if (billable > 0) {
// 计算批次数：每1000个Token为一批，每批扣费单价
int batches = billable / threshold;
BigDecimal numberCost = unitPrice
.multiply(BigDecimal.valueOf(billable))
.multiply(BigDecimal.valueOf(batches))
.setScale(2, RoundingMode.HALF_UP);
log.debug("deductToken->按token扣费，结算token数量: {}，单价: {}，费用: {}", billable, unitPrice, numberCost);
log.debug("deductToken->按token扣费，结算token数量: {}，批次数: {}，单价: {}，费用: {}",
billable, batches, unitPrice, numberCost);
try {
// 先尝试扣费
deductUserBalance(chatRequest.getUserId(), numberCost.doubleValue());
@@ -119,9 +122,9 @@ public class ChatCostServiceImpl implements IChatCostService {
chatToken.setToken(remainder);
chatTokenService.editToken(chatToken);
log.debug("deductToken->扣费成功，更新余数: {}", remainder);
// 记录账单消息
saveBillingRecord(chatRequest, billable, numberCost.doubleValue(), chatModelVo.getModelType());
// 更新消息的计费信息到备注
updateMessageBilling(chatRequest, billable, numberCost.doubleValue(), chatModelVo.getModelType());
} catch (ServiceException e) {
// 余额不足时，不更新token累计，保持原有累计数
log.warn("deductToken->余额不足，本次token累计保持不变: {}", totalTokens);
@@ -134,11 +137,15 @@ public class ChatCostServiceImpl implements IChatCostService {
chatToken.setUserId(chatRequest.getUserId());
chatToken.setToken(totalTokens);
chatTokenService.editToken(chatToken);
// 虽未扣费，但要更新消息的基本信息（实际token数、计费类型等）
updateMessageWithoutBilling(chatRequest, tokens, chatModelVo.getModelType());
}
}
/**
* 保存聊天消息记录（不进行计费）
* 保存成功后将消息ID设置到ChatRequest中，供后续扣费使用
*/
@Override
public void saveMessage(ChatRequest chatRequest) {
@@ -146,45 +153,32 @@ public class ChatCostServiceImpl implements IChatCostService {
log.warn("saveMessage->用户ID或会话ID为空，跳过保存消息");
return;
}
// 验证消息内容
if (chatRequest.getPrompt() == null || chatRequest.getPrompt().trim().isEmpty()) {
log.warn("saveMessage->消息内容为空，跳过保存");
return;
}
ChatMessageBo chatMessageBo = new ChatMessageBo();
chatMessageBo.setUserId(chatRequest.getUserId());
chatMessageBo.setSessionId(chatRequest.getSessionId());
chatMessageBo.setRole(chatRequest.getRole());
chatMessageBo.setContent(chatRequest.getPrompt().trim());
chatMessageBo.setModelName(chatRequest.getModel());
// 计算并保存本次消息的token数
int tokens = TikTokensUtil.tokens(chatRequest.getModel(), chatRequest.getPrompt());
chatMessageBo.setTotalTokens(tokens);
// 普通消息不涉及扣费，deductCost保持null
chatMessageBo.setDeductCost(null);
chatMessageBo.setRemark("用户消息");
// 设置计费类型（根据模型配置获取计费类型）
try {
ChatModelVo chatModelVo = chatModelService.selectModelByName(chatRequest.getModel());
if (chatModelVo != null) {
chatMessageBo.setBillingType(chatModelVo.getModelType());
} else {
chatMessageBo.setBillingType(null); // 模型不存在时设为null
}
} catch (Exception e) {
log.warn("saveMessage->获取模型计费类型失败，设为null，模型: {}", chatRequest.getModel());
chatMessageBo.setBillingType(null);
}
// // 基础消息信息，计费相关数据（tokens、费用、计费类型等）在扣费时统一设置
// chatMessageBo.setTotalTokens(0); // 初始设为0，扣费时更新
// chatMessageBo.setDeductCost(null);
// chatMessageBo.setBillingType(null);
// chatMessageBo.setRemark("用户消息");
try {
chatMessageService.insertByBo(chatMessageBo);
log.debug("saveMessage->成功保存消息，用户ID: {}, 会话ID: {}, tokens: {}",
chatRequest.getUserId(), chatRequest.getSessionId(), tokens);
// 保存成功后，将生成的消息ID设置到ChatRequest中
chatRequest.setMessageId(chatMessageBo.getId());
log.debug("saveMessage->成功保存消息，消息ID: {}, 用户ID: {}, 会话ID: {}",
chatMessageBo.getId(), chatRequest.getUserId(), chatRequest.getSessionId());
} catch (Exception e) {
log.error("saveMessage->保存消息失败", e);
throw new ServiceException("保存消息失败");
@@ -195,28 +189,29 @@ public class ChatCostServiceImpl implements IChatCostService {
@Override
public void publishBillingEvent(ChatRequest chatRequest) {
log.debug("publishBillingEvent->发布计费事件，用户ID: {}，会话ID: {}，模型: {}",
log.debug("publishBillingEvent->发布计费事件，用户ID: {}，会话ID: {}，模型: {}",
chatRequest.getUserId(), chatRequest.getSessionId(), chatRequest.getModel());
// 预检查：评估可能的扣费金额，如果余额不足则直接抛异常
try {
preCheckBalance(chatRequest);
} catch (ServiceException e) {
log.warn("publishBillingEvent->预检查余额不足，用户ID: {}，模型: {}",
log.warn("publishBillingEvent->预检查余额不足，用户ID: {}，模型: {}",
chatRequest.getUserId(), chatRequest.getModel());
throw e; // 直接抛出，阻止消息保存和对话继续
}
eventPublisher.publishEvent(new ChatMessageCreatedEvent(
chatRequest.getUserId(),
chatRequest.getSessionId(),
chatRequest.getModel(),
chatRequest.getRole(),
chatRequest.getPrompt()
chatRequest.getPrompt(),
chatRequest.getMessageId()
));
log.debug("publishBillingEvent->计费事件发布完成");
}
/**
* 预检查用户余额是否足够支付可能的费用
*/
@@ -224,34 +219,36 @@ public class ChatCostServiceImpl implements IChatCostService {
if (chatRequest.getUserId() == null) {
return;
}
int tokens = TikTokensUtil.tokens(chatRequest.getModel(), chatRequest.getPrompt());
String modelName = chatRequest.getModel();
ChatModelVo chatModelVo = chatModelService.selectModelByName(modelName);
BigDecimal unitPrice = BigDecimal.valueOf(chatModelVo.getModelPrice());
// 按次计费：直接检查单次费用
if (BillingType.TIMES.getCode().equals(chatModelVo.getModelType())) {
BigDecimal numberCost = unitPrice.setScale(2, RoundingMode.HALF_UP);
checkUserBalanceWithoutDeduct(chatRequest.getUserId(), numberCost.doubleValue());
return;
}
// 按token计费：检查累计后可能的费用
final int threshold = 100;
final int threshold = 1000;
ChatUsageToken chatToken = chatTokenService.queryByUserId(chatRequest.getUserId(), modelName);
int previousUnpaid = (chatToken == null) ? 0 : chatToken.getToken();
int totalTokens = previousUnpaid + tokens;
int billable = (totalTokens / threshold) * threshold;
if (billable > 0) {
// 计算批次数：每1000个Token为一批，每批扣费单价
int batches = billable / threshold;
BigDecimal numberCost = unitPrice
.multiply(BigDecimal.valueOf(billable))
.multiply(BigDecimal.valueOf(batches))
.setScale(2, RoundingMode.HALF_UP);
checkUserBalanceWithoutDeduct(chatRequest.getUserId(), numberCost.doubleValue());
}
}
/**
* 检查用户余额是否足够，但不扣除
*/
@@ -260,69 +257,97 @@ public class ChatCostServiceImpl implements IChatCostService {
if (sysUser == null) {
throw new ServiceException("用户不存在");
}
BigDecimal userBalance = BigDecimal.valueOf(sysUser.getUserBalance() == null ? 0D : sysUser.getUserBalance())
.setScale(2, RoundingMode.HALF_UP);
BigDecimal cost = BigDecimal.valueOf(numberCost == null ? 0D : numberCost)
.setScale(2, RoundingMode.HALF_UP);
if (userBalance.compareTo(cost) < 0 || userBalance.compareTo(BigDecimal.ZERO) == 0) {
throw new ServiceException("余额不足, 请充值。当前余额: " + userBalance + "，需要: " + cost);
}
}
/**
* 保存账单记录
* 更新消息的基本信息（不涉及扣费）
*/
private void saveBillingRecord(ChatRequest chatRequest, int billedTokens, double cost, String billingTypeCode) {
private void updateMessageWithoutBilling(ChatRequest chatRequest, int actualTokens, String billingTypeCode) {
// 检查是否有消息ID可以更新
if (chatRequest.getMessageId() == null) {
log.warn("updateMessageWithoutBilling->消息ID为空，无法更新基本信息");
return;
}
try {
ChatMessageBo billingMessage = new ChatMessageBo();
billingMessage.setUserId(chatRequest.getUserId());
billingMessage.setSessionId(chatRequest.getSessionId());
billingMessage.setRole("system"); // 系统账单消息
billingMessage.setModelName(chatRequest.getModel());
billingMessage.setTotalTokens(billedTokens);
billingMessage.setDeductCost(cost);
billingMessage.setRemark(getBillingTypeName(billingTypeCode));
// 设置计费类型和构建消息内容
setBillingTypeAndContent(billingMessage, billingTypeCode, billedTokens, cost);
chatMessageService.insertByBo(billingMessage);
log.debug("saveBillingRecord->保存账单记录成功，用户ID: {}, 计费类型: {}, 费用: {}",
chatRequest.getUserId(), billingTypeCode, cost);
// 创建更新对象，只更新基本信息，不涉及扣费
ChatMessageBo updateMessage = new ChatMessageBo();
updateMessage.setId(chatRequest.getMessageId());
updateMessage.setTotalTokens(actualTokens); // 设置实际token数
updateMessage.setBillingType(billingTypeCode); // 设置计费类型
updateMessage.setRemark("用户消息（累计中，未达扣费阈值）"); // 说明状态
// 更新消息
chatMessageService.updateByBo(updateMessage);
log.debug("updateMessageWithoutBilling->更新消息基本信息成功，消息ID: {}, 实际tokens: {}, 计费类型: {}",
chatRequest.getMessageId(), actualTokens, billingTypeCode);
} catch (Exception e) {
log.error("saveBillingRecord->保存账单记录失败", e);
// 账单记录失败不影响主流程，只记录错误日志
log.error("updateMessageWithoutBilling->更新消息基本信息失败，消息ID: {}", chatRequest.getMessageId(), e);
// 更新失败不影响主流程，只记录错误日志
}
}
/**
* 设置计费类型和构建消息内容
* 更新消息的计费信息到备注字段
*/
private void setBillingTypeAndContent(ChatMessageBo billingMessage, String billingTypeCode, int billedTokens, double cost) {
billingMessage.setBillingType(billingTypeCode);
// 使用枚举获取计费类型并构建消息内容
private void updateMessageBilling(ChatRequest chatRequest, int billedTokens, double cost, String billingTypeCode) {
// 检查是否有消息ID可以更新
if (chatRequest.getMessageId() == null) {
log.warn("updateMessageBilling->消息ID为空，无法更新计费信息");
return;
}
try {
// 计算本次消息的实际token数
int actualTokens = TikTokensUtil.tokens(chatRequest.getModel(), chatRequest.getPrompt());
// 构建计费信息
String billingInfo = buildBillingInfo(billingTypeCode, billedTokens, cost);
// 创建更新对象
ChatMessageBo updateMessage = new ChatMessageBo();
updateMessage.setId(chatRequest.getMessageId());
updateMessage.setTotalTokens(actualTokens); // 设置实际token数
updateMessage.setDeductCost(cost);
updateMessage.setRemark(billingInfo);
updateMessage.setBillingType(billingTypeCode);
// 更新消息
chatMessageService.updateByBo(updateMessage);
log.debug("updateMessageBilling->更新消息计费信息成功，消息ID: {}, 实际tokens: {}, 计费tokens: {}, 费用: {}",
chatRequest.getMessageId(), actualTokens, billedTokens, cost);
} catch (Exception e) {
log.error("updateMessageBilling->更新消息计费信息失败，消息ID: {}", chatRequest.getMessageId(), e);
// 更新失败不影响主流程，只记录错误日志
}
}
/**
* 构建计费信息字符串
*/
private String buildBillingInfo(String billingTypeCode, int billedTokens, double cost) {
// 使用枚举获取计费类型并构建计费信息
BillingType billingType = BillingType.fromCode(billingTypeCode);
if (billingType != null) {
String content = switch (billingType) {
return switch (billingType) {
case TIMES -> String.format("%s：消耗 %d tokens，扣费 %.2f 元", billingType.getDescription(), billedTokens, cost);
case TOKEN -> String.format("%s：结算 %d tokens，扣费 %.2f 元", billingType.getDescription(), billedTokens, cost);
};
billingMessage.setContent(content);
} else {
billingMessage.setContent(String.format("系统计费：处理 %d tokens，扣费 %.2f 元", billedTokens, cost));
return String.format("系统计费：处理 %d tokens，扣费 %.2f 元", billedTokens, cost);
}
}
/**
* 获取计费类型名称（用于remark字段）
*/
private String getBillingTypeName(String billingTypeCode) {
BillingType billingType = BillingType.fromCode(billingTypeCode);
return billingType != null ? billingType.getDescription() : "系统计费";
}
/**
* 从用户余额中扣除费用
@@ -377,6 +402,7 @@ public class ChatCostServiceImpl implements IChatCostService {
chatMessageBo.setContent(prompt);
chatMessageBo.setDeductCost(cost);
chatMessageBo.setTotalTokens(0);
chatMessageBo.setRemark(String.format("任务计费：%s，扣费 %.2f 元", type, cost));
chatMessageService.insertByBo(chatMessageBo);
}

6
ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/chat/service/chat/impl/DifyServiceImpl.java
View File

@@ -91,6 +91,8 @@ public class DifyServiceImpl implements IChatService {
public void onMessageEnd(MessageEndEvent event) {
emitter.complete();
log.info("消息结束，完整消息ID: {}", event.getMessageId());
// 扣除费用
ChatRequest chatRequestResponse = new ChatRequest();
// 更新conversationId
if (StrUtil.isBlank(sessionInfo.getConversationId())) {
String conversationId = event.getConversationId();
@@ -104,9 +106,9 @@ public class DifyServiceImpl implements IChatService {
chatSessionBo.setSessionContent(sessionInfo.getSessionContent());
chatSessionBo.setRemark(sessionInfo.getRemark());
chatSessionService.updateByBo(chatSessionBo);
chatRequestResponse.setMessageId(chatSessionBo.getId());
}
// 扣除费用
ChatRequest chatRequestResponse = new ChatRequest();
// 设置对话角色
chatRequestResponse.setRole(Message.Role.ASSISTANT.getName());
chatRequestResponse.setModel(chatRequest.getModel());