new-api/setting/ratio_setting/model_pricing.go

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|video_input)$`)
	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 {
	// Keep in sync with matrixUsageCapabilities in relay/helper/matrix_usage_capability.go
	switch modelName {
	case "seedance-2", "doubao-video", "doubao-video-seedance",
		"doubao-seedance-2-0-260128", "doubao-seedance-2-0-fast-260128",
		"Doubao-Seedance-2.0", "Doubao-Seedance-2.0-fast",
		"kling-v1", "kling-v1-6", "kling-v2-6", "kling-v3",
		"kling-video-o1", "kling-v3-omni":
		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)
}