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ů.
 
 
 

303 řádky
10 KiB

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