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"` } 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 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] switch m.NameRule { case NameRuleExact: metaMap[m.ModelName] = m case NameRulePrefix: prefixList = append(prefixList, m) case NameRuleSuffix: suffixList = append(suffixList, m) case NameRuleContains: containsList = append(containsList, m) } } for _, m := range prefixList { for _, pricingModel := range enableAbilities { if strings.HasPrefix(pricingModel.Model, m.ModelName) { if _, exists := metaMap[pricingModel.Model]; !exists { metaMap[pricingModel.Model] = m } } } } for _, m := range suffixList { for _, pricingModel := range enableAbilities { if strings.HasSuffix(pricingModel.Model, m.ModelName) { if _, exists := metaMap[pricingModel.Model]; !exists { metaMap[pricingModel.Model] = m } } } } for _, m := range containsList { for _, pricingModel := range enableAbilities { if strings.Contains(pricingModel.Model, m.ModelName) { if _, exists := metaMap[pricingModel.Model]; !exists { metaMap[pricingModel.Model] = m } } } } var vendors []Vendor _ = DB.Find(&vendors).Error vendorMap := make(map[int]*Vendor) for i := range vendors { vendorMap[vendors[i].Id] = &vendors[i] } initDefaultVendorMapping(metaMap, vendorMap, enableAbilities) vendorOrderMap := make(map[int]int) 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 } var raw map[string]interface{} if err := common.Unmarshal([]byte(meta.Endpoints), &raw); err == nil { endpoints := make([]string, 0, len(raw)) for k, v := range raw { switch v.(type) { case string, map[string]interface{}: if !common.StringsContains(endpoints, k) { endpoints = append(endpoints, k) } } } if len(endpoints) > 0 { modelSupportEndpointsStr[modelName] = endpoints } } } modelSupportEndpointTypes = make(map[string][]constant.EndpointType) for model, endpoints := range modelSupportEndpointsStr { supportedEndpoints := make([]constant.EndpointType, 0) 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) { 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) }