Nelze vybrat více než 25 témat Téma musí začínat písmenem nebo číslem, může obsahovat pomlčky („-“) a může být dlouhé až 35 znaků.
 
 
 

208 řádky
6.4 KiB

  1. package model
  2. import (
  3. "fmt"
  4. "sync"
  5. "time"
  6. "github.com/QuantumNous/new-api/common"
  7. "gorm.io/gorm"
  8. "gorm.io/gorm/clause"
  9. )
  10. // 渠道定价缓存
  11. var (
  12. channelPricingCache = make(map[string]*ChannelPricing) // key: "modelName:channelId"
  13. channelPricingCacheLock sync.RWMutex
  14. channelPricingCacheTime time.Time
  15. channelPricingCacheTTL = time.Minute * 5 // 缓存5分钟
  16. )
  17. // QuotaType 计费类型
  18. const (
  19. QuotaTypeByTokens = 0 // 按量计费
  20. QuotaTypeByCall = 1 // 按次计费
  21. )
  22. // ChannelPricing 渠道定价表
  23. // 支持同一模型在不同渠道设置不同价格
  24. type ChannelPricing struct {
  25. Id int `json:"id" gorm:"primaryKey"`
  26. ModelName string `json:"model_name" gorm:"size:128;not null;uniqueIndex:idx_model_channel,priority:1"`
  27. ChannelId int `json:"channel_id" gorm:"not null;uniqueIndex:idx_model_channel,priority:2;index"`
  28. QuotaType int `json:"quota_type" gorm:"default:0"` // 0=按量, 1=按次
  29. ModelRatio float64 `json:"model_ratio" gorm:"default:0"`
  30. CompletionRatio float64 `json:"completion_ratio" gorm:"default:0"`
  31. ModelPrice float64 `json:"model_price" gorm:"default:0"`
  32. TagIds string `json:"tag_ids" gorm:"type:varchar(255)"` // 逗号分隔的标签ID
  33. CreatedTime int64 `json:"created_time" gorm:"bigint"`
  34. UpdatedTime int64 `json:"updated_time" gorm:"bigint"`
  35. DeletedAt gorm.DeletedAt `json:"-" gorm:"index"`
  36. }
  37. func (cp *ChannelPricing) Insert() error {
  38. now := common.GetTimestamp()
  39. cp.CreatedTime = now
  40. cp.UpdatedTime = now
  41. err := DB.Create(cp).Error
  42. if err == nil {
  43. InvalidateChannelPricingCache()
  44. }
  45. return err
  46. }
  47. func (cp *ChannelPricing) Update() error {
  48. cp.UpdatedTime = common.GetTimestamp()
  49. err := DB.Model(&ChannelPricing{}).Where("id = ?", cp.Id).Updates(map[string]interface{}{
  50. "quota_type": cp.QuotaType,
  51. "model_ratio": cp.ModelRatio,
  52. "completion_ratio": cp.CompletionRatio,
  53. "model_price": cp.ModelPrice,
  54. "tag_ids": cp.TagIds,
  55. "updated_time": cp.UpdatedTime,
  56. }).Error
  57. if err == nil {
  58. InvalidateChannelPricingCache()
  59. }
  60. return err
  61. }
  62. func (cp *ChannelPricing) Delete() error {
  63. err := DB.Delete(cp).Error
  64. if err == nil {
  65. InvalidateChannelPricingCache()
  66. }
  67. return err
  68. }
  69. // GetChannelPricing 获取指定模型在指定渠道的定价
  70. func GetChannelPricing(modelName string, channelId int) (*ChannelPricing, error) {
  71. var cp ChannelPricing
  72. err := DB.Where("model_name = ? AND channel_id = ?", modelName, channelId).First(&cp).Error
  73. if err != nil {
  74. return nil, err
  75. }
  76. return &cp, nil
  77. }
  78. // GetChannelPricingByModel 获取指定模型的所有渠道定价
  79. func GetChannelPricingByModel(modelName string) ([]*ChannelPricing, error) {
  80. var list []*ChannelPricing
  81. err := DB.Where("model_name = ?", modelName).Find(&list).Error
  82. return list, err
  83. }
  84. // GetAllChannelPricing 获取所有渠道定价(分页)
  85. func GetAllChannelPricing(offset int, limit int) ([]*ChannelPricing, int64, error) {
  86. var list []*ChannelPricing
  87. var total int64
  88. if err := DB.Model(&ChannelPricing{}).Count(&total).Error; err != nil {
  89. return nil, 0, err
  90. }
  91. err := DB.Order("id DESC").Offset(offset).Limit(limit).Find(&list).Error
  92. return list, total, err
  93. }
  94. // BatchUpsertChannelPricing 批量创建或更新渠道定价
  95. func BatchUpsertChannelPricing(pricings []*ChannelPricing) error {
  96. if len(pricings) == 0 {
  97. return nil
  98. }
  99. now := common.GetTimestamp()
  100. for _, cp := range pricings {
  101. cp.UpdatedTime = now
  102. // 仅在 CreatedTime 为空时设置(新记录)
  103. if cp.CreatedTime == 0 {
  104. cp.CreatedTime = now
  105. }
  106. }
  107. // 使用 GORM 的 OnConflict 实现 upsert
  108. // 唯一索引为 idx_model_channel (model_name, channel_id)
  109. return DB.Clauses(clause.OnConflict{
  110. Columns: []clause.Column{
  111. {Name: "model_name"},
  112. {Name: "channel_id"},
  113. },
  114. DoUpdates: clause.AssignmentColumns([]string{
  115. "quota_type",
  116. "model_ratio",
  117. "completion_ratio",
  118. "model_price",
  119. "tag_ids",
  120. "updated_time",
  121. }),
  122. }).Create(&pricings).Error
  123. }
  124. // getChannelPricingCacheKey 生成缓存键
  125. func getChannelPricingCacheKey(modelName string, channelId int) string {
  126. return fmt.Sprintf("%s:%d", modelName, channelId)
  127. }
  128. // GetEffectivePricing 获取有效定价(优先渠道定价,回退全局定价)
  129. // 返回: modelRatio, completionRatio, modelPrice, usePrice, found
  130. func GetEffectivePricing(modelName string, channelId int) (modelRatio, completionRatio, modelPrice float64, usePrice, found bool) {
  131. cacheKey := getChannelPricingCacheKey(modelName, channelId)
  132. // 首先检查缓存
  133. channelPricingCacheLock.RLock()
  134. // 检查缓存是否过期
  135. if time.Since(channelPricingCacheTime) < channelPricingCacheTTL {
  136. if cp, ok := channelPricingCache[cacheKey]; ok {
  137. channelPricingCacheLock.RUnlock()
  138. return cp.ModelRatio, cp.CompletionRatio, cp.ModelPrice, true, true
  139. }
  140. }
  141. channelPricingCacheLock.RUnlock()
  142. // 缓存未命中或已过期,查询数据库
  143. var cp ChannelPricing
  144. err := DB.Where("model_name = ? AND channel_id = ?", modelName, channelId).First(&cp).Error
  145. if err != nil {
  146. // 未找到渠道定价,返回 false 让调用者使用全局定价
  147. return 0, 0, 0, false, false
  148. }
  149. // 更新缓存
  150. channelPricingCacheLock.Lock()
  151. if channelPricingCacheTime.IsZero() || time.Since(channelPricingCacheTime) >= channelPricingCacheTTL {
  152. // 缓存过期,清空并更新时间
  153. channelPricingCache = make(map[string]*ChannelPricing)
  154. channelPricingCacheTime = time.Now()
  155. }
  156. channelPricingCache[cacheKey] = &cp
  157. channelPricingCacheLock.Unlock()
  158. return cp.ModelRatio, cp.CompletionRatio, cp.ModelPrice, true, true
  159. }
  160. // RefreshChannelPricingCache 刷新渠道定价缓存
  161. func RefreshChannelPricingCache() {
  162. channelPricingCacheLock.Lock()
  163. defer channelPricingCacheLock.Unlock()
  164. // 清空缓存
  165. channelPricingCache = make(map[string]*ChannelPricing)
  166. channelPricingCacheTime = time.Now()
  167. // 预加载所有渠道定价
  168. var pricings []*ChannelPricing
  169. if err := DB.Find(&pricings).Error; err != nil {
  170. return
  171. }
  172. for _, cp := range pricings {
  173. cacheKey := getChannelPricingCacheKey(cp.ModelName, cp.ChannelId)
  174. channelPricingCache[cacheKey] = cp
  175. }
  176. }
  177. // InvalidateChannelPricingCache 使渠道定价缓存失效
  178. func InvalidateChannelPricingCache() {
  179. channelPricingCacheLock.Lock()
  180. defer channelPricingCacheLock.Unlock()
  181. channelPricingCache = make(map[string]*ChannelPricing)
  182. channelPricingCacheTime = time.Time{} // 重置为零值
  183. }