# Conflicts: # model/pricing_test.gomaster
| @@ -8,6 +8,8 @@ build | |||||
| *.db-journal | *.db-journal | ||||
| logs | logs | ||||
| web/dist | web/dist | ||||
| !web/dist/ | |||||
| !web/dist/index.html | |||||
| .env | .env | ||||
| one-api | one-api | ||||
| new-api | new-api | ||||
| @@ -177,15 +177,14 @@ func formatDiscount(ratio float64) string { | |||||
| if ratio <= 0 { | if ratio <= 0 { | ||||
| return "免费" | return "免费" | ||||
| } | } | ||||
| discount := int(math.Round(ratio * 10)) | |||||
| if discount >= 10 { | |||||
| rawDiscount := math.Round(ratio*100) / 10 | |||||
| if rawDiscount >= 10 { | |||||
| return "" | return "" | ||||
| } | } | ||||
| remainder := int(math.Round(ratio*100)) % 10 | |||||
| if remainder == 0 { | |||||
| return fmt.Sprintf("%d折", discount) | |||||
| if rawDiscount == math.Trunc(rawDiscount) { | |||||
| return fmt.Sprintf("%.0f折", rawDiscount) | |||||
| } | } | ||||
| return fmt.Sprintf("%.1f折", ratio*10) | |||||
| return fmt.Sprintf("%.1f折", rawDiscount) | |||||
| } | } | ||||
| func ResetModelRatio(c *gin.Context) { | func ResetModelRatio(c *gin.Context) { | ||||
| @@ -1,330 +0,0 @@ | |||||
| package model | |||||
| import ( | |||||
| "fmt" | |||||
| "strconv" | |||||
| "strings" | |||||
| "sync" | |||||
| "github.com/QuantumNous/new-api/common" | |||||
| "github.com/QuantumNous/new-api/setting/ratio_setting" | |||||
| "gorm.io/gorm" | |||||
| "gorm.io/gorm/clause" | |||||
| ) | |||||
| // 渠道定价缓存 | |||||
| var ( | |||||
| channelPricingCache = make(map[string]*ChannelPricing) // key: "modelName:channelId" | |||||
| channelPricingCacheLock sync.RWMutex | |||||
| ) | |||||
| // QuotaType 计费类型 | |||||
| const ( | |||||
| QuotaTypeByTokens = 0 // 按量计费 | |||||
| QuotaTypeByCall = 1 // 按次计费 | |||||
| ) | |||||
| // ChannelPricing 渠道定价表 | |||||
| // 支持同一模型在不同渠道设置不同价格 | |||||
| type ChannelPricing struct { | |||||
| Id int `json:"id" gorm:"primaryKey"` | |||||
| ModelName string `json:"model_name" gorm:"size:128;not null;uniqueIndex:idx_model_channel,priority:1"` | |||||
| ChannelId int `json:"channel_id" gorm:"not null;uniqueIndex:idx_model_channel,priority:2;index"` | |||||
| QuotaType int `json:"quota_type" gorm:"default:0"` // 0=按量, 1=按次 | |||||
| ModelRatio float64 `json:"model_ratio" gorm:"default:0"` | |||||
| CompletionRatio float64 `json:"completion_ratio" gorm:"default:0"` | |||||
| ModelPrice float64 `json:"model_price" gorm:"default:0"` | |||||
| TagIds string `json:"tag_ids" gorm:"type:varchar(255)"` // 逗号分隔的标签ID | |||||
| CreatedTime int64 `json:"created_time" gorm:"bigint"` | |||||
| UpdatedTime int64 `json:"updated_time" gorm:"bigint"` | |||||
| DeletedAt gorm.DeletedAt `json:"-" gorm:"index"` | |||||
| // === 新增字段(0 = 未设置,回退全局值) === | |||||
| CacheRatio float64 `json:"cache_ratio" gorm:"default:0"` | |||||
| CacheCreationRatio float64 `json:"cache_creation_ratio" gorm:"default:0"` | |||||
| ImageRatio float64 `json:"image_ratio" gorm:"default:0"` | |||||
| AudioRatio float64 `json:"audio_ratio" gorm:"default:0"` | |||||
| AudioCompletionRatio float64 `json:"audio_completion_ratio" gorm:"default:0"` | |||||
| IsDefault bool `json:"is_default" gorm:"default:false;index"` | |||||
| } | |||||
| // setCache 写穿透缓存 | |||||
| func setCache(key string, cp *ChannelPricing) { | |||||
| channelPricingCacheLock.Lock() | |||||
| channelPricingCache[key] = cp | |||||
| channelPricingCacheLock.Unlock() | |||||
| } | |||||
| func removeCache(key string) { | |||||
| channelPricingCacheLock.Lock() | |||||
| delete(channelPricingCache, key) | |||||
| channelPricingCacheLock.Unlock() | |||||
| } | |||||
| // ApplyFields 批量设置定价字段(消除 controller 层的重复赋值) | |||||
| func (cp *ChannelPricing) ApplyFields(quotaType int, modelRatio, completionRatio, modelPrice float64, tagIds string, cacheRatio, cacheCreationRatio, imageRatio, audioRatio, audioCompletionRatio float64) { | |||||
| cp.QuotaType = quotaType | |||||
| cp.ModelRatio = modelRatio | |||||
| cp.CompletionRatio = completionRatio | |||||
| cp.ModelPrice = modelPrice | |||||
| cp.TagIds = tagIds | |||||
| cp.CacheRatio = cacheRatio | |||||
| cp.CacheCreationRatio = cacheCreationRatio | |||||
| cp.ImageRatio = imageRatio | |||||
| cp.AudioRatio = audioRatio | |||||
| cp.AudioCompletionRatio = audioCompletionRatio | |||||
| } | |||||
| func (cp *ChannelPricing) Insert() error { | |||||
| now := common.GetTimestamp() | |||||
| cp.CreatedTime = now | |||||
| cp.UpdatedTime = now | |||||
| err := DB.Create(cp).Error | |||||
| if err == nil { | |||||
| setCache(getChannelPricingCacheKey(cp.ModelName, cp.ChannelId), cp) | |||||
| } | |||||
| return err | |||||
| } | |||||
| func (cp *ChannelPricing) Update() error { | |||||
| cp.UpdatedTime = common.GetTimestamp() | |||||
| err := DB.Model(&ChannelPricing{}).Where("id = ?", cp.Id). | |||||
| Select("quota_type", "model_ratio", "completion_ratio", "model_price", | |||||
| "tag_ids", "cache_ratio", "cache_creation_ratio", "image_ratio", | |||||
| "audio_ratio", "audio_completion_ratio", "is_default", "updated_time"). | |||||
| Updates(cp).Error | |||||
| if err == nil { | |||||
| setCache(getChannelPricingCacheKey(cp.ModelName, cp.ChannelId), cp) | |||||
| } | |||||
| return err | |||||
| } | |||||
| func (cp *ChannelPricing) Delete() error { | |||||
| var existing ChannelPricing | |||||
| if err := DB.First(&existing, cp.Id).Error; err != nil { | |||||
| return err | |||||
| } | |||||
| err := DB.Delete(cp).Error | |||||
| if err == nil { | |||||
| removeCache(getChannelPricingCacheKey(existing.ModelName, existing.ChannelId)) | |||||
| } | |||||
| return err | |||||
| } | |||||
| // GetChannelPricing 获取指定模型在指定渠道的定价 | |||||
| func GetChannelPricing(modelName string, channelId int) (*ChannelPricing, error) { | |||||
| var cp ChannelPricing | |||||
| err := DB.Where("model_name = ? AND channel_id = ?", modelName, channelId).First(&cp).Error | |||||
| if err != nil { | |||||
| return nil, err | |||||
| } | |||||
| return &cp, nil | |||||
| } | |||||
| // GetChannelPricingByModel 获取指定模型的所有渠道定价 | |||||
| func GetChannelPricingByModel(modelName string) ([]*ChannelPricing, error) { | |||||
| var list []*ChannelPricing | |||||
| err := DB.Where("model_name = ?", modelName).Find(&list).Error | |||||
| return list, err | |||||
| } | |||||
| // GetAllChannelPricing 获取所有渠道定价(分页) | |||||
| func GetAllChannelPricing(offset int, limit int) ([]*ChannelPricing, int64, error) { | |||||
| var list []*ChannelPricing | |||||
| var total int64 | |||||
| if err := DB.Model(&ChannelPricing{}).Count(&total).Error; err != nil { | |||||
| return nil, 0, err | |||||
| } | |||||
| err := DB.Order("id DESC").Offset(offset).Limit(limit).Find(&list).Error | |||||
| return list, total, err | |||||
| } | |||||
| // BatchUpsertChannelPricing 批量创建或更新渠道定价 | |||||
| func BatchUpsertChannelPricing(pricings []*ChannelPricing) error { | |||||
| if len(pricings) == 0 { | |||||
| return nil | |||||
| } | |||||
| now := common.GetTimestamp() | |||||
| for _, cp := range pricings { | |||||
| cp.UpdatedTime = now | |||||
| // 仅在 CreatedTime 为空时设置(新记录) | |||||
| if cp.CreatedTime == 0 { | |||||
| cp.CreatedTime = now | |||||
| } | |||||
| } | |||||
| // 使用 GORM 的 OnConflict 实现 upsert | |||||
| // 唯一索引为 idx_model_channel (model_name, channel_id) | |||||
| err := DB.Clauses(clause.OnConflict{ | |||||
| Columns: []clause.Column{ | |||||
| {Name: "model_name"}, | |||||
| {Name: "channel_id"}, | |||||
| }, | |||||
| DoUpdates: clause.AssignmentColumns([]string{ | |||||
| "quota_type", | |||||
| "model_ratio", | |||||
| "completion_ratio", | |||||
| "model_price", | |||||
| "tag_ids", | |||||
| "cache_ratio", | |||||
| "cache_creation_ratio", | |||||
| "image_ratio", | |||||
| "audio_ratio", | |||||
| "audio_completion_ratio", | |||||
| "updated_time", | |||||
| }), | |||||
| }).Create(&pricings).Error | |||||
| if err == nil { | |||||
| for _, cp := range pricings { | |||||
| setCache(getChannelPricingCacheKey(cp.ModelName, cp.ChannelId), cp) | |||||
| } | |||||
| } | |||||
| return err | |||||
| } | |||||
| // getChannelPricingCacheKey 生成缓存键 | |||||
| func getChannelPricingCacheKey(modelName string, channelId int) string { | |||||
| return fmt.Sprintf("%s:%d", modelName, channelId) | |||||
| } | |||||
| // GetEffectivePricing 获取有效定价(纯内存查找) | |||||
| func GetEffectivePricing(modelName string, channelId int) (*ChannelPricing, bool) { | |||||
| key := getChannelPricingCacheKey(modelName, channelId) | |||||
| channelPricingCacheLock.RLock() | |||||
| cp, ok := channelPricingCache[key] | |||||
| channelPricingCacheLock.RUnlock() | |||||
| if !ok { | |||||
| return nil, false | |||||
| } | |||||
| return cp, true | |||||
| } | |||||
| // ParseTagIds 解析逗号分隔的标签ID字符串为 PricingTag 切片 | |||||
| func ParseTagIds(tagIds string, tagMap map[int]*PricingTag) []*PricingTag { | |||||
| if tagIds == "" { | |||||
| return nil | |||||
| } | |||||
| tags := make([]*PricingTag, 0) | |||||
| for _, idStr := range strings.Split(tagIds, ",") { | |||||
| if id, err := strconv.Atoi(strings.TrimSpace(idStr)); err == nil { | |||||
| if tag, ok := tagMap[id]; ok { | |||||
| tags = append(tags, tag) | |||||
| } | |||||
| } | |||||
| } | |||||
| return tags | |||||
| } | |||||
| // LoadChannelPricingCache 全量加载渠道定价到内存(启动时调用) | |||||
| func LoadChannelPricingCache() { | |||||
| var pricings []*ChannelPricing | |||||
| if err := DB.Find(&pricings).Error; err != nil { | |||||
| common.SysError("[ChannelPricing] LoadChannelPricingCache failed: " + err.Error()) | |||||
| return | |||||
| } | |||||
| channelPricingCacheLock.Lock() | |||||
| channelPricingCache = make(map[string]*ChannelPricing, len(pricings)) | |||||
| for _, cp := range pricings { | |||||
| key := getChannelPricingCacheKey(cp.ModelName, cp.ChannelId) | |||||
| channelPricingCache[key] = cp | |||||
| } | |||||
| channelPricingCacheLock.Unlock() | |||||
| common.SysLog(fmt.Sprintf("[ChannelPricing] cache loaded %d records", len(pricings))) | |||||
| } | |||||
| // ChannelPricingWithChannel 带渠道信息的定价响应 | |||||
| type ChannelPricingWithChannel struct { | |||||
| Id int `json:"id"` | |||||
| ChannelId int `json:"channel_id"` | |||||
| ChannelName string `json:"channel_name"` | |||||
| ChannelPublicName string `json:"channel_public_name"` | |||||
| ChannelType int `json:"channel_type"` | |||||
| TagIds string `json:"tag_ids" gorm:"column:tag_ids"` // 渠道定价的标签ID列表(逗号分隔) | |||||
| Tags []*PricingTag `json:"tags" gorm:"-"` // 渠道定价的标签详情(不参与数据库扫描) | |||||
| QuotaType int `json:"quota_type"` // 0=按量, 1=按次 | |||||
| ModelRatio float64 `json:"model_ratio"` | |||||
| CompletionRatio float64 `json:"completion_ratio"` | |||||
| ModelPrice float64 `json:"model_price"` | |||||
| HasCustomPricing bool `json:"has_custom_pricing"` // 是否有自定义定价 | |||||
| CacheRatio float64 `json:"cache_ratio"` | |||||
| CacheCreationRatio float64 `json:"cache_creation_ratio"` | |||||
| ImageRatio float64 `json:"image_ratio"` | |||||
| AudioRatio float64 `json:"audio_ratio"` | |||||
| AudioCompletionRatio float64 `json:"audio_completion_ratio"` | |||||
| IsDefault bool `json:"is_default"` | |||||
| UserRatio float64 `json:"user_ratio"` | |||||
| } | |||||
| // GetChannelPricingByModelWithChannelInfo 获取指定模型的渠道定价(带渠道信息) | |||||
| // 返回所有支持该模型的渠道,对于没有渠道定价的渠道使用全局默认价格 | |||||
| func GetChannelPricingByModelWithChannelInfo(modelName string) ([]*ChannelPricingWithChannel, error) { | |||||
| var results []*ChannelPricingWithChannel | |||||
| // 获取全局默认价格 | |||||
| globalModelRatio, hasRatio, _ := ratio_setting.GetModelRatio(modelName) | |||||
| globalModelPrice, hasPrice := ratio_setting.GetModelPrice(modelName, false) | |||||
| globalCompletionRatio := ratio_setting.GetCompletionRatio(modelName) | |||||
| // 确定默认计费类型 | |||||
| var defaultQuotaType int | |||||
| if hasPrice { | |||||
| defaultQuotaType = QuotaTypeByCall | |||||
| } else { | |||||
| defaultQuotaType = QuotaTypeByTokens | |||||
| } | |||||
| // 如果没有全局价格,设置默认值 | |||||
| if !hasRatio { | |||||
| globalModelRatio = 0 | |||||
| } | |||||
| if !hasPrice { | |||||
| globalModelPrice = 0 | |||||
| } | |||||
| // 查询所有支持该模型的渠道,左连接渠道定价表 | |||||
| // 高级字段(cache/image/audio)不回退全局值,直接返回 0 | |||||
| err := DB.Table("abilities"). | |||||
| Select(`abilities.channel_id, channels.name as channel_name, channels.public_name as channel_public_name, channels.type as channel_type, | |||||
| COALESCE(channel_pricings.quota_type, ?) as quota_type, | |||||
| COALESCE(channel_pricings.model_ratio, ?) as model_ratio, | |||||
| COALESCE(channel_pricings.completion_ratio, ?) as completion_ratio, | |||||
| COALESCE(channel_pricings.model_price, ?) as model_price, | |||||
| channel_pricings.id as id, | |||||
| channel_pricings.tag_ids as tag_ids, | |||||
| COALESCE(channel_pricings.cache_ratio, 0) as cache_ratio, | |||||
| COALESCE(channel_pricings.cache_creation_ratio, 0) as cache_creation_ratio, | |||||
| COALESCE(channel_pricings.image_ratio, 0) as image_ratio, | |||||
| COALESCE(channel_pricings.audio_ratio, 0) as audio_ratio, | |||||
| COALESCE(channel_pricings.audio_completion_ratio, 0) as audio_completion_ratio, | |||||
| COALESCE(channel_pricings.is_default, false) as is_default, | |||||
| (channel_pricings.id IS NOT NULL) as has_custom_pricing`, | |||||
| defaultQuotaType, globalModelRatio, globalCompletionRatio, globalModelPrice). | |||||
| Joins("LEFT JOIN channels ON abilities.channel_id = channels.id"). | |||||
| Joins("LEFT JOIN channel_pricings ON abilities.channel_id = channel_pricings.channel_id AND channel_pricings.model_name = ? AND channel_pricings.deleted_at IS NULL", modelName). | |||||
| Where("abilities.model = ?", modelName). | |||||
| Where("abilities.enabled = ?", true). | |||||
| Where("channels.status = ?", 1). // 只显示启用的渠道 | |||||
| Group("abilities.channel_id, channels.name, channels.public_name, channels.type, channel_pricings.quota_type, channel_pricings.model_ratio, channel_pricings.completion_ratio, channel_pricings.model_price, channel_pricings.id, channel_pricings.tag_ids, channel_pricings.cache_ratio, channel_pricings.cache_creation_ratio, channel_pricings.image_ratio, channel_pricings.audio_ratio, channel_pricings.audio_completion_ratio, channel_pricings.is_default"). | |||||
| Scan(&results).Error | |||||
| if err != nil { | |||||
| return nil, err | |||||
| } | |||||
| // 获取所有定价标签 | |||||
| allTags, err := GetAllPricingTags() | |||||
| if err != nil { | |||||
| return results, nil // 如果获取标签失败,仍然返回基础结果 | |||||
| } | |||||
| // 建立标签ID到标签的映射 | |||||
| tagMap := make(map[int]*PricingTag) | |||||
| for _, tag := range allTags { | |||||
| tagMap[tag.Id] = tag | |||||
| } | |||||
| // 为每个渠道定价填充标签 | |||||
| for _, result := range results { | |||||
| result.Tags = ParseTagIds(result.TagIds, tagMap) | |||||
| } | |||||
| return results, nil | |||||
| } | |||||
| @@ -207,7 +207,6 @@ func InitDB() (err error) { | |||||
| return err | return err | ||||
| } | } | ||||
| LoadEmailQuotaCache() | LoadEmailQuotaCache() | ||||
| LoadChannelPricingCache() | |||||
| LoadUserChannelRatioCache() | LoadUserChannelRatioCache() | ||||
| return nil | return nil | ||||
| } else { | } else { | ||||
| @@ -284,8 +283,6 @@ func migrateDB() error { | |||||
| &SubscriptionPreConsumeRecord{}, | &SubscriptionPreConsumeRecord{}, | ||||
| &CustomOAuthProvider{}, | &CustomOAuthProvider{}, | ||||
| &UserOAuthBinding{}, | &UserOAuthBinding{}, | ||||
| &ChannelPricing{}, | |||||
| &PricingTag{}, | |||||
| &UserChannelRatio{}, | &UserChannelRatio{}, | ||||
| &PendingSyncRecord{}, | &PendingSyncRecord{}, | ||||
| &QuotaSyncLog{}, | &QuotaSyncLog{}, | ||||
| @@ -295,6 +292,12 @@ func migrateDB() error { | |||||
| if err != nil { | if err != nil { | ||||
| return err | return err | ||||
| } | } | ||||
| if err := DB.Exec("DROP TABLE IF EXISTS channel_pricings").Error; err != nil { | |||||
| return err | |||||
| } | |||||
| if err := DB.Exec("DROP TABLE IF EXISTS pricing_tags").Error; err != nil { | |||||
| return err | |||||
| } | |||||
| if common.UsingSQLite { | if common.UsingSQLite { | ||||
| if err := ensureSubscriptionPlanTableSQLite(); err != nil { | if err := ensureSubscriptionPlanTableSQLite(); err != nil { | ||||
| return err | return err | ||||
| @@ -303,14 +306,14 @@ func migrateDB() error { | |||||
| if err := DB.AutoMigrate(&SubscriptionPlan{}); err != nil { | if err := DB.AutoMigrate(&SubscriptionPlan{}); err != nil { | ||||
| return err | return err | ||||
| } | } | ||||
| } | |||||
| // 将现有 sort_order=0 的模型和供应商更新为默认大数 | |||||
| DB.Model(&Model{}).Where("sort_order = 0").Update("sort_order", 999999) | |||||
| DB.Model(&Vendor{}).Where("sort_order = 0").Update("sort_order", 999999) | |||||
| migrateChannelPublicName() | |||||
| return nil | |||||
| } | } | ||||
| // 将现有 sort_order=0 的模型和供应商更新为默认大数 | |||||
| DB.Model(&Model{}).Where("sort_order = 0").Update("sort_order", 999999) | |||||
| DB.Model(&Vendor{}).Where("sort_order = 0").Update("sort_order", 999999) | |||||
| migrateChannelPublicName() | |||||
| return nil | |||||
| } | |||||
| func migrateDBFast() error { | func migrateDBFast() error { | ||||
| // Drop bound_channel_id column from tokens table (deprecated field) | // Drop bound_channel_id column from tokens table (deprecated field) | ||||
| @@ -346,8 +349,6 @@ func migrateDBFast() error { | |||||
| {&SubscriptionPreConsumeRecord{}, "SubscriptionPreConsumeRecord"}, | {&SubscriptionPreConsumeRecord{}, "SubscriptionPreConsumeRecord"}, | ||||
| {&CustomOAuthProvider{}, "CustomOAuthProvider"}, | {&CustomOAuthProvider{}, "CustomOAuthProvider"}, | ||||
| {&UserOAuthBinding{}, "UserOAuthBinding"}, | {&UserOAuthBinding{}, "UserOAuthBinding"}, | ||||
| {&ChannelPricing{}, "ChannelPricing"}, | |||||
| {&PricingTag{}, "PricingTag"}, | |||||
| {&UserChannelRatio{}, "UserChannelRatio"}, | {&UserChannelRatio{}, "UserChannelRatio"}, | ||||
| {&PendingSyncRecord{}, "PendingSyncRecord"}, | {&PendingSyncRecord{}, "PendingSyncRecord"}, | ||||
| {&QuotaSyncLog{}, "QuotaSyncLog"}, | {&QuotaSyncLog{}, "QuotaSyncLog"}, | ||||
| @@ -1,11 +1,9 @@ | |||||
| package model | package model | ||||
| import ( | import ( | ||||
| "encoding/json" | |||||
| "fmt" | "fmt" | ||||
| "sort" | "sort" | ||||
| "strings" | "strings" | ||||
| "sync" | "sync" | ||||
| "time" | "time" | ||||
| @@ -41,6 +39,11 @@ type PricingVendor struct { | |||||
| Icon string `json:"icon,omitempty"` | Icon string `json:"icon,omitempty"` | ||||
| } | } | ||||
| const ( | |||||
| QuotaTypeByTokens = 0 | |||||
| QuotaTypeByCall = 1 | |||||
| ) | |||||
| var ( | var ( | ||||
| pricingMap []Pricing | pricingMap []Pricing | ||||
| pricingByModel map[string]*Pricing | pricingByModel map[string]*Pricing | ||||
| @@ -49,7 +52,6 @@ var ( | |||||
| lastGetPricingTime time.Time | lastGetPricingTime time.Time | ||||
| updatePricingLock sync.Mutex | updatePricingLock sync.Mutex | ||||
| // 缓存映射:模型名 -> 启用分组 / 计费类型 | |||||
| modelEnableGroups = make(map[string][]string) | modelEnableGroups = make(map[string][]string) | ||||
| modelQuotaTypeMap = make(map[string]int) | modelQuotaTypeMap = make(map[string]int) | ||||
| modelEnableGroupsLock = sync.RWMutex{} | modelEnableGroupsLock = sync.RWMutex{} | ||||
| @@ -60,7 +62,6 @@ var ( | |||||
| modelSupportEndpointsLock = sync.RWMutex{} | modelSupportEndpointsLock = sync.RWMutex{} | ||||
| ) | ) | ||||
| // SetTestPricing 设置测试用的定价缓存(仅用于测试) | |||||
| func SetTestPricing(pricing []Pricing) { | func SetTestPricing(pricing []Pricing) { | ||||
| updatePricingLock.Lock() | updatePricingLock.Lock() | ||||
| defer updatePricingLock.Unlock() | defer updatePricingLock.Unlock() | ||||
| @@ -74,11 +75,10 @@ func SetTestPricing(pricing []Pricing) { | |||||
| } | } | ||||
| func GetPricing() []Pricing { | func GetPricing() []Pricing { | ||||
| if time.Since(lastGetPricingTime) > time.Minute*1 || len(pricingMap) == 0 { | |||||
| if time.Since(lastGetPricingTime) > time.Minute || len(pricingMap) == 0 { | |||||
| updatePricingLock.Lock() | updatePricingLock.Lock() | ||||
| defer updatePricingLock.Unlock() | defer updatePricingLock.Unlock() | ||||
| // Double check after acquiring the lock | |||||
| if time.Since(lastGetPricingTime) > time.Minute*1 || len(pricingMap) == 0 { | |||||
| if time.Since(lastGetPricingTime) > time.Minute || len(pricingMap) == 0 { | |||||
| modelSupportEndpointsLock.Lock() | modelSupportEndpointsLock.Lock() | ||||
| defer modelSupportEndpointsLock.Unlock() | defer modelSupportEndpointsLock.Unlock() | ||||
| updatePricing() | updatePricing() | ||||
| @@ -87,19 +87,16 @@ func GetPricing() []Pricing { | |||||
| return pricingMap | return pricingMap | ||||
| } | } | ||||
| // GetPricingByModel 从缓存中查找指定模型的定价信息 | |||||
| func GetPricingByModel(modelName string) *Pricing { | func GetPricingByModel(modelName string) *Pricing { | ||||
| GetPricing() // 确保缓存已刷新 | |||||
| GetPricing() | |||||
| modelEnableGroupsLock.RLock() | modelEnableGroupsLock.RLock() | ||||
| p := pricingByModel[modelName] | p := pricingByModel[modelName] | ||||
| modelEnableGroupsLock.RUnlock() | modelEnableGroupsLock.RUnlock() | ||||
| return p | return p | ||||
| } | } | ||||
| // GetVendors 返回当前定价接口使用到的供应商信息 | |||||
| func GetVendors() []PricingVendor { | func GetVendors() []PricingVendor { | ||||
| if time.Since(lastGetPricingTime) > time.Minute*1 || len(pricingMap) == 0 { | |||||
| // 保证先刷新一次 | |||||
| if time.Since(lastGetPricingTime) > time.Minute || len(pricingMap) == 0 { | |||||
| GetPricing() | GetPricing() | ||||
| } | } | ||||
| return vendorsList | return vendorsList | ||||
| @@ -118,13 +115,12 @@ func GetModelSupportEndpointTypes(model string) []constant.EndpointType { | |||||
| } | } | ||||
| func updatePricing() { | func updatePricing() { | ||||
| //modelRatios := common.GetModelRatios() | |||||
| enableAbilities, err := GetAllEnableAbilityWithChannels() | enableAbilities, err := GetAllEnableAbilityWithChannels() | ||||
| if err != nil { | if err != nil { | ||||
| common.SysLog(fmt.Sprintf("GetAllEnableAbilityWithChannels error: %v", err)) | common.SysLog(fmt.Sprintf("GetAllEnableAbilityWithChannels error: %v", err)) | ||||
| return | return | ||||
| } | } | ||||
| // 预加载模型元数据与供应商一次,避免循环查询 | |||||
| var allMeta []Model | var allMeta []Model | ||||
| _ = DB.Find(&allMeta).Error | _ = DB.Find(&allMeta).Error | ||||
| metaMap := make(map[string]*Model) | metaMap := make(map[string]*Model) | ||||
| @@ -133,21 +129,18 @@ func updatePricing() { | |||||
| containsList := make([]*Model, 0) | containsList := make([]*Model, 0) | ||||
| for i := range allMeta { | for i := range allMeta { | ||||
| m := &allMeta[i] | m := &allMeta[i] | ||||
| if m.NameRule == NameRuleExact { | |||||
| switch m.NameRule { | |||||
| case NameRuleExact: | |||||
| metaMap[m.ModelName] = m | metaMap[m.ModelName] = m | ||||
| } else { | |||||
| switch m.NameRule { | |||||
| case NameRulePrefix: | |||||
| prefixList = append(prefixList, m) | |||||
| case NameRuleSuffix: | |||||
| suffixList = append(suffixList, m) | |||||
| case NameRuleContains: | |||||
| containsList = append(containsList, m) | |||||
| } | |||||
| case NameRulePrefix: | |||||
| prefixList = append(prefixList, m) | |||||
| case NameRuleSuffix: | |||||
| suffixList = append(suffixList, m) | |||||
| case NameRuleContains: | |||||
| containsList = append(containsList, m) | |||||
| } | } | ||||
| } | } | ||||
| // 将非精确规则模型匹配到 metaMap | |||||
| for _, m := range prefixList { | for _, m := range prefixList { | ||||
| for _, pricingModel := range enableAbilities { | for _, pricingModel := range enableAbilities { | ||||
| if strings.HasPrefix(pricingModel.Model, m.ModelName) { | if strings.HasPrefix(pricingModel.Model, m.ModelName) { | ||||
| @@ -176,7 +169,6 @@ func updatePricing() { | |||||
| } | } | ||||
| } | } | ||||
| // 预加载供应商 | |||||
| var vendors []Vendor | var vendors []Vendor | ||||
| _ = DB.Find(&vendors).Error | _ = DB.Find(&vendors).Error | ||||
| vendorMap := make(map[int]*Vendor) | vendorMap := make(map[int]*Vendor) | ||||
| @@ -184,10 +176,8 @@ func updatePricing() { | |||||
| vendorMap[vendors[i].Id] = &vendors[i] | vendorMap[vendors[i].Id] = &vendors[i] | ||||
| } | } | ||||
| // 初始化默认供应商映射 | |||||
| initDefaultVendorMapping(metaMap, vendorMap, enableAbilities) | initDefaultVendorMapping(metaMap, vendorMap, enableAbilities) | ||||
| // 构建对前端友好的供应商列表 | |||||
| vendorOrderMap := make(map[int]int) | vendorOrderMap := make(map[int]int) | ||||
| for _, v := range vendorMap { | for _, v := range vendorMap { | ||||
| vendorOrderMap[v.Id] = v.SortOrder | vendorOrderMap[v.Id] = v.SortOrder | ||||
| @@ -212,7 +202,6 @@ func updatePricing() { | |||||
| }) | }) | ||||
| modelGroupsMap := make(map[string]*types.Set[string]) | modelGroupsMap := make(map[string]*types.Set[string]) | ||||
| for _, ability := range enableAbilities { | for _, ability := range enableAbilities { | ||||
| groups, ok := modelGroupsMap[ability.Model] | groups, ok := modelGroupsMap[ability.Model] | ||||
| if !ok { | if !ok { | ||||
| @@ -222,10 +211,7 @@ func updatePricing() { | |||||
| groups.Add(ability.Group) | groups.Add(ability.Group) | ||||
| } | } | ||||
| //这里使用切片而不是Set,因为一个模型可能支持多个端点类型,并且第一个端点是优先使用端点 | |||||
| modelSupportEndpointsStr := make(map[string][]string) | modelSupportEndpointsStr := make(map[string][]string) | ||||
| // 先根据已有能力填充原生端点 | |||||
| for _, ability := range enableAbilities { | for _, ability := range enableAbilities { | ||||
| endpoints := modelSupportEndpointsStr[ability.Model] | endpoints := modelSupportEndpointsStr[ability.Model] | ||||
| channelTypes := common.GetEndpointTypesByChannelType(ability.ChannelType, ability.Model) | channelTypes := common.GetEndpointTypesByChannelType(ability.ChannelType, ability.Model) | ||||
| @@ -237,13 +223,12 @@ func updatePricing() { | |||||
| modelSupportEndpointsStr[ability.Model] = endpoints | modelSupportEndpointsStr[ability.Model] = endpoints | ||||
| } | } | ||||
| // 再补充模型自定义端点:若配置有效则替换默认端点,不做合并 | |||||
| for modelName, meta := range metaMap { | for modelName, meta := range metaMap { | ||||
| if strings.TrimSpace(meta.Endpoints) == "" { | if strings.TrimSpace(meta.Endpoints) == "" { | ||||
| continue | continue | ||||
| } | } | ||||
| var raw map[string]interface{} | var raw map[string]interface{} | ||||
| if err := json.Unmarshal([]byte(meta.Endpoints), &raw); err == nil { | |||||
| if err := common.Unmarshal([]byte(meta.Endpoints), &raw); err == nil { | |||||
| endpoints := make([]string, 0, len(raw)) | endpoints := make([]string, 0, len(raw)) | ||||
| for k, v := range raw { | for k, v := range raw { | ||||
| switch v.(type) { | switch v.(type) { | ||||
| @@ -263,15 +248,12 @@ func updatePricing() { | |||||
| for model, endpoints := range modelSupportEndpointsStr { | for model, endpoints := range modelSupportEndpointsStr { | ||||
| supportedEndpoints := make([]constant.EndpointType, 0) | supportedEndpoints := make([]constant.EndpointType, 0) | ||||
| for _, endpointStr := range endpoints { | for _, endpointStr := range endpoints { | ||||
| endpointType := constant.EndpointType(endpointStr) | |||||
| supportedEndpoints = append(supportedEndpoints, endpointType) | |||||
| supportedEndpoints = append(supportedEndpoints, constant.EndpointType(endpointStr)) | |||||
| } | } | ||||
| modelSupportEndpointTypes[model] = supportedEndpoints | modelSupportEndpointTypes[model] = supportedEndpoints | ||||
| } | } | ||||
| // 构建全局 supportedEndpointMap(默认 + 自定义覆盖) | |||||
| supportedEndpointMap = make(map[string]common.EndpointInfo) | supportedEndpointMap = make(map[string]common.EndpointInfo) | ||||
| // 1. 默认端点 | |||||
| for _, endpoints := range modelSupportEndpointTypes { | for _, endpoints := range modelSupportEndpointTypes { | ||||
| for _, et := range endpoints { | for _, et := range endpoints { | ||||
| if info, ok := common.GetDefaultEndpointInfo(et); ok { | if info, ok := common.GetDefaultEndpointInfo(et); ok { | ||||
| @@ -281,13 +263,12 @@ func updatePricing() { | |||||
| } | } | ||||
| } | } | ||||
| } | } | ||||
| // 2. 自定义端点(models 表)覆盖默认 | |||||
| for _, meta := range metaMap { | for _, meta := range metaMap { | ||||
| if strings.TrimSpace(meta.Endpoints) == "" { | if strings.TrimSpace(meta.Endpoints) == "" { | ||||
| continue | continue | ||||
| } | } | ||||
| var raw map[string]interface{} | var raw map[string]interface{} | ||||
| if err := json.Unmarshal([]byte(meta.Endpoints), &raw); err == nil { | |||||
| if err := common.Unmarshal([]byte(meta.Endpoints), &raw); err == nil { | |||||
| for k, v := range raw { | for k, v := range raw { | ||||
| switch val := v.(type) { | switch val := v.(type) { | ||||
| case string: | case string: | ||||
| @@ -301,42 +282,19 @@ func updatePricing() { | |||||
| ep.Method = strings.ToUpper(m) | ep.Method = strings.ToUpper(m) | ||||
| } | } | ||||
| supportedEndpointMap[k] = ep | supportedEndpointMap[k] = ep | ||||
| default: | |||||
| // ignore unsupported types | |||||
| } | } | ||||
| } | } | ||||
| } | } | ||||
| } | } | ||||
| // 从渠道定价表加载实际定价数据(仅启用渠道),同时获取渠道名称 | |||||
| var allCPs []struct { | |||||
| ChannelPricing | |||||
| ChannelName string | |||||
| ChannelPublicName string | |||||
| } | |||||
| DB.Table("channel_pricings"). | |||||
| Select("channel_pricings.*, channels.name as channel_name, channels.public_name as channel_public_name"). | |||||
| Joins("JOIN channels ON channel_pricings.channel_id = channels.id"). | |||||
| Where("channels.status = 1 AND channel_pricings.deleted_at IS NULL"). | |||||
| Find(&allCPs) | |||||
| cpMap := make(map[string][]ChannelPricing) | |||||
| channelNameMap := make(map[int]string) | |||||
| for i := range allCPs { | |||||
| cpMap[allCPs[i].ModelName] = append(cpMap[allCPs[i].ModelName], allCPs[i].ChannelPricing) | |||||
| channelNameMap[allCPs[i].ChannelId] = ChannelDisplayName(allCPs[i].ChannelPublicName, allCPs[i].ChannelName) | |||||
| } | |||||
| pricingMap = make([]Pricing, 0) | |||||
| pricingMap = make([]Pricing, 0, len(modelGroupsMap)) | |||||
| for model, groups := range modelGroupsMap { | for model, groups := range modelGroupsMap { | ||||
| pricing := Pricing{ | pricing := Pricing{ | ||||
| ModelName: model, | ModelName: model, | ||||
| EnableGroup: groups.Items(), | EnableGroup: groups.Items(), | ||||
| SupportedEndpointTypes: modelSupportEndpointTypes[model], | SupportedEndpointTypes: modelSupportEndpointTypes[model], | ||||
| } | } | ||||
| // 补充模型元数据(描述、标签、供应商、状态) | |||||
| if meta, ok := metaMap[model]; ok { | if meta, ok := metaMap[model]; ok { | ||||
| // 若模型被禁用(status==0),则直接跳过,不返回给前端 | |||||
| if meta.Status == 0 { | if meta.Status == 0 { | ||||
| continue | continue | ||||
| } | } | ||||
| @@ -346,14 +304,10 @@ func updatePricing() { | |||||
| pricing.VendorID = meta.VendorID | pricing.VendorID = meta.VendorID | ||||
| pricing.Type = meta.Type | pricing.Type = meta.Type | ||||
| } | } | ||||
| // 使用渠道定价表中的实际数据,选取最便宜的渠道 | |||||
| applyBestChannelPricing(&pricing, cpMap[model], model) | |||||
| applyGlobalDefault(&pricing, model) | |||||
| pricingMap = append(pricingMap, pricing) | pricingMap = append(pricingMap, pricing) | ||||
| } | } | ||||
| // 按 sort_order 排序 pricingMap,999999 视为未设置 | |||||
| sort.Slice(pricingMap, func(i, j int) bool { | sort.Slice(pricingMap, func(i, j int) bool { | ||||
| mi, okI := metaMap[pricingMap[i].ModelName] | mi, okI := metaMap[pricingMap[i].ModelName] | ||||
| mj, okJ := metaMap[pricingMap[j].ModelName] | mj, okJ := metaMap[pricingMap[j].ModelName] | ||||
| @@ -371,12 +325,10 @@ func updatePricing() { | |||||
| return pricingMap[i].ModelName < pricingMap[j].ModelName | return pricingMap[i].ModelName < pricingMap[j].ModelName | ||||
| }) | }) | ||||
| // 防止大更新后数据不通用 | |||||
| if len(pricingMap) > 0 { | if len(pricingMap) > 0 { | ||||
| pricingMap[0].PricingVersion = "82c4a357505fff6fee8462c3f7ec8a645bb95532669cb73b2cabee6a416ec24f" | pricingMap[0].PricingVersion = "82c4a357505fff6fee8462c3f7ec8a645bb95532669cb73b2cabee6a416ec24f" | ||||
| } | } | ||||
| // 刷新缓存映射,供高并发快速查询 | |||||
| modelEnableGroupsLock.Lock() | modelEnableGroupsLock.Lock() | ||||
| modelEnableGroups = make(map[string][]string) | modelEnableGroups = make(map[string][]string) | ||||
| modelQuotaTypeMap = make(map[string]int) | modelQuotaTypeMap = make(map[string]int) | ||||
| @@ -393,79 +345,24 @@ func updatePricing() { | |||||
| lastGetPricingTime = time.Now() | lastGetPricingTime = time.Now() | ||||
| } | } | ||||
| // GetSupportedEndpointMap 返回全局端点到路径的映射 | |||||
| func GetSupportedEndpointMap() map[string]common.EndpointInfo { | func GetSupportedEndpointMap() map[string]common.EndpointInfo { | ||||
| return supportedEndpointMap | return supportedEndpointMap | ||||
| } | } | ||||
| // applyGlobalDefault 用全局默认值填充 Pricing(无渠道定价时的回退) | |||||
| func applyGlobalDefault(pricing *Pricing, model string) { | func applyGlobalDefault(pricing *Pricing, model string) { | ||||
| modelPrice, findPrice := ratio_setting.GetModelPrice(model, false) | modelPrice, findPrice := ratio_setting.GetModelPrice(model, false) | ||||
| if findPrice { | if findPrice { | ||||
| pricing.ModelPrice = modelPrice | pricing.ModelPrice = modelPrice | ||||
| pricing.QuotaType = 1 | |||||
| pricing.ModelRatio = 0 | |||||
| pricing.CompletionRatio = 0 | |||||
| pricing.QuotaType = QuotaTypeByCall | |||||
| } else { | } else { | ||||
| modelRatio, _, _ := ratio_setting.GetModelRatio(model) | modelRatio, _, _ := ratio_setting.GetModelRatio(model) | ||||
| pricing.ModelPrice = 0 | |||||
| pricing.ModelRatio = modelRatio | pricing.ModelRatio = modelRatio | ||||
| pricing.CompletionRatio = ratio_setting.GetCompletionRatio(model) | pricing.CompletionRatio = ratio_setting.GetCompletionRatio(model) | ||||
| pricing.QuotaType = 0 | |||||
| pricing.QuotaType = QuotaTypeByTokens | |||||
| } | } | ||||
| pricing.CacheRatio, _ = ratio_setting.GetCacheRatio(model) | pricing.CacheRatio, _ = ratio_setting.GetCacheRatio(model) | ||||
| pricing.CacheCreationRatio, _ = ratio_setting.GetCreateCacheRatio(model) | pricing.CacheCreationRatio, _ = ratio_setting.GetCreateCacheRatio(model) | ||||
| } | } | ||||
| // applyBestChannelPricing 从渠道定价中选取最优(最便宜)的价格填充 Pricing | |||||
| // 优先选按量计费(quota_type=0)的渠道,因为缓存价格仅对按量计费有意义 | |||||
| // 扩展比率字段(cache_ratio 等)为 0 表示未设置,需回退到全局默认值 | |||||
| func applyBestChannelPricing(pricing *Pricing, cps []ChannelPricing, model string) { | |||||
| if len(cps) == 0 { | |||||
| applyGlobalDefault(pricing, model) | |||||
| return | |||||
| } | |||||
| // 优先选按量计费 (quota_type=0) 中 model_ratio 最低的渠道 | |||||
| var bestPerToken *ChannelPricing | |||||
| for i := range cps { | |||||
| cp := &cps[i] | |||||
| if cp.QuotaType == 0 { | |||||
| if bestPerToken == nil || cp.ModelRatio < bestPerToken.ModelRatio { | |||||
| bestPerToken = cp | |||||
| } | |||||
| } | |||||
| } | |||||
| if bestPerToken != nil { | |||||
| pricing.QuotaType = 0 | |||||
| pricing.ModelRatio = bestPerToken.ModelRatio | |||||
| pricing.CompletionRatio = bestPerToken.CompletionRatio | |||||
| if bestPerToken.CacheRatio > 0 { | |||||
| pricing.CacheRatio = bestPerToken.CacheRatio | |||||
| } else { | |||||
| pricing.CacheRatio, _ = ratio_setting.GetCacheRatio(model) | |||||
| } | |||||
| if bestPerToken.CacheCreationRatio > 0 { | |||||
| pricing.CacheCreationRatio = bestPerToken.CacheCreationRatio | |||||
| } else { | |||||
| pricing.CacheCreationRatio, _ = ratio_setting.GetCreateCacheRatio(model) | |||||
| } | |||||
| return | |||||
| } | |||||
| // 没有按量渠道,选按次计费 (quota_type=1) 中 model_price 最低的渠道 | |||||
| var bestPerCall *ChannelPricing | |||||
| for i := range cps { | |||||
| cp := &cps[i] | |||||
| if cp.QuotaType == 1 { | |||||
| if bestPerCall == nil || cp.ModelPrice < bestPerCall.ModelPrice { | |||||
| bestPerCall = cp | |||||
| } | |||||
| } | |||||
| } | |||||
| if bestPerCall != nil { | |||||
| pricing.QuotaType = 1 | |||||
| pricing.ModelPrice = bestPerCall.ModelPrice | |||||
| return | |||||
| } | |||||
| applyGlobalDefault(pricing, model) | |||||
| } | |||||
| @@ -1,60 +0,0 @@ | |||||
| package model | |||||
| import ( | |||||
| "github.com/QuantumNous/new-api/common" | |||||
| "gorm.io/gorm" | |||||
| ) | |||||
| // PricingTag 定价标签表 | |||||
| // 用于对渠道定价进行分类和标识 | |||||
| type PricingTag struct { | |||||
| Id int `json:"id" gorm:"primaryKey"` | |||||
| Name string `json:"name" gorm:"size:64;not null;uniqueIndex"` | |||||
| Color string `json:"color" gorm:"size:16;default:'#1890ff'"` | |||||
| Description string `json:"description" gorm:"type:text"` | |||||
| SortOrder int `json:"sort_order" gorm:"default:0"` | |||||
| CreatedTime int64 `json:"created_time" gorm:"bigint"` | |||||
| DeletedAt gorm.DeletedAt `json:"-" gorm:"index"` | |||||
| } | |||||
| func (pt *PricingTag) Insert() error { | |||||
| pt.CreatedTime = common.GetTimestamp() | |||||
| return DB.Create(pt).Error | |||||
| } | |||||
| func (pt *PricingTag) Update() error { | |||||
| return DB.Model(&PricingTag{}).Where("id = ?", pt.Id).Updates(map[string]interface{}{ | |||||
| "name": pt.Name, | |||||
| "color": pt.Color, | |||||
| "description": pt.Description, | |||||
| "sort_order": pt.SortOrder, | |||||
| }).Error | |||||
| } | |||||
| func (pt *PricingTag) Delete() error { | |||||
| return DB.Delete(pt).Error | |||||
| } | |||||
| func GetPricingTagById(id int) (*PricingTag, error) { | |||||
| var pt PricingTag | |||||
| err := DB.First(&pt, id).Error | |||||
| if err != nil { | |||||
| return nil, err | |||||
| } | |||||
| return &pt, nil | |||||
| } | |||||
| func GetAllPricingTags() ([]*PricingTag, error) { | |||||
| var list []*PricingTag | |||||
| err := DB.Order("sort_order ASC, id ASC").Find(&list).Error | |||||
| return list, err | |||||
| } | |||||
| func GetPricingTagByName(name string) (*PricingTag, error) { | |||||
| var pt PricingTag | |||||
| err := DB.Where("name = ?", name).First(&pt).Error | |||||
| if err != nil { | |||||
| return nil, err | |||||
| } | |||||
| return &pt, nil | |||||
| } | |||||
| @@ -2,13 +2,19 @@ package model | |||||
| import ( | import ( | ||||
| "testing" | "testing" | ||||
| "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/setting/ratio_setting" | ||||
| "github.com/glebarez/sqlite" | |||||
| "github.com/stretchr/testify/require" | "github.com/stretchr/testify/require" | ||||
| "gorm.io/gorm" | |||||
| ) | ) | ||||
| const testPricingModel = "test-pricing-model-apply" | |||||
| const testPricingApplyModel = "test-pricing-model-apply" | |||||
| const testPricingCallModel = "test-pricing-call-model" | const testPricingCallModel = "test-pricing-call-model" | ||||
| const testPricingGlobalModel = "test-pricing-global-model" | |||||
| func resetPricingRatios(t *testing.T) { | func resetPricingRatios(t *testing.T) { | ||||
| t.Helper() | t.Helper() | ||||
| @@ -19,15 +25,61 @@ func resetPricingRatios(t *testing.T) { | |||||
| require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{}`)) | require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{}`)) | ||||
| } | } | ||||
| func setupPricingTest(t *testing.T) { | |||||
| t.Helper() | |||||
| resetPricingRatios(t) | |||||
| db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) | |||||
| require.NoError(t, err) | |||||
| sqlDB, err := db.DB() | |||||
| require.NoError(t, err) | |||||
| sqlDB.SetMaxOpenConns(1) | |||||
| origDB := DB | |||||
| origUsingSQLite := common.UsingSQLite | |||||
| origPricingMap := pricingMap | |||||
| origPricingByModel := pricingByModel | |||||
| origVendorsList := vendorsList | |||||
| origSupportedEndpointMap := supportedEndpointMap | |||||
| origModelSupportEndpointTypes := modelSupportEndpointTypes | |||||
| origLastGetPricingTime := lastGetPricingTime | |||||
| DB = db | |||||
| common.UsingSQLite = true | |||||
| pricingMap = nil | |||||
| pricingByModel = nil | |||||
| vendorsList = nil | |||||
| supportedEndpointMap = nil | |||||
| modelSupportEndpointTypes = make(map[string][]constant.EndpointType) | |||||
| lastGetPricingTime = time.Time{} | |||||
| require.NoError(t, db.AutoMigrate(&Model{}, &Vendor{}, &Channel{}, &Ability{})) | |||||
| t.Cleanup(func() { | |||||
| DB = origDB | |||||
| common.UsingSQLite = origUsingSQLite | |||||
| pricingMap = origPricingMap | |||||
| pricingByModel = origPricingByModel | |||||
| vendorsList = origVendorsList | |||||
| supportedEndpointMap = origSupportedEndpointMap | |||||
| modelSupportEndpointTypes = origModelSupportEndpointTypes | |||||
| lastGetPricingTime = origLastGetPricingTime | |||||
| require.NoError(t, sqlDB.Close()) | |||||
| resetPricingRatios(t) | |||||
| }) | |||||
| } | |||||
| func TestApplyGlobalDefault_PerTokenModel(t *testing.T) { | func TestApplyGlobalDefault_PerTokenModel(t *testing.T) { | ||||
| resetPricingRatios(t) | resetPricingRatios(t) | ||||
| require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testPricingModel+`":10}`)) | |||||
| require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testPricingModel+`":3}`)) | |||||
| require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testPricingModel+`":0.5}`)) | |||||
| require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testPricingModel+`":0.75}`)) | |||||
| require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testPricingApplyModel+`":10}`)) | |||||
| require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testPricingApplyModel+`":3}`)) | |||||
| require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testPricingApplyModel+`":0.5}`)) | |||||
| require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testPricingApplyModel+`":0.75}`)) | |||||
| p := &Pricing{} | p := &Pricing{} | ||||
| applyGlobalDefault(p, testPricingModel) | |||||
| applyGlobalDefault(p, testPricingApplyModel) | |||||
| require.Equal(t, QuotaTypeByTokens, p.QuotaType) | require.Equal(t, QuotaTypeByTokens, p.QuotaType) | ||||
| require.Equal(t, 10.0, p.ModelRatio) | require.Equal(t, 10.0, p.ModelRatio) | ||||
| @@ -47,3 +99,41 @@ func TestApplyGlobalDefault_PerCallModel(t *testing.T) { | |||||
| require.Equal(t, 0.3, p.ModelPrice) | require.Equal(t, 0.3, p.ModelPrice) | ||||
| require.Equal(t, 0.0, p.ModelRatio) | require.Equal(t, 0.0, p.ModelRatio) | ||||
| } | } | ||||
| func TestGetPricingByModel_UsesGlobalDefaultOnly(t *testing.T) { | |||||
| setupPricingTest(t) | |||||
| require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testPricingGlobalModel+`":10}`)) | |||||
| require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testPricingGlobalModel+`":3}`)) | |||||
| require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testPricingGlobalModel+`":0.5}`)) | |||||
| require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testPricingGlobalModel+`":0.75}`)) | |||||
| require.NoError(t, (&Model{ | |||||
| ModelName: testPricingGlobalModel, | |||||
| Status: 1, | |||||
| }).Insert()) | |||||
| require.NoError(t, (&Vendor{ | |||||
| Id: 1, | |||||
| Name: "test-vendor", | |||||
| SortOrder: 1, | |||||
| }).Insert()) | |||||
| require.NoError(t, (&Channel{ | |||||
| Id: 1, | |||||
| Type: constant.ChannelTypeOpenAI, | |||||
| Key: "k", | |||||
| Name: "channel-a", | |||||
| Status: common.ChannelStatusEnabled, | |||||
| Models: testPricingGlobalModel, | |||||
| Group: "default", | |||||
| }).Insert()) | |||||
| RefreshPricing() | |||||
| pricing := GetPricingByModel(testPricingGlobalModel) | |||||
| require.NotNil(t, pricing) | |||||
| require.Equal(t, QuotaTypeByTokens, pricing.QuotaType) | |||||
| require.Equal(t, 10.0, pricing.ModelRatio) | |||||
| require.Equal(t, 3.0, pricing.CompletionRatio) | |||||
| require.Equal(t, 0.5, pricing.CacheRatio) | |||||
| require.Equal(t, 0.75, pricing.CacheCreationRatio) | |||||
| } | |||||
| @@ -0,0 +1,45 @@ | |||||
| package model | |||||
| import ( | |||||
| "testing" | |||||
| "github.com/QuantumNous/new-api/common" | |||||
| "github.com/glebarez/sqlite" | |||||
| "github.com/stretchr/testify/require" | |||||
| "gorm.io/gorm" | |||||
| ) | |||||
| type channelPricingMigrationStub struct { | |||||
| ID uint `gorm:"primaryKey"` | |||||
| } | |||||
| func (channelPricingMigrationStub) TableName() string { | |||||
| return "channel_pricings" | |||||
| } | |||||
| type pricingTagMigrationStub struct { | |||||
| ID uint `gorm:"primaryKey"` | |||||
| } | |||||
| func (pricingTagMigrationStub) TableName() string { | |||||
| return "pricing_tags" | |||||
| } | |||||
| func TestMigrateDBDropsChannelPricingTables(t *testing.T) { | |||||
| db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) | |||||
| require.NoError(t, err) | |||||
| origDB := DB | |||||
| origSQLite := common.UsingSQLite | |||||
| DB = db | |||||
| common.UsingSQLite = true | |||||
| t.Cleanup(func() { | |||||
| DB = origDB | |||||
| common.UsingSQLite = origSQLite | |||||
| }) | |||||
| require.NoError(t, db.AutoMigrate(&channelPricingMigrationStub{}, &pricingTagMigrationStub{})) | |||||
| require.NoError(t, migrateDB()) | |||||
| require.False(t, DB.Migrator().HasTable("channel_pricings")) | |||||
| require.False(t, DB.Migrator().HasTable("pricing_tags")) | |||||
| } | |||||
| @@ -0,0 +1,11 @@ | |||||
| <!doctype html> | |||||
| <html lang="zh-CN"> | |||||
| <head> | |||||
| <meta charset="UTF-8" /> | |||||
| <meta name="viewport" content="width=device-width, initial-scale=1.0" /> | |||||
| <title>new-api</title> | |||||
| </head> | |||||
| <body> | |||||
| <div id="root">Frontend assets are not built in this checkout.</div> | |||||
| </body> | |||||
| </html> | |||||