From db8eacd558dafae98afcab7059d312a4adfe59f7 Mon Sep 17 00:00:00 2001 From: fengsilin Date: Wed, 17 Jun 2026 18:25:11 +0800 Subject: [PATCH] feat(task-pricing): persist matrix billing snapshots --- controller/relay.go | 112 +++++++++++++++++++++++++++++++++++----- model/task.go | 16 +++--- relay/relay_task.go | 65 +++++++++++------------ service/task_billing.go | 42 +++++++++++++++ service/task_polling.go | 19 +++++++ types/model_pricing.go | 61 ++++++++++++++++++++++ 6 files changed, 263 insertions(+), 52 deletions(-) diff --git a/controller/relay.go b/controller/relay.go index 4ca1497..9c2d1c4 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -24,6 +24,7 @@ import ( "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/QuantumNous/new-api/types" "github.com/bytedance/gopkg/util/gopool" @@ -204,7 +205,6 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { break } - addUsedChannel(c, channel.Id) bodyStorage, bodyErr := common.GetBodyStorage(c) if bodyErr != nil { @@ -488,6 +488,95 @@ func RelayTaskFetch(c *gin.Context) { } } +func preloadTaskPricingConfig(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskError { + modelName := strings.TrimSpace(info.OriginModelName) + action := strings.TrimSpace(info.Action) + contentType := c.Request.Header.Get("Content-Type") + + storage, err := common.GetBodyStorage(c) + if err != nil { + status := http.StatusBadRequest + if common.IsRequestBodyTooLargeError(err) || errors.Is(err, common.ErrRequestBodyTooLarge) { + status = http.StatusRequestEntityTooLarge + } + return service.TaskErrorWrapperLocal(err, "read_request_body_failed", status) + } + defer func() { + _, _ = storage.Seek(0, io.SeekStart) + c.Request.Body = io.NopCloser(storage) + }() + + switch { + case strings.HasPrefix(contentType, "application/json"): + body, err := storage.Bytes() + if err != nil { + return service.TaskErrorWrapperLocal(err, "read_request_body_failed", http.StatusBadRequest) + } + if strings.TrimSpace(string(body)) != "" { + var payload map[string]any + if err := common.Unmarshal(body, &payload); err != nil { + return service.TaskErrorWrapperLocal(err, "invalid_json", http.StatusBadRequest) + } + if raw, ok := payload["model"].(string); ok && strings.TrimSpace(raw) != "" { + modelName = strings.TrimSpace(raw) + } + if raw, ok := payload["action"].(string); ok && strings.TrimSpace(raw) != "" { + action = strings.TrimSpace(raw) + } + } + case strings.Contains(contentType, gin.MIMEMultipartPOSTForm): + form, err := common.ParseMultipartFormReusable(c) + if err != nil { + return service.TaskErrorWrapperLocal(err, "invalid_multipart_form", http.StatusBadRequest) + } + if vals := form.Value["model"]; len(vals) > 0 && strings.TrimSpace(vals[0]) != "" { + modelName = strings.TrimSpace(vals[0]) + } + if vals := form.Value["action"]; len(vals) > 0 && strings.TrimSpace(vals[0]) != "" { + action = strings.TrimSpace(vals[0]) + } + } + + if modelName == "" && action != "" { + platform := constant.TaskPlatform(c.GetString("platform")) + modelName = service.CoverTaskActionToModelName(platform, action) + } + if modelName != "" { + info.OriginModelName = modelName + } + if action != "" { + info.Action = action + } + + info.PricingConfigSnapshotLoaded = true + info.PricingConfigSnapshot = ratio_setting.GetPricingConfig(info.OriginModelName) + info.RequireMatrixUsageBilling = info.PricingConfigSnapshot != nil && + info.PricingConfigSnapshot.BillingUnit == types.BillingUnitPer1MTokens + return nil +} + +func buildTaskBillingContext(info *relaycommon.RelayInfo) *model.TaskBillingContext { + bc := &model.TaskBillingContext{ + ModelPrice: info.PriceData.ModelPrice, + GroupRatio: info.PriceData.GroupRatioInfo.GroupRatio, + ModelRatio: info.PriceData.ModelRatio, + OtherRatios: info.PriceData.OtherRatios, + OriginModelName: info.OriginModelName, + PerCallBilling: common.StringsContains(constant.TaskPricePatches, info.OriginModelName), + } + if decision := info.PricingDecisionFrozen; decision != nil && decision.BillingMode == types.BillingModeMatrix { + bc.BillingMode = decision.BillingMode + bc.PricingSnapshot = types.CloneMapAny(decision.Snapshot) + bc.BillingUnit = decision.BillingUnit + bc.TokenUnitPriceUSD = decision.TokenUnitPriceUSD + bc.PerCallBilling = decision.PerCallBilling + bc.ModelPrice = decision.PriceUSD + bc.GroupRatio = decision.GroupRatioInfo.GroupRatio + bc.OtherRatios = types.CloneRatios(decision.OtherRatios) + } + return bc +} + func RelayTask(c *gin.Context) { relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil) if err != nil { @@ -503,6 +592,10 @@ func RelayTask(c *gin.Context) { respondTaskError(c, taskErr) return } + if taskErr := preloadTaskPricingConfig(c, relayInfo); taskErr != nil { + respondTaskError(c, taskErr) + return + } var result *relay.TaskSubmitResult var taskErr *dto.TaskError @@ -524,11 +617,9 @@ func RelayTask(c *gin.Context) { if lockedCh, ok := relayInfo.LockedChannel.(*model.Channel); ok && lockedCh != nil { channel = lockedCh - if retryParam.GetRetry() > 0 { - if setupErr := middleware.SetupContextForSelectedChannel(c, channel, relayInfo.OriginModelName); setupErr != nil { - taskErr = service.TaskErrorWrapperLocal(setupErr.Err, "setup_locked_channel_failed", http.StatusInternalServerError) - break - } + if setupErr := middleware.SetupContextForSelectedChannel(c, channel, relayInfo.OriginModelName); setupErr != nil { + taskErr = service.TaskErrorWrapperLocal(setupErr.Err, "setup_locked_channel_failed", http.StatusInternalServerError) + break } } else { var channelErr *types.NewAPIError @@ -587,14 +678,7 @@ func RelayTask(c *gin.Context) { task.PrivateData.BillingSource = relayInfo.BillingSource task.PrivateData.SubscriptionId = relayInfo.SubscriptionId task.PrivateData.TokenId = relayInfo.TokenId - task.PrivateData.BillingContext = &model.TaskBillingContext{ - ModelPrice: relayInfo.PriceData.ModelPrice, - GroupRatio: relayInfo.PriceData.GroupRatioInfo.GroupRatio, - ModelRatio: relayInfo.PriceData.ModelRatio, - OtherRatios: relayInfo.PriceData.OtherRatios, - OriginModelName: relayInfo.OriginModelName, - PerCallBilling: common.StringsContains(constant.TaskPricePatches, relayInfo.OriginModelName), - } + task.PrivateData.BillingContext = buildTaskBillingContext(relayInfo) task.Quota = result.Quota task.Data = result.TaskData task.Action = relayInfo.Action diff --git a/model/task.go b/model/task.go index 9844450..12072be 100644 --- a/model/task.go +++ b/model/task.go @@ -109,12 +109,16 @@ type TaskPrivateData struct { // TaskBillingContext 记录任务提交时的计费参数，以便轮询阶段可以重新计算额度。 type TaskBillingContext struct { - ModelPrice float64 `json:"model_price,omitempty"` // 模型单价 - GroupRatio float64 `json:"group_ratio,omitempty"` // 分组倍率 - ModelRatio float64 `json:"model_ratio,omitempty"` // 模型倍率 - OtherRatios map[string]float64 `json:"other_ratios,omitempty"` // 附加倍率（时长、分辨率等） - OriginModelName string `json:"origin_model_name,omitempty"` // 模型名称，必须为OriginModelName - PerCallBilling bool `json:"per_call_billing,omitempty"` // 按次计费：跳过轮询阶段的差额结算 + ModelPrice float64 `json:"model_price,omitempty"` // 模型单价 + GroupRatio float64 `json:"group_ratio,omitempty"` // 分组倍率 + ModelRatio float64 `json:"model_ratio,omitempty"` // 模型倍率 + OtherRatios map[string]float64 `json:"other_ratios,omitempty"` // 附加倍率（时长、分辨率等） + OriginModelName string `json:"origin_model_name,omitempty"` // 模型名称，必须为OriginModelName + PerCallBilling bool `json:"per_call_billing,omitempty"` // 按次计费：跳过轮询阶段的差额结算 + BillingMode string `json:"billing_mode,omitempty"` + PricingSnapshot map[string]any `json:"pricing_snapshot,omitempty"` + BillingUnit string `json:"billing_unit,omitempty"` + TokenUnitPriceUSD float64 `json:"token_unit_price_usd,omitempty"` } // GetUpstreamTaskID 获取上游真实 task ID（用于与 provider 通信） diff --git a/relay/relay_task.go b/relay/relay_task.go index 4040ec7..47ed96a 100644 --- a/relay/relay_task.go +++ b/relay/relay_task.go @@ -19,6 +19,7 @@ import ( relayconstant "github.com/QuantumNous/new-api/relay/constant" "github.com/QuantumNous/new-api/relay/helper" "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" ) @@ -87,28 +88,19 @@ func ResolveOriginTask(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskErr } info.LockedChannel = ch - if originTask.ChannelId != info.ChannelId { - key, _, newAPIError := ch.GetNextEnabledKey() - if newAPIError != nil { - return service.TaskErrorWrapper(newAPIError, "channel_no_available_key", newAPIError.StatusCode) - } - common.SetContextKey(c, constant.ContextKeyChannelKey, key) - common.SetContextKey(c, constant.ContextKeyChannelType, ch.Type) - common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, ch.GetBaseURL()) - common.SetContextKey(c, constant.ContextKeyChannelId, originTask.ChannelId) - - info.ChannelBaseUrl = ch.GetBaseURL() - info.ChannelId = originTask.ChannelId - info.ChannelType = ch.Type - info.ApiKey = key - } - // 提取 remix 参数（时长、分辨率 → OtherRatios） if info.Action == constant.TaskActionRemix { if originTask.PrivateData.BillingContext != nil { // 新的 remix 逻辑：直接从原始任务的 BillingContext 中提取 OtherRatios（如果存在） - for s, f := range originTask.PrivateData.BillingContext.OtherRatios { - info.PriceData.AddOtherRatio(s, f) + bc := originTask.PrivateData.BillingContext + info.OriginPricing = &types.OriginPricingSnapshot{ + BillingMode: bc.BillingMode, + BillingUnit: bc.BillingUnit, + PriceUSD: bc.ModelPrice, + TokenUnitPriceUSD: bc.TokenUnitPriceUSD, + Snapshot: types.CloneMapAny(bc.PricingSnapshot), + OtherRatios: types.CloneRatios(bc.OtherRatios), + PerCallBilling: bc.PerCallBilling, } } else { // 旧的 remix 逻辑：直接从 task data 解析 seconds 和 size（如果存在） @@ -120,13 +112,16 @@ func ResolveOriginTask(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskErr seconds = 4 } sizeStr, _ := taskData["size"].(string) - if info.PriceData.OtherRatios == nil { - info.PriceData.OtherRatios = map[string]float64{} + otherRatios := map[string]float64{ + "seconds": float64(seconds), + "size": 1, } - info.PriceData.OtherRatios["seconds"] = float64(seconds) - info.PriceData.OtherRatios["size"] = 1 if sizeStr == "1792x1024" || sizeStr == "1024x1792" { - info.PriceData.OtherRatios["size"] = 1.666667 + otherRatios["size"] = 1.666667 + } + info.OriginPricing = &types.OriginPricingSnapshot{ + OtherRatios: otherRatios, + PerCallBilling: false, } } } @@ -181,18 +176,22 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe return nil, pricingErr } info.PriceData = priceData + matrixOrOriginPricing := info.OriginPricing != nil || + (info.PricingDecisionFrozen != nil && info.PricingDecisionFrozen.BillingMode == types.BillingModeMatrix) // 5. 计费估算：让适配器根据用户请求提供 OtherRatios（时长、分辨率等） // 必须在 ModelPriceHelperPerCall 之后调用（它会重建 PriceData）。 // ResolveOriginTask 可能已在 remix 路径中预设了 OtherRatios，此处合并。 - if estimatedRatios := adaptor.EstimateBilling(c, info); len(estimatedRatios) > 0 { - for k, v := range estimatedRatios { - info.PriceData.AddOtherRatio(k, v) + if !matrixOrOriginPricing { + if estimatedRatios := adaptor.EstimateBilling(c, info); len(estimatedRatios) > 0 { + for k, v := range estimatedRatios { + info.PriceData.AddOtherRatio(k, v) + } } } // 6. 将 OtherRatios 应用到基础额度 - if !common.StringsContains(constant.TaskPricePatches, modelName) { + if !matrixOrOriginPricing && !common.StringsContains(constant.TaskPricePatches, modelName) { for _, ra := range info.PriceData.OtherRatios { if ra != 1.0 { info.PriceData.Quota = int(float64(info.PriceData.Quota) * ra) @@ -240,11 +239,13 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe // 11. 提交后计费调整：让适配器根据上游实际返回调整 OtherRatios finalQuota := info.PriceData.Quota - if adjustedRatios := adaptor.AdjustBillingOnSubmit(info, taskData); len(adjustedRatios) > 0 { - // 基于调整后的 ratios 重新计算 quota - finalQuota = recalcQuotaFromRatios(info, adjustedRatios) - info.PriceData.OtherRatios = adjustedRatios - info.PriceData.Quota = finalQuota + if !matrixOrOriginPricing { + if adjustedRatios := adaptor.AdjustBillingOnSubmit(info, taskData); len(adjustedRatios) > 0 { + // 基于调整后的 ratios 重新计算 quota + finalQuota = recalcQuotaFromRatios(info, adjustedRatios) + info.PriceData.OtherRatios = adjustedRatios + info.PriceData.Quota = finalQuota + } } return &TaskSubmitResult{ diff --git a/service/task_billing.go b/service/task_billing.go index 0da4cf4..87eef0b 100644 --- a/service/task_billing.go +++ b/service/task_billing.go @@ -118,6 +118,12 @@ func taskBillingOther(task *model.Task) map[string]interface{} { if bc := task.PrivateData.BillingContext; bc != nil { other["model_price"] = bc.ModelPrice other["group_ratio"] = bc.GroupRatio + if bc.BillingMode != "" { + other["billing_mode"] = bc.BillingMode + other["billing_unit"] = bc.BillingUnit + other["token_unit_price_usd"] = bc.TokenUnitPriceUSD + other["pricing_snapshot"] = bc.PricingSnapshot + } if len(bc.OtherRatios) > 0 { for k, v := range bc.OtherRatios { other[k] = v @@ -241,6 +247,42 @@ func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int // RecalculateTaskQuotaByTokens 根据实际 token 消耗重新计费（异步差额结算）。 // 当任务成功且返回了 totalTokens 时，根据模型倍率和分组倍率重新计算实际扣费额度， // 与预扣费的差额进行补扣或退还。支持钱包和订阅计费来源。 + +func RecalculateTaskQuotaWithOther(ctx context.Context, task *model.Task, actualQuota int, reason string, extra map[string]any) { + if len(extra) == 0 { + RecalculateTaskQuota(ctx, task, actualQuota, reason) + return + } + if task.PrivateData.BillingContext == nil { + task.PrivateData.BillingContext = &model.TaskBillingContext{} + } + if task.PrivateData.BillingContext.PricingSnapshot == nil { + task.PrivateData.BillingContext.PricingSnapshot = map[string]any{} + } + for k, v := range extra { + task.PrivateData.BillingContext.PricingSnapshot[k] = v + } + RecalculateTaskQuota(ctx, task, actualQuota, reason) +} + +func RecalculateMatrixUsageTaskQuota(ctx context.Context, task *model.Task, totalTokens int) { + if totalTokens <= 0 { + return + } + bc := task.PrivateData.BillingContext + if bc == nil { + return + } + actualQuota := int(float64(totalTokens) * bc.TokenUnitPriceUSD / 1_000_000 * common.QuotaPerUnit * bc.GroupRatio) + if actualQuota <= 0 && totalTokens > 0 && bc.TokenUnitPriceUSD > 0 && bc.GroupRatio > 0 { + actualQuota = 1 + } + reason := fmt.Sprintf("matrix usage billing: tokens=%d, token_unit_price_usd=%.6f, group_ratio=%.4f", totalTokens, bc.TokenUnitPriceUSD, bc.GroupRatio) + RecalculateTaskQuotaWithOther(ctx, task, actualQuota, reason, map[string]any{ + "total_tokens": totalTokens, + }) +} + func RecalculateTaskQuotaByTokens(ctx context.Context, task *model.Task, totalTokens int) { if totalTokens <= 0 { return diff --git a/service/task_polling.go b/service/task_polling.go index 3c5cab5..3186905 100644 --- a/service/task_polling.go +++ b/service/task_polling.go @@ -17,6 +17,7 @@ import ( "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/relay/channel/task/taskcommon" relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" "github.com/samber/lo" ) @@ -520,6 +521,24 @@ func truncateBase64(s string) string { // 2. taskResult.TotalTokens > 0 → 按 token 重算 // 3. 都不满足 → 保持预扣额度不变 func settleTaskBillingOnComplete(ctx context.Context, adaptor TaskPollingAdaptor, task *model.Task, taskResult *relaycommon.TaskInfo) { + if bc := task.PrivateData.BillingContext; bc != nil && bc.BillingMode == types.BillingModeMatrix { + _, billingUnit, _, err := types.ResolveMatrixPricingSnapshot(bc.BillingUnit, bc.PricingSnapshot) + if err != nil { + logger.LogError(ctx, fmt.Sprintf("task %s matrix pricing snapshot invalid: %s", task.TaskID, err.Error())) + return + } + if billingUnit == types.BillingUnitPerCall { + logger.LogInfo(ctx, fmt.Sprintf("task %s matrix per_call billing, skip settlement", task.TaskID)) + return + } + if billingUnit == types.BillingUnitPer1MTokens { + RecalculateMatrixUsageTaskQuota(ctx, task, taskResult.TotalTokens) + return + } + logger.LogError(ctx, fmt.Sprintf("task %s unknown matrix billing_unit: %s", task.TaskID, billingUnit)) + return + } + // 0. 按次计费的任务不做差额结算 if bc := task.PrivateData.BillingContext; bc != nil && bc.PerCallBilling { logger.LogInfo(ctx, fmt.Sprintf("任务 %s 按次计费，跳过差额结算", task.TaskID)) diff --git a/types/model_pricing.go b/types/model_pricing.go index 3d3a301..19cd893 100644 --- a/types/model_pricing.go +++ b/types/model_pricing.go @@ -1,6 +1,7 @@ package types import ( + "errors" "fmt" "github.com/QuantumNous/new-api/common" @@ -131,3 +132,63 @@ func NormalizePricingConfigDefaults(cfg *PricingConfig) { cfg.PreconsumeStrategy = PreconsumeStrategyExact } } + +func ResolveMatrixPricingSnapshot(billingUnit string, snapshot map[string]any) (price float64, resolvedBillingUnit string, pricingMode string, err error) { + if snapshot == nil { + return 0, "", "", errors.New("pricing snapshot is required") + } + rawPrice, ok := snapshot["price"] + if !ok { + return 0, "", "", errors.New("pricing snapshot missing price") + } + price, ok = snapshotNumber(rawPrice) + if !ok { + return 0, "", "", errors.New("pricing snapshot price must be numeric") + } + rawMode, ok := snapshot["pricing_mode"].(string) + if !ok || rawMode == "" { + return 0, "", "", errors.New("pricing snapshot pricing_mode must be string") + } + switch PricingMode(rawMode) { + case PricingModeMatched, PricingModeFallbackMax, PricingModeFallbackDefault: + pricingMode = rawMode + default: + return 0, "", "", fmt.Errorf("unsupported pricing_mode %q", rawMode) + } + + snapshotUnit, _ := snapshot["billing_unit"].(string) + if snapshotUnit != "" && snapshotUnit != BillingUnitPerCall && snapshotUnit != BillingUnitPer1MTokens { + return 0, "", "", fmt.Errorf("unsupported snapshot billing_unit %q", snapshotUnit) + } + if billingUnit != "" && billingUnit != BillingUnitPerCall && billingUnit != BillingUnitPer1MTokens { + return 0, "", "", fmt.Errorf("unsupported billing_unit %q", billingUnit) + } + switch { + case billingUnit != "" && snapshotUnit != "" && billingUnit != snapshotUnit: + return 0, "", "", fmt.Errorf("billing_unit %q does not match snapshot billing_unit %q", billingUnit, snapshotUnit) + case billingUnit != "": + resolvedBillingUnit = billingUnit + case snapshotUnit != "": + resolvedBillingUnit = snapshotUnit + default: + resolvedBillingUnit = BillingUnitPerCall + } + return price, resolvedBillingUnit, pricingMode, nil +} + +func snapshotNumber(v any) (float64, bool) { + switch n := v.(type) { + case float64: + return n, true + case float32: + return float64(n), true + case int: + return float64(n), true + case int64: + return float64(n), true + case int32: + return float64(n), true + default: + return 0, false + } +}