|
- 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 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) {
- 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)
- }
|