|
- 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 // 按量计费
- 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, 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{} // 重置为零值
- }
|