Nevar pievienot vairāk kā 25 tēmas Tēmai ir jāsākas ar burtu vai ciparu, tā var saturēt domu zīmes ('-') un var būt līdz 35 simboliem gara.
 
 
 

269 rindas
9.1 KiB

  1. package model
  2. import (
  3. "fmt"
  4. "sync"
  5. "time"
  6. "github.com/QuantumNous/new-api/common"
  7. "github.com/QuantumNous/new-api/setting/ratio_setting"
  8. "gorm.io/gorm"
  9. "gorm.io/gorm/clause"
  10. )
  11. // 渠道定价缓存
  12. var (
  13. channelPricingCache = make(map[string]*ChannelPricing) // key: "modelName:channelId"
  14. channelPricingCacheLock sync.RWMutex
  15. channelPricingCacheTime time.Time
  16. channelPricingCacheTTL = time.Minute * 5 // 缓存5分钟
  17. )
  18. // QuotaType 计费类型
  19. const (
  20. QuotaTypeByTokens = 0 // 按量计费
  21. QuotaTypeByCall = 1 // 按次计费
  22. )
  23. // ChannelPricing 渠道定价表
  24. // 支持同一模型在不同渠道设置不同价格
  25. type ChannelPricing struct {
  26. Id int `json:"id" gorm:"primaryKey"`
  27. ModelName string `json:"model_name" gorm:"size:128;not null;uniqueIndex:idx_model_channel,priority:1"`
  28. ChannelId int `json:"channel_id" gorm:"not null;uniqueIndex:idx_model_channel,priority:2;index"`
  29. QuotaType int `json:"quota_type" gorm:"default:0"` // 0=按量, 1=按次
  30. ModelRatio float64 `json:"model_ratio" gorm:"default:0"`
  31. CompletionRatio float64 `json:"completion_ratio" gorm:"default:0"`
  32. ModelPrice float64 `json:"model_price" gorm:"default:0"`
  33. TagIds string `json:"tag_ids" gorm:"type:varchar(255)"` // 逗号分隔的标签ID
  34. CreatedTime int64 `json:"created_time" gorm:"bigint"`
  35. UpdatedTime int64 `json:"updated_time" gorm:"bigint"`
  36. DeletedAt gorm.DeletedAt `json:"-" gorm:"index"`
  37. }
  38. func (cp *ChannelPricing) Insert() error {
  39. now := common.GetTimestamp()
  40. cp.CreatedTime = now
  41. cp.UpdatedTime = now
  42. err := DB.Create(cp).Error
  43. if err == nil {
  44. InvalidateChannelPricingCache()
  45. }
  46. return err
  47. }
  48. func (cp *ChannelPricing) Update() error {
  49. cp.UpdatedTime = common.GetTimestamp()
  50. err := DB.Model(&ChannelPricing{}).Where("id = ?", cp.Id).Updates(map[string]interface{}{
  51. "quota_type": cp.QuotaType,
  52. "model_ratio": cp.ModelRatio,
  53. "completion_ratio": cp.CompletionRatio,
  54. "model_price": cp.ModelPrice,
  55. "tag_ids": cp.TagIds,
  56. "updated_time": cp.UpdatedTime,
  57. }).Error
  58. if err == nil {
  59. InvalidateChannelPricingCache()
  60. }
  61. return err
  62. }
  63. func (cp *ChannelPricing) Delete() error {
  64. err := DB.Delete(cp).Error
  65. if err == nil {
  66. InvalidateChannelPricingCache()
  67. }
  68. return err
  69. }
  70. // GetChannelPricing 获取指定模型在指定渠道的定价
  71. func GetChannelPricing(modelName string, channelId int) (*ChannelPricing, error) {
  72. var cp ChannelPricing
  73. err := DB.Where("model_name = ? AND channel_id = ?", modelName, channelId).First(&cp).Error
  74. if err != nil {
  75. return nil, err
  76. }
  77. return &cp, nil
  78. }
  79. // GetChannelPricingByModel 获取指定模型的所有渠道定价
  80. func GetChannelPricingByModel(modelName string) ([]*ChannelPricing, error) {
  81. var list []*ChannelPricing
  82. err := DB.Where("model_name = ?", modelName).Find(&list).Error
  83. return list, err
  84. }
  85. // GetAllChannelPricing 获取所有渠道定价(分页)
  86. func GetAllChannelPricing(offset int, limit int) ([]*ChannelPricing, int64, error) {
  87. var list []*ChannelPricing
  88. var total int64
  89. if err := DB.Model(&ChannelPricing{}).Count(&total).Error; err != nil {
  90. return nil, 0, err
  91. }
  92. err := DB.Order("id DESC").Offset(offset).Limit(limit).Find(&list).Error
  93. return list, total, err
  94. }
  95. // BatchUpsertChannelPricing 批量创建或更新渠道定价
  96. func BatchUpsertChannelPricing(pricings []*ChannelPricing) error {
  97. if len(pricings) == 0 {
  98. return nil
  99. }
  100. now := common.GetTimestamp()
  101. for _, cp := range pricings {
  102. cp.UpdatedTime = now
  103. // 仅在 CreatedTime 为空时设置(新记录)
  104. if cp.CreatedTime == 0 {
  105. cp.CreatedTime = now
  106. }
  107. }
  108. // 使用 GORM 的 OnConflict 实现 upsert
  109. // 唯一索引为 idx_model_channel (model_name, channel_id)
  110. return DB.Clauses(clause.OnConflict{
  111. Columns: []clause.Column{
  112. {Name: "model_name"},
  113. {Name: "channel_id"},
  114. },
  115. DoUpdates: clause.AssignmentColumns([]string{
  116. "quota_type",
  117. "model_ratio",
  118. "completion_ratio",
  119. "model_price",
  120. "tag_ids",
  121. "updated_time",
  122. }),
  123. }).Create(&pricings).Error
  124. }
  125. // getChannelPricingCacheKey 生成缓存键
  126. func getChannelPricingCacheKey(modelName string, channelId int) string {
  127. return fmt.Sprintf("%s:%d", modelName, channelId)
  128. }
  129. // GetEffectivePricing 获取有效定价(优先渠道定价,回退全局定价)
  130. // 返回: modelRatio, completionRatio, modelPrice, usePrice, found
  131. func GetEffectivePricing(modelName string, channelId int) (modelRatio, completionRatio, modelPrice float64, usePrice, found bool) {
  132. cacheKey := getChannelPricingCacheKey(modelName, channelId)
  133. // 首先检查缓存
  134. channelPricingCacheLock.RLock()
  135. // 检查缓存是否过期
  136. if time.Since(channelPricingCacheTime) < channelPricingCacheTTL {
  137. if cp, ok := channelPricingCache[cacheKey]; ok {
  138. channelPricingCacheLock.RUnlock()
  139. return cp.ModelRatio, cp.CompletionRatio, cp.ModelPrice, cp.QuotaType == QuotaTypeByCall, true
  140. }
  141. }
  142. channelPricingCacheLock.RUnlock()
  143. // 缓存未命中或已过期,查询数据库
  144. var cp ChannelPricing
  145. err := DB.Where("model_name = ? AND channel_id = ?", modelName, channelId).First(&cp).Error
  146. if err != nil {
  147. // 未找到渠道定价,返回 false 让调用者使用全局定价
  148. return 0, 0, 0, false, false
  149. }
  150. // 更新缓存
  151. channelPricingCacheLock.Lock()
  152. if channelPricingCacheTime.IsZero() || time.Since(channelPricingCacheTime) >= channelPricingCacheTTL {
  153. // 缓存过期,清空并更新时间
  154. channelPricingCache = make(map[string]*ChannelPricing)
  155. channelPricingCacheTime = time.Now()
  156. }
  157. channelPricingCache[cacheKey] = &cp
  158. channelPricingCacheLock.Unlock()
  159. return cp.ModelRatio, cp.CompletionRatio, cp.ModelPrice, cp.QuotaType == QuotaTypeByCall, true
  160. }
  161. // RefreshChannelPricingCache 刷新渠道定价缓存
  162. func RefreshChannelPricingCache() {
  163. channelPricingCacheLock.Lock()
  164. defer channelPricingCacheLock.Unlock()
  165. // 清空缓存
  166. channelPricingCache = make(map[string]*ChannelPricing)
  167. channelPricingCacheTime = time.Now()
  168. // 预加载所有渠道定价
  169. var pricings []*ChannelPricing
  170. if err := DB.Find(&pricings).Error; err != nil {
  171. return
  172. }
  173. for _, cp := range pricings {
  174. cacheKey := getChannelPricingCacheKey(cp.ModelName, cp.ChannelId)
  175. channelPricingCache[cacheKey] = cp
  176. }
  177. }
  178. // InvalidateChannelPricingCache 使渠道定价缓存失效
  179. func InvalidateChannelPricingCache() {
  180. channelPricingCacheLock.Lock()
  181. defer channelPricingCacheLock.Unlock()
  182. channelPricingCache = make(map[string]*ChannelPricing)
  183. channelPricingCacheTime = time.Time{} // 重置为零值
  184. }
  185. // ChannelPricingWithChannel 带渠道信息的定价响应
  186. type ChannelPricingWithChannel struct {
  187. Id int `json:"id"`
  188. ChannelId int `json:"channel_id"`
  189. ChannelName string `json:"channel_name"`
  190. ChannelType int `json:"channel_type"`
  191. QuotaType int `json:"quota_type"` // 0=按量, 1=按次
  192. ModelRatio float64 `json:"model_ratio"`
  193. CompletionRatio float64 `json:"completion_ratio"`
  194. ModelPrice float64 `json:"model_price"`
  195. HasCustomPricing bool `json:"has_custom_pricing"` // 是否有自定义定价
  196. }
  197. // GetChannelPricingByModelWithChannelInfo 获取指定模型的渠道定价(带渠道信息)
  198. // 返回所有支持该模型的渠道,对于没有渠道定价的渠道使用全局默认价格
  199. func GetChannelPricingByModelWithChannelInfo(modelName string) ([]*ChannelPricingWithChannel, error) {
  200. var results []*ChannelPricingWithChannel
  201. // 获取全局默认价格
  202. globalModelRatio, hasRatio, _ := ratio_setting.GetModelRatio(modelName)
  203. globalModelPrice, hasPrice := ratio_setting.GetModelPrice(modelName, false)
  204. globalCompletionRatio := ratio_setting.GetCompletionRatio(modelName)
  205. // 确定默认计费类型
  206. var defaultQuotaType int
  207. if hasPrice {
  208. defaultQuotaType = QuotaTypeByCall
  209. } else {
  210. defaultQuotaType = QuotaTypeByTokens
  211. }
  212. // 如果没有全局价格,设置默认值
  213. if !hasRatio {
  214. globalModelRatio = 0
  215. }
  216. if !hasPrice {
  217. globalModelPrice = 0
  218. }
  219. // 查询所有支持该模型的渠道,左连接渠道定价表
  220. err := DB.Table("abilities").
  221. Select(`abilities.channel_id, channels.name as channel_name, channels.type as channel_type,
  222. COALESCE(channel_pricings.quota_type, ?) as quota_type,
  223. COALESCE(channel_pricings.model_ratio, ?) as model_ratio,
  224. COALESCE(channel_pricings.completion_ratio, ?) as completion_ratio,
  225. COALESCE(channel_pricings.model_price, ?) as model_price,
  226. channel_pricings.id as id,
  227. (channel_pricings.id IS NOT NULL) as has_custom_pricing`,
  228. defaultQuotaType, globalModelRatio, globalCompletionRatio, globalModelPrice).
  229. Joins("LEFT JOIN channels ON abilities.channel_id = channels.id").
  230. Joins("LEFT JOIN channel_pricings ON abilities.channel_id = channel_pricings.channel_id AND channel_pricings.model_name = ? AND channel_pricings.deleted_at IS NULL", modelName).
  231. Where("abilities.model = ?", modelName).
  232. Where("abilities.enabled = ?", true).
  233. Where("channels.status = ?", 1). // 只显示启用的渠道
  234. Group("abilities.channel_id"). // 去重(同一渠道可能有多个分组)
  235. Scan(&results).Error
  236. return results, err
  237. }