feat(task-pricing): validate and cache model pricing rules

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":

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)
+}

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"))
+}