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 // 默认通道缓存:modelName → channelId defaultChannelCache = make(map[string]int) defaultChannelCacheLock 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) if cp.IsDefault { setDefaultChannelCache(cp.ModelName, cp.ChannelId) } } 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) if cp.IsDefault { setDefaultChannelCache(cp.ModelName, cp.ChannelId) } else { clearDefaultChannelCacheIfMatch(cp.ModelName, cp.Id) } } 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)) if existing.IsDefault { clearDefaultChannelCache(existing.ModelName) } } 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() rebuildDefaultChannelCache(pricings) 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"` 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"` } // 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.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.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 } // === 默认通道缓存 === // syncIsDefaultToCache 将默认标记的变更同步到 channelPricingCache func syncIsDefaultToCache(modelName string, channelId int, isDefault bool) { channelPricingCacheLock.Lock() for key, cp := range channelPricingCache { if cp.ModelName == modelName { if isDefault { cp.IsDefault = cp.ChannelId == channelId } else { cp.IsDefault = false } } channelPricingCache[key] = cp } channelPricingCacheLock.Unlock() } // setDefaultChannelCache 设置默认通道缓存 func setDefaultChannelCache(modelName string, channelId int) { defaultChannelCacheLock.Lock() defaultChannelCache[modelName] = channelId defaultChannelCacheLock.Unlock() } // clearDefaultChannelCache 清除指定模型的默认通道缓存 func clearDefaultChannelCache(modelName string) { defaultChannelCacheLock.Lock() delete(defaultChannelCache, modelName) defaultChannelCacheLock.Unlock() } // clearDefaultChannelCacheIfMatch 如果默认通道的定价记录 ID 匹配则清除 func clearDefaultChannelCacheIfMatch(modelName string, pricingId int) { defaultChannelCacheLock.RLock() cachedId, ok := defaultChannelCache[modelName] defaultChannelCacheLock.RUnlock() if !ok { return } // 需要通过缓存找到对应的 pricing 来比对 key := getChannelPricingCacheKey(modelName, cachedId) channelPricingCacheLock.RLock() cp, exists := channelPricingCache[key] channelPricingCacheLock.RUnlock() if exists && cp.Id == pricingId { clearDefaultChannelCache(modelName) } } // rebuildDefaultChannelCache 从全量数据构建默认通道缓存(启动时调用) func rebuildDefaultChannelCache(pricings []*ChannelPricing) { defaultChannelCacheLock.Lock() defaultChannelCache = make(map[string]int) for _, cp := range pricings { if cp.IsDefault { defaultChannelCache[cp.ModelName] = cp.ChannelId } } defaultChannelCacheLock.Unlock() common.SysLog(fmt.Sprintf("[ChannelPricing] default channel cache loaded %d records", len(defaultChannelCache))) } // GetDefaultChannelId 获取指定模型的默认通道 ID(纯内存读) func GetDefaultChannelId(modelName string) (int, bool) { defaultChannelCacheLock.RLock() id, ok := defaultChannelCache[modelName] defaultChannelCacheLock.RUnlock() return id, ok } // SetDefaultChannel 设置指定模型的默认通道(事务保证互斥) func SetDefaultChannel(modelName string, channelId int) error { return DB.Transaction(func(tx *gorm.DB) error { // 清除该模型所有现有的默认标记 if err := tx.Model(&ChannelPricing{}). Where("model_name = ? AND is_default = ?", modelName, true). Update("is_default", false).Error; err != nil { return err } // 设置新的默认 if err := tx.Model(&ChannelPricing{}). Where("model_name = ? AND channel_id = ?", modelName, channelId). Update("is_default", true).Error; err != nil { return err } // 更新缓存 setDefaultChannelCache(modelName, channelId) syncIsDefaultToCache(modelName, channelId, true) return nil }) } // ClearDefaultChannel 清除指定模型的默认通道标记 func ClearDefaultChannel(modelName string) error { err := DB.Model(&ChannelPricing{}). Where("model_name = ? AND is_default = ?", modelName, true). Update("is_default", false).Error if err == nil { clearDefaultChannelCache(modelName) syncIsDefaultToCache(modelName, 0, false) } return err }