From b4301a53ae255f4bb03b291b489e6fb1f1ed6942 Mon Sep 17 00:00:00 2001 From: fengsilin Date: Wed, 17 Jun 2026 17:27:18 +0800 Subject: [PATCH] feat(task-pricing): validate and cache model pricing rules --- model/option.go | 3 + setting/ratio_setting/model_pricing.go | 377 ++++++++++++++++++++ setting/ratio_setting/model_pricing_test.go | 161 +++++++++ 3 files changed, 541 insertions(+) create mode 100644 setting/ratio_setting/model_pricing.go create mode 100644 setting/ratio_setting/model_pricing_test.go diff --git a/model/option.go b/model/option.go index b2e583c..056bfac 100644 --- a/model/option.go +++ b/model/option.go @@ -133,6 +133,7 @@ func InitOptionMap() { common.OptionMap["ModelRequestRateLimitGroup"] = setting.ModelRequestRateLimitGroup2JSONString() common.OptionMap["ModelRatio"] = ratio_setting.ModelRatio2JSONString() common.OptionMap["ModelPrice"] = ratio_setting.ModelPrice2JSONString() + common.OptionMap["ModelPricingRules"] = ratio_setting.ModelPricingRules2JSONString() common.OptionMap["CacheRatio"] = ratio_setting.CacheRatio2JSONString() common.OptionMap["CreateCacheRatio"] = ratio_setting.CreateCacheRatio2JSONString() common.OptionMap["GroupRatio"] = ratio_setting.GroupRatio2JSONString() @@ -511,6 +512,8 @@ func updateOptionMap(key string, value string) (err error) { err = ratio_setting.UpdateCompletionRatioByJSONString(value) case "ModelPrice": err = ratio_setting.UpdateModelPriceByJSONString(value) + case "ModelPricingRules": + err = ratio_setting.UpdateModelPricingRulesByJSONString(value) case "CacheRatio": err = ratio_setting.UpdateCacheRatioByJSONString(value) case "CreateCacheRatio": diff --git a/setting/ratio_setting/model_pricing.go b/setting/ratio_setting/model_pricing.go new file mode 100644 index 0000000..64a23d3 --- /dev/null +++ b/setting/ratio_setting/model_pricing.go @@ -0,0 +1,377 @@ +package ratio_setting + +import ( + "fmt" + "math" + "regexp" + "sort" + "strconv" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/types" +) + +const ModelPricingRulesOptionKey = "ModelPricingRules" + +var ( + pricingConfigMap = types.NewRWMap[string, *types.PricingConfig]() + dimensionKeyRE = regexp.MustCompile(`^[a-z][a-z0-9_]{0,31}$`) + requestSourceRE = regexp.MustCompile(`^request\.[a-zA-Z_][a-zA-Z0-9_.]*$`) + derivedSourceRE = regexp.MustCompile(`^derived\.(megapixels|aspect_ratio|input_mode)$`) + reservedDimKeys = map[string]bool{ + "price": true, + "source": true, + "__fallback__": true, + "price_formula": true, + "pricing_mode": true, + "fallback_strategy": true, + "schema_version": true, + "scope": true, + "billing_unit": true, + "preconsume_strategy": true, + } +) + +func GetPricingConfig(model string) *types.PricingConfig { + cfg, ok := pricingConfigMap.Get(model) + if !ok { + return nil + } + return types.ClonePricingConfig(cfg) +} + +func GetPricingConfigCopy() map[string]*types.PricingConfig { + src := pricingConfigMap.ReadAll() + dst := make(map[string]*types.PricingConfig, len(src)) + for k, v := range src { + dst[k] = types.ClonePricingConfig(v) + } + return dst +} + +func SetPricingConfig(model string, cfg *types.PricingConfig) error { + if strings.TrimSpace(model) == "" { + return fmt.Errorf("model name is required") + } + clone := types.ClonePricingConfig(cfg) + if err := ValidatePricingConfig(model, clone); err != nil { + return err + } + pricingConfigMap.Set(model, clone) + InvalidateExposedDataCache() + return nil +} + +func DeletePricingConfig(model string) error { + all := pricingConfigMap.ReadAll() + delete(all, model) + pricingConfigMap.Clear() + pricingConfigMap.AddAll(all) + InvalidateExposedDataCache() + return nil +} + +func ModelPricingRules2JSONString() string { + bytes, err := common.Marshal(GetPricingConfigCopy()) + if err != nil { + return "{}" + } + return string(bytes) +} + +func UpdateModelPricingRulesByJSONString(jsonStr string) error { + if strings.TrimSpace(jsonStr) == "" { + jsonStr = "{}" + } + raw := map[string]*types.PricingConfig{} + if err := common.UnmarshalJsonStr(jsonStr, &raw); err != nil { + return err + } + next := make(map[string]*types.PricingConfig, len(raw)) + for model, cfg := range raw { + clone := types.ClonePricingConfig(cfg) + if err := ValidatePricingConfig(model, clone); err != nil { + return fmt.Errorf("model %s: %w", model, err) + } + next[model] = clone + } + pricingConfigMap.Clear() + pricingConfigMap.AddAll(next) + InvalidateExposedDataCache() + return nil +} + +func ValidatePricingConfig(model string, cfg *types.PricingConfig) error { + if cfg == nil { + return fmt.Errorf("pricing config is required") + } + if cfg.SchemaVersion != types.PricingSchemaVersion { + return fmt.Errorf("schema_version must be %d", types.PricingSchemaVersion) + } + if cfg.Scope != types.PricingScopeModel { + return fmt.Errorf("scope must be %q", types.PricingScopeModel) + } + types.NormalizePricingConfigDefaults(cfg) + if !isAllowedBillingUnit(cfg.BillingUnit) { + return fmt.Errorf("billing_unit must be per_call or per_1m_tokens") + } + if !isAllowedPreconsumeStrategy(cfg.PreconsumeStrategy) { + return fmt.Errorf("preconsume_strategy must be exact or minimum") + } + if cfg.BillingUnit == types.BillingUnitPerCall && cfg.PreconsumeStrategy != types.PreconsumeStrategyExact { + return fmt.Errorf("billing_unit per_call requires preconsume_strategy exact") + } + if cfg.BillingUnit == types.BillingUnitPer1MTokens { + if cfg.PreconsumeStrategy != types.PreconsumeStrategyMinimum { + return fmt.Errorf("billing_unit per_1m_tokens requires preconsume_strategy minimum") + } + if !SupportsAnyMatrixUsageBillingModel(model) { + return fmt.Errorf("model %s does not support matrix usage billing", model) + } + } + if len(cfg.Dimensions) == 0 { + return fmt.Errorf("dimensions must not be empty") + } + if len(cfg.Table) == 0 { + return fmt.Errorf("table must not be empty") + } + if err := validateFallback(cfg.Fallback); err != nil { + return err + } + if err := validateDimensions(cfg.Dimensions); err != nil { + return err + } + if err := validateRows(cfg.Dimensions, cfg.Table); err != nil { + return err + } + return nil +} + +func SupportsAnyMatrixUsageBillingModel(modelName string) bool { + switch modelName { + case "seedance-2", "doubao-video", "doubao-video-seedance": + return true + default: + return false + } +} + +func isAllowedBillingUnit(unit string) bool { + return unit == types.BillingUnitPerCall || unit == types.BillingUnitPer1MTokens +} + +func isAllowedPreconsumeStrategy(strategy string) bool { + return strategy == types.PreconsumeStrategyExact || strategy == types.PreconsumeStrategyMinimum +} + +func validateFallback(fallback types.PricingFallback) error { + switch fallback.Strategy { + case types.PricingFallbackReject, types.PricingFallbackMax: + return nil + case types.PricingFallbackDefault: + if fallback.DefaultPrice < 0 { + return fmt.Errorf("fallback.default_price must be >= 0") + } + return nil + default: + return fmt.Errorf("fallback.strategy must be reject, max, or default") + } +} + +func validateDimensions(dimensions []types.PricingDimension) error { + seen := map[string]bool{} + for _, dim := range dimensions { + if !dimensionKeyRE.MatchString(dim.Key) { + return fmt.Errorf("dimension key %q is invalid", dim.Key) + } + if reservedDimKeys[dim.Key] { + return fmt.Errorf("dimension key %q is reserved", dim.Key) + } + if seen[dim.Key] { + return fmt.Errorf("dimension key %q is duplicated", dim.Key) + } + seen[dim.Key] = true + if dim.Type != "string" && dim.Type != "number" && dim.Type != "boolean" { + return fmt.Errorf("dimension %q type must be string, number, or boolean", dim.Key) + } + if !requestSourceRE.MatchString(dim.Source) && !derivedSourceRE.MatchString(dim.Source) { + return fmt.Errorf("dimension %q source is not allowed", dim.Key) + } + if dim.Optional { + if dim.Default == nil { + return fmt.Errorf("dimension %q default required when optional=true", dim.Key) + } + if !valueMatchesType(dim.Default, dim.Type) { + return fmt.Errorf("dimension %q default type must match %s", dim.Key, dim.Type) + } + } else if dim.Default != nil { + return fmt.Errorf("dimension %q default must be null when optional=false", dim.Key) + } + } + return nil +} + +func validateRows(dimensions []types.PricingDimension, rows []types.PricingRow) error { + for i, row := range rows { + if err := validateRowFields(dimensions, row); err != nil { + return fmt.Errorf("row %d: %w", i, err) + } + } + if err := validateUniqueRows(dimensions, rows); err != nil { + return err + } + if err := validateAmbiguousRows(dimensions, rows); err != nil { + return err + } + return nil +} + +func validateRowFields(dimensions []types.PricingDimension, row types.PricingRow) error { + allowed := map[string]bool{"price": true, "source": true} + for _, dim := range dimensions { + allowed[dim.Key] = true + } + for key := range row { + if !allowed[key] { + return fmt.Errorf("unexpected field %q", key) + } + } + for _, dim := range dimensions { + v, ok := row[dim.Key] + if !ok { + return fmt.Errorf("missing dimension %q", dim.Key) + } + if s, ok := v.(string); ok && s == "*" { + continue + } + if !valueMatchesType(v, dim.Type) { + return fmt.Errorf("dimension %q value type must match %s or wildcard", dim.Key, dim.Type) + } + } + price, ok := numberValue(row["price"]) + if !ok || price < 0 { + return fmt.Errorf("price must be >= 0") + } + source, ok := row["source"].(string) + if !ok || (source != types.PricingRowSourceManual && source != types.PricingRowSourceGenerated) { + return fmt.Errorf("source must be manual or generated") + } + return nil +} + +func validateUniqueRows(dimensions []types.PricingDimension, rows []types.PricingRow) error { + seen := map[string]bool{} + for _, row := range rows { + key := rowSignature(dimensions, row) + if seen[key] { + return fmt.Errorf("duplicate pricing row %s", key) + } + seen[key] = true + } + return nil +} + +func validateAmbiguousRows(dimensions []types.PricingDimension, rows []types.PricingRow) error { + for i := 0; i < len(rows); i++ { + for j := i + 1; j < len(rows); j++ { + if wildcardCount(dimensions, rows[i]) != wildcardCount(dimensions, rows[j]) { + continue + } + if rowsIntersect(dimensions, rows[i], rows[j]) { + return fmt.Errorf("ambiguous rows can both match with same wildcard count") + } + } + } + return nil +} + +func rowSignature(dimensions []types.PricingDimension, row types.PricingRow) string { + parts := make([]string, 0, len(dimensions)) + for _, dim := range dimensions { + parts = append(parts, dim.Key+"="+literalValue(row[dim.Key])) + } + sort.Strings(parts) + return strings.Join(parts, "|") +} + +func wildcardCount(dimensions []types.PricingDimension, row types.PricingRow) int { + count := 0 + for _, dim := range dimensions { + if s, ok := row[dim.Key].(string); ok && s == "*" { + count++ + } + } + return count +} + +func rowsIntersect(dimensions []types.PricingDimension, a, b types.PricingRow) bool { + for _, dim := range dimensions { + av := a[dim.Key] + bv := b[dim.Key] + as, awild := av.(string) + bs, bwild := bv.(string) + if awild && as == "*" || bwild && bs == "*" { + continue + } + if !sameLiteralValue(av, bv) { + return false + } + } + return true +} + +func valueMatchesType(v any, typ string) bool { + switch typ { + case "string": + _, ok := v.(string) + return ok + case "number": + _, ok := numberValue(v) + return ok + case "boolean": + _, ok := v.(bool) + return ok + default: + return false + } +} + +func numberValue(v any) (float64, bool) { + switch n := v.(type) { + case float64: + return n, !math.IsNaN(n) && !math.IsInf(n, 0) + case float32: + f := float64(n) + return f, !math.IsNaN(f) && !math.IsInf(f, 0) + case int: + return float64(n), true + case int64: + return float64(n), true + case int32: + return float64(n), true + case jsonNumber: + f, err := strconv.ParseFloat(string(n), 64) + return f, err == nil && !math.IsNaN(f) && !math.IsInf(f, 0) + default: + return 0, false + } +} + +type jsonNumber string + +func literalValue(v any) string { + if f, ok := numberValue(v); ok { + return strconv.FormatFloat(f, 'f', -1, 64) + } + return fmt.Sprintf("%v", v) +} + +func sameLiteralValue(a, b any) bool { + af, aok := numberValue(a) + bf, bok := numberValue(b) + if aok && bok { + return af == bf + } + return fmt.Sprintf("%v", a) == fmt.Sprintf("%v", b) +} diff --git a/setting/ratio_setting/model_pricing_test.go b/setting/ratio_setting/model_pricing_test.go new file mode 100644 index 0000000..e243e25 --- /dev/null +++ b/setting/ratio_setting/model_pricing_test.go @@ -0,0 +1,161 @@ +package ratio_setting + +import ( + "testing" + + "github.com/QuantumNous/new-api/types" + "github.com/stretchr/testify/require" +) + +func validPricingConfig() *types.PricingConfig { + return &types.PricingConfig{ + SchemaVersion: types.PricingSchemaVersion, + Scope: types.PricingScopeModel, + Dimensions: []types.PricingDimension{ + {Key: "resolution", Source: "request.resolution", Type: "string", Default: nil}, + {Key: "duration", Source: "request.duration", Type: "number", Default: nil}, + }, + Table: []types.PricingRow{ + {"resolution": "720P", "duration": float64(5), "price": float64(0.10), "source": types.PricingRowSourceManual}, + {"resolution": "720P", "duration": float64(10), "price": float64(0.20), "source": types.PricingRowSourceGenerated}, + }, + Fallback: types.PricingFallback{Strategy: types.PricingFallbackReject}, + } +} + +func TestValidatePricingConfigValidDefaults(t *testing.T) { + cfg := validPricingConfig() + require.NoError(t, ValidatePricingConfig("hailuo-test", cfg)) + require.Equal(t, types.BillingUnitPerCall, cfg.BillingUnit) + require.Equal(t, types.PreconsumeStrategyExact, cfg.PreconsumeStrategy) +} + +func TestValidatePricingConfigSchemaRules(t *testing.T) { + tests := []struct { + name string + mutate func(*types.PricingConfig) + wantErr string + }{ + {"schema_version", func(c *types.PricingConfig) { c.SchemaVersion = 2 }, "schema_version"}, + {"scope", func(c *types.PricingConfig) { c.Scope = "channel" }, "scope"}, + {"dimensions", func(c *types.PricingConfig) { c.Dimensions = nil }, "dimensions"}, + {"table", func(c *types.PricingConfig) { c.Table = nil }, "table"}, + {"fallback", func(c *types.PricingConfig) { c.Fallback.Strategy = "nearest" }, "fallback.strategy"}, + {"default_price", func(c *types.PricingConfig) { + c.Fallback = types.PricingFallback{Strategy: types.PricingFallbackDefault, DefaultPrice: -1} + }, "default_price"}, + {"per_call_strategy", func(c *types.PricingConfig) { + c.BillingUnit = types.BillingUnitPerCall + c.PreconsumeStrategy = types.PreconsumeStrategyMinimum + }, "per_call"}, + {"usage_strategy", func(c *types.PricingConfig) { + c.BillingUnit = types.BillingUnitPer1MTokens + c.PreconsumeStrategy = types.PreconsumeStrategyExact + }, "per_1m_tokens"}, + {"usage_allowlist", func(c *types.PricingConfig) { + c.BillingUnit = types.BillingUnitPer1MTokens + c.PreconsumeStrategy = types.PreconsumeStrategyMinimum + }, "does not support"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := validPricingConfig() + tt.mutate(cfg) + err := ValidatePricingConfig("hailuo-test", cfg) + require.Error(t, err) + require.Contains(t, err.Error(), tt.wantErr) + }) + } +} + +func TestValidatePricingConfigUsageAllowlist(t *testing.T) { + cfg := validPricingConfig() + cfg.BillingUnit = types.BillingUnitPer1MTokens + cfg.PreconsumeStrategy = types.PreconsumeStrategyMinimum + require.NoError(t, ValidatePricingConfig("seedance-2", cfg)) +} + +func TestValidatePricingConfigDimensionRules(t *testing.T) { + tests := []struct { + name string + dim types.PricingDimension + wantErr string + }{ + {"bad_key", types.PricingDimension{Key: "Bad", Source: "request.x", Type: "string", Default: nil}, "invalid"}, + {"reserved", types.PricingDimension{Key: "price", Source: "request.x", Type: "string", Default: nil}, "reserved"}, + {"bad_type", types.PricingDimension{Key: "quality", Source: "request.x", Type: "object", Default: nil}, "type"}, + {"bad_source", types.PricingDimension{Key: "quality", Source: "response.usage.total_tokens", Type: "number", Default: nil}, "source"}, + {"optional_default_missing", types.PricingDimension{Key: "quality", Source: "request.quality", Type: "string", Optional: true, Default: nil}, "default required"}, + {"optional_default_type", types.PricingDimension{Key: "quality", Source: "request.quality", Type: "number", Optional: true, Default: "fast"}, "default type"}, + {"required_default", types.PricingDimension{Key: "quality", Source: "request.quality", Type: "string", Default: "fast"}, "default must be null"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := validPricingConfig() + cfg.Dimensions = []types.PricingDimension{tt.dim} + cfg.Table = []types.PricingRow{{tt.dim.Key: "x", "price": float64(1), "source": types.PricingRowSourceManual}} + if tt.dim.Type == "number" { + cfg.Table[0][tt.dim.Key] = float64(1) + } + err := ValidatePricingConfig("hailuo-test", cfg) + require.Error(t, err) + require.Contains(t, err.Error(), tt.wantErr) + }) + } +} + +func TestValidatePricingConfigRowRules(t *testing.T) { + tests := []struct { + name string + mutate func(*types.PricingConfig) + wantErr string + }{ + {"missing_dim", func(c *types.PricingConfig) { delete(c.Table[0], "duration") }, "missing dimension"}, + {"extra_field", func(c *types.PricingConfig) { c.Table[0]["price_formula"] = "x" }, "unexpected field"}, + {"bad_dim_type", func(c *types.PricingConfig) { c.Table[0]["duration"] = "five" }, "value type"}, + {"bad_price", func(c *types.PricingConfig) { c.Table[0]["price"] = -1.0 }, "price"}, + {"bad_source", func(c *types.PricingConfig) { c.Table[0]["source"] = "api" }, "source"}, + {"duplicate", func(c *types.PricingConfig) { c.Table = append(c.Table, types.PricingRow(types.CloneMapAny(c.Table[0]))) }, "duplicate"}, + {"ambiguous", func(c *types.PricingConfig) { + c.Dimensions = []types.PricingDimension{ + {Key: "res", Source: "request.res", Type: "string", Default: nil}, + {Key: "dur", Source: "request.dur", Type: "number", Default: nil}, + } + c.Table = []types.PricingRow{ + {"res": "*", "dur": float64(5), "price": float64(0.10), "source": types.PricingRowSourceManual}, + {"res": "720P", "dur": "*", "price": float64(0.12), "source": types.PricingRowSourceManual}, + } + }, "ambiguous"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := validPricingConfig() + tt.mutate(cfg) + err := ValidatePricingConfig("hailuo-test", cfg) + require.Error(t, err) + require.Contains(t, err.Error(), tt.wantErr) + }) + } +} + +func TestModelPricingRulesRoundTripAndClone(t *testing.T) { + require.NoError(t, UpdateModelPricingRulesByJSONString(`{}`)) + cfg := validPricingConfig() + require.NoError(t, SetPricingConfig("hailuo-test", cfg)) + + cfg.Table[0]["price"] = float64(999) + got := GetPricingConfig("hailuo-test") + require.NotNil(t, got) + require.Equal(t, float64(0.10), got.Table[0]["price"]) + + jsonStr := ModelPricingRules2JSONString() + require.Contains(t, jsonStr, "hailuo-test") + require.NoError(t, UpdateModelPricingRulesByJSONString(jsonStr)) + require.NotNil(t, GetPricingConfig("hailuo-test")) + + require.NoError(t, DeletePricingConfig("hailuo-test")) + require.Nil(t, GetPricingConfig("hailuo-test")) +}