From 6c9ee5ec2b0f46ce7e4b400016262d182714a9a4 Mon Sep 17 00:00:00 2001 From: fengsilin Date: Tue, 24 Mar 2026 11:12:36 +0800 Subject: [PATCH] feat: add channel pricing cache and GetEffectivePricing function Co-Authored-By: Claude Opus 4.6 --- model/channel_pricing.go | 84 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 84 insertions(+) diff --git a/model/channel_pricing.go b/model/channel_pricing.go index 6680de0..dbb0a11 100644 --- a/model/channel_pricing.go +++ b/model/channel_pricing.go @@ -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{} // 重置为零值 +}