new-api/model/pricing.go

package model

import (
	"fmt"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/setting/ratio_setting"
	"github.com/QuantumNous/new-api/types"
)

type Pricing struct {
	ModelName              string                  `json:"model_name"`
	Description            string                  `json:"description,omitempty"`
	Icon                   string                  `json:"icon,omitempty"`
	Tags                   string                  `json:"tags,omitempty"`
	VendorID               int                     `json:"vendor_id,omitempty"`
	QuotaType              int                     `json:"quota_type"`
	ModelRatio             float64                 `json:"model_ratio"`
	ModelPrice             float64                 `json:"model_price"`
	OwnerBy                string                  `json:"owner_by"`
	CompletionRatio        float64                 `json:"completion_ratio"`
	CacheRatio             float64                 `json:"cache_ratio"`
	CacheCreationRatio     float64                 `json:"cache_creation_ratio"`
	EnableGroup            []string                `json:"enable_groups"`
	SupportedEndpointTypes []constant.EndpointType `json:"supported_endpoint_types"`
	PricingVersion         string                  `json:"pricing_version,omitempty"`
	Type                   int                     `json:"type"`
	PricingConfig          *types.PricingConfig    `json:"pricing_config,omitempty"`
}

type PricingVendor struct {
	ID          int    `json:"id"`
	Name        string `json:"name"`
	Description string `json:"description,omitempty"`
	Icon        string `json:"icon,omitempty"`
}

const (
	QuotaTypeByTokens = 0
	QuotaTypeByCall   = 1
)

var (
	pricingMap           []Pricing
	pricingByModel       map[string]*Pricing
	vendorsList          []PricingVendor
	supportedEndpointMap map[string]common.EndpointInfo
	lastGetPricingTime   time.Time
	updatePricingLock    sync.Mutex

	modelEnableGroups     = make(map[string][]string)
	modelQuotaTypeMap     = make(map[string]int)
	modelEnableGroupsLock = sync.RWMutex{}
)

var (
	modelSupportEndpointTypes = make(map[string][]constant.EndpointType)
	modelSupportEndpointsLock = sync.RWMutex{}
)

func SetTestPricing(pricing []Pricing) {
	updatePricingLock.Lock()
	defer updatePricingLock.Unlock()
	pricingMap = pricing
	idx := make(map[string]*Pricing, len(pricing))
	for i := range pricing {
		idx[pricing[i].ModelName] = &pricingMap[i]
	}
	pricingByModel = idx
	lastGetPricingTime = time.Now()
}

func GetPricing() []Pricing {
	if time.Since(lastGetPricingTime) > time.Minute || len(pricingMap) == 0 {
		updatePricingLock.Lock()
		defer updatePricingLock.Unlock()
		if time.Since(lastGetPricingTime) > time.Minute || len(pricingMap) == 0 {
			modelSupportEndpointsLock.Lock()
			defer modelSupportEndpointsLock.Unlock()
			updatePricing()
		}
	}
	return pricingMap
}

func GetPricingByModel(modelName string) *Pricing {
	GetPricing()
	modelEnableGroupsLock.RLock()
	p := pricingByModel[modelName]
	modelEnableGroupsLock.RUnlock()
	return p
}

func GetVendors() []PricingVendor {
	if time.Since(lastGetPricingTime) > time.Minute || len(pricingMap) == 0 {
		GetPricing()
	}
	return vendorsList
}

func GetModelSupportEndpointTypes(model string) []constant.EndpointType {
	if model == "" {
		return make([]constant.EndpointType, 0)
	}
	modelSupportEndpointsLock.RLock()
	defer modelSupportEndpointsLock.RUnlock()
	if endpoints, ok := modelSupportEndpointTypes[model]; ok {
		return endpoints
	}
	return make([]constant.EndpointType, 0)
}

func parseModelEndpointTypes(rawValue string) []string {
	trimmed := strings.TrimSpace(rawValue)
	if trimmed == "" {
		return nil
	}

	var rawObject map[string]interface{}
	if err := common.Unmarshal([]byte(trimmed), &rawObject); err == nil {
		endpoints := make([]string, 0, len(rawObject))
		for k, v := range rawObject {
			switch v.(type) {
			case string, map[string]interface{}:
				if !common.StringsContains(endpoints, k) {
					endpoints = append(endpoints, k)
				}
			}
		}
		if len(endpoints) == 0 {
			return nil
		}
		return endpoints
	}

	var rawList []string
	if err := common.Unmarshal([]byte(trimmed), &rawList); err == nil {
		endpoints := make([]string, 0, len(rawList))
		for _, endpoint := range rawList {
			endpoint = strings.TrimSpace(endpoint)
			if endpoint == "" || common.StringsContains(endpoints, endpoint) {
				continue
			}
			endpoints = append(endpoints, endpoint)
		}
		if len(endpoints) == 0 {
			return nil
		}
		return endpoints
	}

	return nil
}

func updatePricing() {
	enableAbilities, err := GetAllEnableAbilityWithChannels()
	if err != nil {
		common.SysLog(fmt.Sprintf("GetAllEnableAbilityWithChannels error: %v", err))
		return
	}

	var allMeta []Model
	_ = DB.Find(&allMeta).Error
	metaMap := make(map[string]*Model)
	prefixList := make([]*Model, 0)
	suffixList := make([]*Model, 0)
	containsList := make([]*Model, 0)
	for i := range allMeta {
		m := &allMeta[i]
		if m.NameRule == NameRuleExact {
			metaMap[m.ModelName] = m
			continue
		}
		switch m.NameRule {
		case NameRulePrefix:
			prefixList = append(prefixList, m)
		case NameRuleSuffix:
			suffixList = append(suffixList, m)
		case NameRuleContains:
			containsList = append(containsList, m)
		}
	}

	for _, m := range prefixList {
		for _, ability := range enableAbilities {
			if strings.HasPrefix(ability.Model, m.ModelName) {
				if _, exists := metaMap[ability.Model]; !exists {
					metaMap[ability.Model] = m
				}
			}
		}
	}
	for _, m := range suffixList {
		for _, ability := range enableAbilities {
			if strings.HasSuffix(ability.Model, m.ModelName) {
				if _, exists := metaMap[ability.Model]; !exists {
					metaMap[ability.Model] = m
				}
			}
		}
	}
	for _, m := range containsList {
		for _, ability := range enableAbilities {
			if strings.Contains(ability.Model, m.ModelName) {
				if _, exists := metaMap[ability.Model]; !exists {
					metaMap[ability.Model] = m
				}
			}
		}
	}

	var vendors []Vendor
	_ = DB.Find(&vendors).Error
	vendorMap := make(map[int]*Vendor, len(vendors))
	for i := range vendors {
		vendorMap[vendors[i].Id] = &vendors[i]
	}

	initDefaultVendorMapping(metaMap, vendorMap, enableAbilities)

	vendorOrderMap := make(map[int]int, len(vendorMap))
	for _, v := range vendorMap {
		vendorOrderMap[v.Id] = v.SortOrder
	}

	vendorsList = make([]PricingVendor, 0, len(vendorMap))
	for _, v := range vendorMap {
		vendorsList = append(vendorsList, PricingVendor{
			ID:          v.Id,
			Name:        v.Name,
			Description: v.Description,
			Icon:        v.Icon,
		})
	}
	sort.Slice(vendorsList, func(i, j int) bool {
		oi := vendorOrderMap[vendorsList[i].ID]
		oj := vendorOrderMap[vendorsList[j].ID]
		if oi != oj {
			return oi < oj
		}
		return vendorsList[i].ID < vendorsList[j].ID
	})

	modelGroupsMap := make(map[string]*types.Set[string])
	for _, ability := range enableAbilities {
		groups, ok := modelGroupsMap[ability.Model]
		if !ok {
			groups = types.NewSet[string]()
			modelGroupsMap[ability.Model] = groups
		}
		groups.Add(ability.Group)
	}

	modelSupportEndpointsStr := make(map[string][]string)
	for _, ability := range enableAbilities {
		endpoints := modelSupportEndpointsStr[ability.Model]
		channelTypes := common.GetEndpointTypesByChannelType(ability.ChannelType, ability.Model)
		for _, channelType := range channelTypes {
			if !common.StringsContains(endpoints, string(channelType)) {
				endpoints = append(endpoints, string(channelType))
			}
		}
		modelSupportEndpointsStr[ability.Model] = endpoints
	}

	for modelName, meta := range metaMap {
		if strings.TrimSpace(meta.Endpoints) == "" {
			continue
		}
		if endpoints := parseModelEndpointTypes(meta.Endpoints); len(endpoints) > 0 {
			modelSupportEndpointsStr[modelName] = endpoints
		}
	}

	modelSupportEndpointTypes = make(map[string][]constant.EndpointType)
	for model, endpoints := range modelSupportEndpointsStr {
		supportedEndpoints := make([]constant.EndpointType, 0, len(endpoints))
		for _, endpointStr := range endpoints {
			supportedEndpoints = append(supportedEndpoints, constant.EndpointType(endpointStr))
		}
		modelSupportEndpointTypes[model] = supportedEndpoints
	}

	supportedEndpointMap = make(map[string]common.EndpointInfo)
	for _, endpoints := range modelSupportEndpointTypes {
		for _, et := range endpoints {
			if info, ok := common.GetDefaultEndpointInfo(et); ok {
				if _, exists := supportedEndpointMap[string(et)]; !exists {
					supportedEndpointMap[string(et)] = info
				}
			}
		}
	}
	for _, meta := range metaMap {
		if strings.TrimSpace(meta.Endpoints) == "" {
			continue
		}
		var raw map[string]interface{}
		if err := common.Unmarshal([]byte(meta.Endpoints), &raw); err == nil {
			for k, v := range raw {
				switch val := v.(type) {
				case string:
					supportedEndpointMap[k] = common.EndpointInfo{Path: val, Method: "POST"}
				case map[string]interface{}:
					ep := common.EndpointInfo{Method: "POST"}
					if p, ok := val["path"].(string); ok {
						ep.Path = p
					}
					if m, ok := val["method"].(string); ok {
						ep.Method = strings.ToUpper(m)
					}
					supportedEndpointMap[k] = ep
				}
			}
		}
	}

	pricingMap = make([]Pricing, 0, len(modelGroupsMap))
	for model, groups := range modelGroupsMap {
		pricing := Pricing{
			ModelName:              model,
			EnableGroup:            groups.Items(),
			SupportedEndpointTypes: modelSupportEndpointTypes[model],
		}

		if meta, ok := metaMap[model]; ok {
			if meta.Status == 0 {
				continue
			}
			pricing.Description = meta.Description
			pricing.Icon = meta.Icon
			pricing.Tags = meta.Tags
			pricing.VendorID = meta.VendorID
			pricing.Type = meta.Type
		}

		applyGlobalDefault(&pricing, model)
		pricingMap = append(pricingMap, pricing)
	}

	sort.Slice(pricingMap, func(i, j int) bool {
		mi, okI := metaMap[pricingMap[i].ModelName]
		mj, okJ := metaMap[pricingMap[j].ModelName]
		si := 999999
		sj := 999999
		if okI {
			si = mi.SortOrder
		}
		if okJ {
			sj = mj.SortOrder
		}
		if si != sj {
			return si < sj
		}
		return pricingMap[i].ModelName < pricingMap[j].ModelName
	})

	if len(pricingMap) > 0 {
		pricingMap[0].PricingVersion = "82c4a357505fff6fee8462c3f7ec8a645bb95532669cb73b2cabee6a416ec24f"
	}

	modelEnableGroupsLock.Lock()
	modelEnableGroups = make(map[string][]string)
	modelQuotaTypeMap = make(map[string]int)
	byModel := make(map[string]*Pricing, len(pricingMap))
	for i := range pricingMap {
		p := &pricingMap[i]
		modelEnableGroups[p.ModelName] = p.EnableGroup
		modelQuotaTypeMap[p.ModelName] = p.QuotaType
		byModel[p.ModelName] = p
	}
	pricingByModel = byModel
	modelEnableGroupsLock.Unlock()

	lastGetPricingTime = time.Now()
}

func GetSupportedEndpointMap() map[string]common.EndpointInfo {
	return supportedEndpointMap
}

func applyGlobalDefault(pricing *Pricing, model string) {
	if applyMatrixPricingConfig(pricing, model) {
		return
	}
	modelPrice, findPrice := ratio_setting.GetModelPrice(model, false)
	if findPrice {
		pricing.ModelPrice = modelPrice
		pricing.ModelRatio = 0
		pricing.CompletionRatio = 0
		pricing.QuotaType = QuotaTypeByCall
	} else {
		modelRatio, _, _ := ratio_setting.GetModelRatio(model)
		pricing.ModelPrice = 0
		pricing.ModelRatio = modelRatio
		pricing.CompletionRatio = ratio_setting.GetCompletionRatio(model)
		pricing.QuotaType = QuotaTypeByTokens
	}
	pricing.CacheRatio, _ = ratio_setting.GetCacheRatio(model)
	pricing.CacheCreationRatio, _ = ratio_setting.GetCreateCacheRatio(model)
}

func applyMatrixPricingConfig(pricing *Pricing, model string) bool {
	cfg := ratio_setting.GetPricingConfig(model)
	if cfg == nil {
		return false
	}
	pricing.PricingConfig = cfg
	minPrice := minPricingConfigPrice(cfg)
	if cfg.BillingUnit == types.BillingUnitPer1MTokens {
		pricing.ModelPrice = minPrice
		pricing.ModelRatio = 0
		pricing.CompletionRatio = 0
		pricing.QuotaType = QuotaTypeByTokens
		return true
	}
	pricing.ModelPrice = minPrice
	pricing.ModelRatio = 0
	pricing.CompletionRatio = 0
	pricing.QuotaType = QuotaTypeByCall
	return true
}

func minPricingConfigPrice(cfg *types.PricingConfig) float64 {
	if cfg == nil || len(cfg.Table) == 0 {
		return 0
	}
	var min float64
	found := false
	for _, row := range cfg.Table {
		price, ok := pricingRowNumber(row["price"])
		if !ok {
			continue
		}
		if !found || price < min {
			min = price
			found = true
		}
	}
	return min
}

func pricingRowNumber(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
	}
}