Преглед изворни кода

feat: add channel pricing cache and GetEffectivePricing function

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
feat/alipay-payment
fengsilin пре 1 месец
родитељ
комит
6c9ee5ec2b
1 измењених фајлова са 84 додато и 0 уклоњено
  1. +84
    -0
      model/channel_pricing.go

+ 84
- 0
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{} // 重置为零值
}

Loading…
Откажи
Сачувај