mirror of
https://gitcode.com/ageerle/ruoyi-ai.git
synced 2026-03-13 20:53:42 +08:00
fix(billing): 修复Token计费逻辑和消息更新机制
* 修复Token计费算法:按批次计费而非Token数量计费 * 添加ChatRequest.messageId字段支持消息关联更新 * 优化消息保存流程:分离基础信息保存和计费信息更新 * 修复预检查逻辑:统一预检查和实际扣费计算方式 * 调整Token阈值:100 → 1000,减少扣费频次 * 完善事件传递:ChatMessageCreatedEvent增加messageId Fixes: 余额预检查误判、消息计费信息缺失、Token计费不准确
This commit is contained in:
@@ -82,4 +82,9 @@ public class ChatRequest {
|
||||
*/
|
||||
private String token;
|
||||
|
||||
/**
|
||||
* 消息ID(保存消息成功后设置,用于后续扣费更新)
|
||||
*/
|
||||
private Long messageId;
|
||||
|
||||
}
|
||||
|
||||
@@ -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; }
|
||||
}
|
||||
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -79,13 +79,13 @@ public class ChatCostServiceImpl implements IChatCostService {
|
||||
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,10 +105,13 @@ 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 {
|
||||
// 先尝试扣费
|
||||
@@ -120,8 +123,8 @@ public class ChatCostServiceImpl implements IChatCostService {
|
||||
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) {
|
||||
@@ -160,31 +167,18 @@ public class ChatCostServiceImpl implements IChatCostService {
|
||||
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("保存消息失败");
|
||||
@@ -212,7 +206,8 @@ public class ChatCostServiceImpl implements IChatCostService {
|
||||
chatRequest.getSessionId(),
|
||||
chatRequest.getModel(),
|
||||
chatRequest.getRole(),
|
||||
chatRequest.getPrompt()
|
||||
chatRequest.getPrompt(),
|
||||
chatRequest.getMessageId()
|
||||
));
|
||||
log.debug("publishBillingEvent->计费事件发布完成");
|
||||
}
|
||||
@@ -238,15 +233,17 @@ public class ChatCostServiceImpl implements IChatCostService {
|
||||
}
|
||||
|
||||
// 按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());
|
||||
}
|
||||
@@ -272,57 +269,85 @@ public class ChatCostServiceImpl implements IChatCostService {
|
||||
}
|
||||
|
||||
/**
|
||||
* 保存账单记录
|
||||
* 更新消息的基本信息(不涉及扣费)
|
||||
*/
|
||||
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));
|
||||
// 创建更新对象,只更新基本信息,不涉及扣费
|
||||
ChatMessageBo updateMessage = new ChatMessageBo();
|
||||
updateMessage.setId(chatRequest.getMessageId());
|
||||
updateMessage.setTotalTokens(actualTokens); // 设置实际token数
|
||||
updateMessage.setBillingType(billingTypeCode); // 设置计费类型
|
||||
updateMessage.setRemark("用户消息(累计中,未达扣费阈值)"); // 说明状态
|
||||
|
||||
// 设置计费类型和构建消息内容
|
||||
setBillingTypeAndContent(billingMessage, billingTypeCode, billedTokens, cost);
|
||||
|
||||
chatMessageService.insertByBo(billingMessage);
|
||||
log.debug("saveBillingRecord->保存账单记录成功,用户ID: {}, 计费类型: {}, 费用: {}",
|
||||
chatRequest.getUserId(), billingTypeCode, cost);
|
||||
// 更新消息
|
||||
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);
|
||||
}
|
||||
|
||||
|
||||
@@ -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());
|
||||
|
||||
Reference in New Issue
Block a user