|
- 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"`
- 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"`
- }
-
- // 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
- }
-
- // === 默认通道缓存 ===
-
- // 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
- }
|