|
|
|
@@ -1,11 +1,23 @@ |
|
|
|
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 // 按量计费 |
|
|
|
@@ -109,3 +121,75 @@ func BatchUpsertChannelPricing(pricings []*ChannelPricing) error { |
|
|
|
}), |
|
|
|
}).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{} // 重置为零值 |
|
|
|
} |