diff --git a/model/model_extra.go b/model/model_extra.go index 3724346e..6ade6ff0 100644 --- a/model/model_extra.go +++ b/model/model_extra.go @@ -1,24 +1,34 @@ package model // GetModelEnableGroups 返回指定模型名称可用的用户分组列表。 -// 复用缓存的定价映射，避免额外的数据库查询。 +// 使用在 updatePricing() 中维护的缓存映射，O(1) 读取，适合高并发场景。 func GetModelEnableGroups(modelName string) []string { - for _, p := range GetPricing() { - if p.ModelName == modelName { - return p.EnableGroup - } + // 确保缓存最新 + GetPricing() + + if modelName == "" { + return make([]string, 0) } - return make([]string, 0) + + modelEnableGroupsLock.RLock() + groups, ok := modelEnableGroups[modelName] + modelEnableGroupsLock.RUnlock() + if !ok { + return make([]string, 0) + } + return groups } // GetModelQuotaType 返回指定模型的计费类型（quota_type）。 -// 复用缓存的定价映射，避免额外数据库查询。 -// 如果未找到对应模型，默认返回 0。 +// 同样使用缓存映射，避免每次遍历定价切片。 func GetModelQuotaType(modelName string) int { - for _, p := range GetPricing() { - if p.ModelName == modelName { - return p.QuotaType - } + GetPricing() + + modelEnableGroupsLock.RLock() + quota, ok := modelQuotaTypeMap[modelName] + modelEnableGroupsLock.RUnlock() + if !ok { + return 0 } - return 0 + return quota } diff --git a/model/model_meta.go b/model/model_meta.go index 4faf7a84..f90d4831 100644 --- a/model/model_meta.go +++ b/model/model_meta.go @@ -36,7 +36,7 @@ type BoundChannel struct { type Model struct { Id int `json:"id"` - ModelName string `json:"model_name" gorm:"uniqueIndex;size:128;not null"` + ModelName string `json:"model_name" gorm:"size:128;not null;uniqueIndex:uk_model_name,where:deleted_at IS NULL"` Description string `json:"description,omitempty" gorm:"type:text"` Tags string `json:"tags,omitempty" gorm:"type:varchar(255)"` VendorID int `json:"vendor_id,omitempty" gorm:"index"` diff --git a/model/prefill_group.go b/model/prefill_group.go index 7a3a6673..6ebe3b04 100644 --- a/model/prefill_group.go +++ b/model/prefill_group.go @@ -4,6 +4,7 @@ import ( "one-api/common" "gorm.io/datatypes" + "gorm.io/gorm" ) // PrefillGroup 用于存储可复用的“组”信息，例如模型组、标签组、端点组等。 @@ -15,12 +16,13 @@ import ( type PrefillGroup struct { Id int `json:"id"` - Name string `json:"name" gorm:"uniqueIndex;size:64;not null"` + Name string `json:"name" gorm:"size:64;not null;uniqueIndex:uk_prefill_name,where:deleted_at IS NULL"` Type string `json:"type" gorm:"size:32;index;not null"` Items datatypes.JSON `json:"items" gorm:"type:json"` Description string `json:"description,omitempty" gorm:"type:varchar(255)"` CreatedTime int64 `json:"created_time" gorm:"bigint"` UpdatedTime int64 `json:"updated_time" gorm:"bigint"` + DeletedAt gorm.DeletedAt `json:"-" gorm:"index"` } // Insert 新建组 diff --git a/model/pricing.go b/model/pricing.go index 53fd0e89..5f3939da 100644 --- a/model/pricing.go +++ b/model/pricing.go @@ -37,6 +37,11 @@ var ( vendorsList []PricingVendor lastGetPricingTime time.Time updatePricingLock sync.Mutex + + // 缓存映射：模型名 -> 启用分组 / 计费类型 + modelEnableGroups = make(map[string][]string) + modelQuotaTypeMap = make(map[string]int) + modelEnableGroupsLock = sync.RWMutex{} ) var ( @@ -193,30 +198,41 @@ func updatePricing() { } pricingMap = make([]Pricing, 0) - for model, groups := range modelGroupsMap { - pricing := Pricing{ - ModelName: model, - EnableGroup: groups.Items(), - SupportedEndpointTypes: modelSupportEndpointTypes[model], - } + for model, groups := range modelGroupsMap { + pricing := Pricing{ + ModelName: model, + EnableGroup: groups.Items(), + SupportedEndpointTypes: modelSupportEndpointTypes[model], + } - // 补充模型元数据（描述、标签、供应商等） - if meta, ok := metaMap[model]; ok { - pricing.Description = meta.Description - pricing.Tags = meta.Tags - pricing.VendorID = meta.VendorID - } - modelPrice, findPrice := ratio_setting.GetModelPrice(model, false) - if findPrice { - pricing.ModelPrice = modelPrice - pricing.QuotaType = 1 - } else { - modelRatio, _, _ := ratio_setting.GetModelRatio(model) - pricing.ModelRatio = modelRatio - pricing.CompletionRatio = ratio_setting.GetCompletionRatio(model) - pricing.QuotaType = 0 - } - pricingMap = append(pricingMap, pricing) - } - lastGetPricingTime = time.Now() + // 补充模型元数据（描述、标签、供应商等） + if meta, ok := metaMap[model]; ok { + pricing.Description = meta.Description + pricing.Tags = meta.Tags + pricing.VendorID = meta.VendorID + } + modelPrice, findPrice := ratio_setting.GetModelPrice(model, false) + if findPrice { + pricing.ModelPrice = modelPrice + pricing.QuotaType = 1 + } else { + modelRatio, _, _ := ratio_setting.GetModelRatio(model) + pricing.ModelRatio = modelRatio + pricing.CompletionRatio = ratio_setting.GetCompletionRatio(model) + pricing.QuotaType = 0 + } + pricingMap = append(pricingMap, pricing) + } + + // 刷新缓存映射，供高并发快速查询 + modelEnableGroupsLock.Lock() + modelEnableGroups = make(map[string][]string) + modelQuotaTypeMap = make(map[string]int) + for _, p := range pricingMap { + modelEnableGroups[p.ModelName] = p.EnableGroup + modelQuotaTypeMap[p.ModelName] = p.QuotaType + } + modelEnableGroupsLock.Unlock() + + lastGetPricingTime = time.Now() } diff --git a/model/vendor_meta.go b/model/vendor_meta.go index 1dcec351..76bda1f0 100644 --- a/model/vendor_meta.go +++ b/model/vendor_meta.go @@ -14,7 +14,7 @@ import ( type Vendor struct { Id int `json:"id"` - Name string `json:"name" gorm:"uniqueIndex;size:128;not null"` + Name string `json:"name" gorm:"size:128;not null;uniqueIndex:uk_vendor_name,where:deleted_at IS NULL"` Description string `json:"description,omitempty" gorm:"type:text"` Icon string `json:"icon,omitempty" gorm:"type:varchar(128)"` Status int `json:"status" gorm:"default:1"`