new-api/types/model_pricing.go

package types

import (
	"errors"
	"fmt"

	"github.com/QuantumNous/new-api/common"
)

const (
	PricingSchemaVersion = 1

	PricingScopeModel = "model"

	BillingUnitPerCall     = "per_call"
	BillingUnitPer1MTokens = "per_1m_tokens"

	PreconsumeStrategyExact   = "exact"
	PreconsumeStrategyMinimum = "minimum"

	PricingFallbackReject  = "reject"
	PricingFallbackMax     = "max"
	PricingFallbackDefault = "default"

	PricingRowSourceManual    = "manual"
	PricingRowSourceGenerated = "generated"

	BillingModeMatrix = "matrix"
)

type PricingConfig struct {
	SchemaVersion      int                `json:"schema_version"`
	Scope              string             `json:"scope"`
	BillingUnit        string             `json:"billing_unit,omitempty"`
	PreconsumeStrategy string             `json:"preconsume_strategy,omitempty"`
	Dimensions         []PricingDimension `json:"dimensions"`
	Table              []PricingRow       `json:"table"`
	Fallback           PricingFallback    `json:"fallback"`
	GeneratorSnapshot  map[string]any     `json:"generator_snapshot,omitempty"`
}

type PricingDimension struct {
	Key      string `json:"key"`
	Source   string `json:"source"`
	Type     string `json:"type"`
	Optional bool   `json:"optional,omitempty"`
	Default  any    `json:"default"`
}

type PricingRow map[string]any

type PricingFallback struct {
	Strategy     string  `json:"strategy"`
	DefaultPrice float64 `json:"default_price,omitempty"`
}

type OriginPricingSnapshot struct {
	BillingMode       string
	BillingUnit       string
	PriceUSD          float64
	TokenUnitPriceUSD float64
	Snapshot          map[string]any
	OtherRatios       map[string]float64
	PerCallBilling    bool
}

func ClonePricingConfig(src *PricingConfig) *PricingConfig {
	if src == nil {
		return nil
	}
	return &PricingConfig{
		SchemaVersion:      src.SchemaVersion,
		Scope:              src.Scope,
		BillingUnit:        src.BillingUnit,
		PreconsumeStrategy: src.PreconsumeStrategy,
		Dimensions:         append([]PricingDimension(nil), src.Dimensions...),
		Table:              ClonePricingRows(src.Table),
		Fallback:           src.Fallback,
		GeneratorSnapshot:  CloneMapAny(src.GeneratorSnapshot),
	}
}

func ClonePricingRows(src []PricingRow) []PricingRow {
	if src == nil {
		return nil
	}
	dst := make([]PricingRow, len(src))
	for i := range src {
		dst[i] = PricingRow(CloneMapAny(map[string]any(src[i])))
	}
	return dst
}

func CloneMapAny(src map[string]any) map[string]any {
	if src == nil {
		return nil
	}
	bytes, err := common.Marshal(src)
	if err != nil {
		dst := make(map[string]any, len(src))
		for k, v := range src {
			dst[k] = v
		}
		return dst
	}
	var dst map[string]any
	if err := common.Unmarshal(bytes, &dst); err != nil {
		panic(fmt.Sprintf("clone map failed after successful marshal: %v", err))
	}
	return dst
}

func CloneRatios(src map[string]float64) map[string]float64 {
	if src == nil {
		return nil
	}
	dst := make(map[string]float64, len(src))
	for k, v := range src {
		dst[k] = v
	}
	return dst
}

func NormalizePricingConfigDefaults(cfg *PricingConfig) {
	if cfg == nil {
		return
	}
	if cfg.BillingUnit == "" {
		cfg.BillingUnit = BillingUnitPerCall
	}
	if cfg.PreconsumeStrategy == "" {
		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
	}
}