package model import ( "fmt" "sync" "time" "github.com/QuantumNous/new-api/common" "gorm.io/gorm" "gorm.io/gorm/clause" ) // 渠道定价缓存 var ( channelPricingCache = make(map[string]*ChannelPricing) // key: "modelName:channelId" channelPricingCacheLock sync.RWMutex channelPricingCacheTime time.Time channelPricingCacheTTL = time.Minute * 5 // 缓存5分钟 ) // 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"` } func (cp *ChannelPricing) Insert() error { now := common.GetTimestamp() cp.CreatedTime = now cp.UpdatedTime = now return DB.Create(cp).Error } func (cp *ChannelPricing) Update() error { cp.UpdatedTime = common.GetTimestamp() return DB.Model(&ChannelPricing{}).Where("id = ?", cp.Id).Updates(map[string]interface{}{ "quota_type": cp.QuotaType, "model_ratio": cp.ModelRatio, "completion_ratio": cp.CompletionRatio, "model_price": cp.ModelPrice, "tag_ids": cp.TagIds, "updated_time": cp.UpdatedTime, }).Error } func (cp *ChannelPricing) Delete() error { return DB.Delete(cp).Error } // 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) return 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", "updated_time", }), }).Create(&pricings).Error } // getChannelPricingCacheKey 生成缓存键 func getChannelPricingCacheKey(modelName string, channelId int) string { return fmt.Sprintf("%s:%d", modelName, channelId) } // GetEffectivePricing 获取有效定价(优先渠道定价,回退全局定价) // 返回: modelRatio, completionRatio, modelPrice, usePrice, found func GetEffectivePricing(modelName string, channelId int) (modelRatio, completionRatio, modelPrice float64, usePrice, found bool) { cacheKey := getChannelPricingCacheKey(modelName, channelId) // 首先检查缓存 channelPricingCacheLock.RLock() // 检查缓存是否过期 if time.Since(channelPricingCacheTime) < channelPricingCacheTTL { if cp, ok := channelPricingCache[cacheKey]; ok { channelPricingCacheLock.RUnlock() return cp.ModelRatio, cp.CompletionRatio, cp.ModelPrice, true, true } } channelPricingCacheLock.RUnlock() // 缓存未命中或已过期,查询数据库 var cp ChannelPricing err := DB.Where("model_name = ? AND channel_id = ?", modelName, channelId).First(&cp).Error if err != nil { // 未找到渠道定价,返回 false 让调用者使用全局定价 return 0, 0, 0, false, false } // 更新缓存 channelPricingCacheLock.Lock() if channelPricingCacheTime.IsZero() || time.Since(channelPricingCacheTime) >= channelPricingCacheTTL { // 缓存过期,清空并更新时间 channelPricingCache = make(map[string]*ChannelPricing) channelPricingCacheTime = time.Now() } channelPricingCache[cacheKey] = &cp channelPricingCacheLock.Unlock() return cp.ModelRatio, cp.CompletionRatio, cp.ModelPrice, true, true } // RefreshChannelPricingCache 刷新渠道定价缓存 func RefreshChannelPricingCache() { channelPricingCacheLock.Lock() defer channelPricingCacheLock.Unlock() // 清空缓存 channelPricingCache = make(map[string]*ChannelPricing) channelPricingCacheTime = time.Now() // 预加载所有渠道定价 var pricings []*ChannelPricing if err := DB.Find(&pricings).Error; err != nil { return } for _, cp := range pricings { cacheKey := getChannelPricingCacheKey(cp.ModelName, cp.ChannelId) channelPricingCache[cacheKey] = cp } } // InvalidateChannelPricingCache 使渠道定价缓存失效 func InvalidateChannelPricingCache() { channelPricingCacheLock.Lock() defer channelPricingCacheLock.Unlock() channelPricingCache = make(map[string]*ChannelPricing) channelPricingCacheTime = time.Time{} // 重置为零值 }