package model import ( "fmt" "strconv" "strings" "sync" "time" "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 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 err := DB.Create(cp).Error if err == nil { InvalidateChannelPricingCache() } return err } func (cp *ChannelPricing) Update() error { cp.UpdatedTime = common.GetTimestamp() err := 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 if err == nil { InvalidateChannelPricingCache() } return err } func (cp *ChannelPricing) Delete() error { err := DB.Delete(cp).Error if err == nil { InvalidateChannelPricingCache() } 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) 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, cp.QuotaType == QuotaTypeByCall, 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, cp.QuotaType == QuotaTypeByCall, 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{} // 重置为零值 } // 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"` // 是否有自定义定价 } // 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 } // 查询所有支持该模型的渠道,左连接渠道定价表 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, (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"). 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 { if result.TagIds != "" { result.Tags = make([]*PricingTag, 0) for _, idStr := range strings.Split(result.TagIds, ",") { if id, err := strconv.Atoi(strings.TrimSpace(idStr)); err == nil { if tag, ok := tagMap[id]; ok { result.Tags = append(result.Tags, tag) } } } } } return results, nil }