mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-03-30 04:40:59 +00:00
feat(task): add adaptor billing interface and async settlement framework
Add three billing lifecycle methods to the TaskAdaptor interface: - EstimateBilling: compute OtherRatios from user request before pricing - AdjustBillingOnSubmit: adjust ratios from upstream submit response - AdjustBillingOnComplete: determine final quota at task terminal state Introduce BaseBilling as embeddable no-op default for adaptors without custom billing. Move Sora/Ali OtherRatios logic from shared validation into per-adaptor EstimateBilling implementations. Add TaskBillingContext to persist pricing params (model_price, group_ratio, other_ratios) in task private data for async polling settlement. Extract RecalculateTaskQuota as a general-purpose delta settlement function and unify polling billing via settleTaskBillingOnComplete (adaptor-first, then token-based fallback).
This commit is contained in:
@@ -130,6 +130,58 @@ func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) {
|
||||
model.RecordLog(task.UserId, model.LogTypeSystem, logContent)
|
||||
}
|
||||
|
||||
// RecalculateTaskQuota 通用的异步差额结算。
|
||||
// actualQuota 是任务完成后的实际应扣额度,与预扣额度 (task.Quota) 做差额结算。
|
||||
// reason 用于日志记录(例如 "token重算" 或 "adaptor调整")。
|
||||
func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int, reason string) {
|
||||
if actualQuota <= 0 {
|
||||
return
|
||||
}
|
||||
preConsumedQuota := task.Quota
|
||||
quotaDelta := actualQuota - preConsumedQuota
|
||||
|
||||
if quotaDelta == 0 {
|
||||
logger.LogInfo(ctx, fmt.Sprintf("任务 %s 预扣费准确(%s,%s)",
|
||||
task.TaskID, logger.LogQuota(actualQuota), reason))
|
||||
return
|
||||
}
|
||||
|
||||
logger.LogInfo(ctx, fmt.Sprintf("任务 %s 差额结算:delta=%s(实际:%s,预扣:%s,%s)",
|
||||
task.TaskID,
|
||||
logger.LogQuota(quotaDelta),
|
||||
logger.LogQuota(actualQuota),
|
||||
logger.LogQuota(preConsumedQuota),
|
||||
reason,
|
||||
))
|
||||
|
||||
// 调整资金来源
|
||||
if err := taskAdjustFunding(task, quotaDelta); err != nil {
|
||||
logger.LogError(ctx, fmt.Sprintf("差额结算资金调整失败 task %s: %s", task.TaskID, err.Error()))
|
||||
return
|
||||
}
|
||||
|
||||
// 调整令牌额度
|
||||
taskAdjustTokenQuota(ctx, task, quotaDelta)
|
||||
|
||||
// 更新统计(仅补扣时更新,退还不影响已用统计)
|
||||
if quotaDelta > 0 {
|
||||
model.UpdateUserUsedQuotaAndRequestCount(task.UserId, quotaDelta)
|
||||
model.UpdateChannelUsedQuota(task.ChannelId, quotaDelta)
|
||||
}
|
||||
task.Quota = actualQuota
|
||||
|
||||
var action string
|
||||
if quotaDelta > 0 {
|
||||
action = "补扣费"
|
||||
} else {
|
||||
action = "退还"
|
||||
}
|
||||
logContent := fmt.Sprintf("异步任务成功%s,预扣费 %s,实际扣费 %s,原因:%s",
|
||||
action,
|
||||
logger.LogQuota(preConsumedQuota), logger.LogQuota(actualQuota), reason)
|
||||
model.RecordLog(task.UserId, model.LogTypeSystem, logContent)
|
||||
}
|
||||
|
||||
// RecalculateTaskQuotaByTokens 根据实际 token 消耗重新计费(异步差额结算)。
|
||||
// 当任务成功且返回了 totalTokens 时,根据模型倍率和分组倍率重新计算实际扣费额度,
|
||||
// 与预扣费的差额进行补扣或退还。支持钱包和订阅计费来源。
|
||||
@@ -180,48 +232,6 @@ func RecalculateTaskQuotaByTokens(ctx context.Context, task *model.Task, totalTo
|
||||
// 计算实际应扣费额度: totalTokens * modelRatio * groupRatio
|
||||
actualQuota := int(float64(totalTokens) * modelRatio * finalGroupRatio)
|
||||
|
||||
// 计算差额(正数=需要补扣,负数=需要退还)
|
||||
preConsumedQuota := task.Quota
|
||||
quotaDelta := actualQuota - preConsumedQuota
|
||||
|
||||
if quotaDelta == 0 {
|
||||
logger.LogInfo(ctx, fmt.Sprintf("视频任务 %s 预扣费准确(%s,tokens:%d)",
|
||||
task.TaskID, logger.LogQuota(actualQuota), totalTokens))
|
||||
return
|
||||
}
|
||||
|
||||
logger.LogInfo(ctx, fmt.Sprintf("视频任务 %s 差额结算:delta=%s(实际:%s,预扣:%s,tokens:%d)",
|
||||
task.TaskID,
|
||||
logger.LogQuota(quotaDelta),
|
||||
logger.LogQuota(actualQuota),
|
||||
logger.LogQuota(preConsumedQuota),
|
||||
totalTokens,
|
||||
))
|
||||
|
||||
// 调整资金来源
|
||||
if err := taskAdjustFunding(task, quotaDelta); err != nil {
|
||||
logger.LogError(ctx, fmt.Sprintf("差额结算资金调整失败 task %s: %s", task.TaskID, err.Error()))
|
||||
return
|
||||
}
|
||||
|
||||
// 调整令牌额度
|
||||
taskAdjustTokenQuota(ctx, task, quotaDelta)
|
||||
|
||||
// 更新统计(仅补扣时更新,退还不影响已用统计)
|
||||
if quotaDelta > 0 {
|
||||
model.UpdateUserUsedQuotaAndRequestCount(task.UserId, quotaDelta)
|
||||
model.UpdateChannelUsedQuota(task.ChannelId, quotaDelta)
|
||||
}
|
||||
task.Quota = actualQuota
|
||||
|
||||
var action string
|
||||
if quotaDelta > 0 {
|
||||
action = "补扣费"
|
||||
} else {
|
||||
action = "退还"
|
||||
}
|
||||
logContent := fmt.Sprintf("视频任务成功%s,模型倍率 %.2f,分组倍率 %.2f,tokens %d,预扣费 %s,实际扣费 %s",
|
||||
action, modelRatio, finalGroupRatio, totalTokens,
|
||||
logger.LogQuota(preConsumedQuota), logger.LogQuota(actualQuota))
|
||||
model.RecordLog(task.UserId, model.LogTypeSystem, logContent)
|
||||
reason := fmt.Sprintf("token重算:tokens=%d, modelRatio=%.2f, groupRatio=%.2f", totalTokens, modelRatio, finalGroupRatio)
|
||||
RecalculateTaskQuota(ctx, task, actualQuota, reason)
|
||||
}
|
||||
|
||||
@@ -26,6 +26,9 @@ type TaskPollingAdaptor interface {
|
||||
Init(info *relaycommon.RelayInfo)
|
||||
FetchTask(baseURL string, key string, body map[string]any, proxy string) (*http.Response, error)
|
||||
ParseTaskResult(body []byte) (*relaycommon.TaskInfo, error)
|
||||
// AdjustBillingOnComplete 在任务到达终态(成功/失败)时由轮询循环调用。
|
||||
// 返回正数触发差额结算(补扣/退还),返回 0 保持预扣费金额不变。
|
||||
AdjustBillingOnComplete(task *model.Task, taskResult *relaycommon.TaskInfo) int
|
||||
}
|
||||
|
||||
// GetTaskAdaptorFunc 由 main 包注入,用于获取指定平台的任务适配器。
|
||||
@@ -372,10 +375,8 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *
|
||||
task.PrivateData.ResultURL = taskcommon.BuildProxyURL(task.TaskID)
|
||||
}
|
||||
|
||||
// 如果返回了 total_tokens,根据模型倍率重新计费
|
||||
if taskResult.TotalTokens > 0 {
|
||||
RecalculateTaskQuotaByTokens(ctx, task, taskResult.TotalTokens)
|
||||
}
|
||||
// 完成时计费调整:优先由 adaptor 计算,回退到 token 重算
|
||||
settleTaskBillingOnComplete(ctx, adaptor, task, taskResult)
|
||||
case model.TaskStatusFailure:
|
||||
logger.LogJson(ctx, fmt.Sprintf("Task %s failed", taskId), task)
|
||||
task.Status = model.TaskStatusFailure
|
||||
@@ -444,3 +445,22 @@ func truncateBase64(s string) string {
|
||||
}
|
||||
return s[:maxKeep] + "..."
|
||||
}
|
||||
|
||||
// settleTaskBillingOnComplete 任务完成时的统一计费调整。
|
||||
// 优先级:1. adaptor.AdjustBillingOnComplete 返回正数 → 使用 adaptor 计算的额度
|
||||
//
|
||||
// 2. taskResult.TotalTokens > 0 → 按 token 重算
|
||||
// 3. 都不满足 → 保持预扣额度不变
|
||||
func settleTaskBillingOnComplete(ctx context.Context, adaptor TaskPollingAdaptor, task *model.Task, taskResult *relaycommon.TaskInfo) {
|
||||
// 1. 优先让 adaptor 决定最终额度
|
||||
if actualQuota := adaptor.AdjustBillingOnComplete(task, taskResult); actualQuota > 0 {
|
||||
RecalculateTaskQuota(ctx, task, actualQuota, "adaptor计费调整")
|
||||
return
|
||||
}
|
||||
// 2. 回退到 token 重算
|
||||
if taskResult.TotalTokens > 0 {
|
||||
RecalculateTaskQuotaByTokens(ctx, task, taskResult.TotalTokens)
|
||||
return
|
||||
}
|
||||
// 3. 无调整,保持预扣额度
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user