new-api/types/model_pricing.go

package types

import (
	"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
	}
}