Du kannst nicht mehr als 25 Themen auswählen Themen müssen entweder mit einem Buchstaben oder einer Ziffer beginnen. Sie können Bindestriche („-“) enthalten und bis zu 35 Zeichen lang sein.
 
 
 

329 Zeilen
12 KiB

  1. package model
  2. import (
  3. "fmt"
  4. "strconv"
  5. "strings"
  6. "sync"
  7. "github.com/QuantumNous/new-api/common"
  8. "github.com/QuantumNous/new-api/setting/ratio_setting"
  9. "gorm.io/gorm"
  10. "gorm.io/gorm/clause"
  11. )
  12. // 渠道定价缓存
  13. var (
  14. channelPricingCache = make(map[string]*ChannelPricing) // key: "modelName:channelId"
  15. channelPricingCacheLock sync.RWMutex
  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. // === 新增字段(0 = 未设置,回退全局值) ===
  37. CacheRatio float64 `json:"cache_ratio" gorm:"default:0"`
  38. CacheCreationRatio float64 `json:"cache_creation_ratio" gorm:"default:0"`
  39. ImageRatio float64 `json:"image_ratio" gorm:"default:0"`
  40. AudioRatio float64 `json:"audio_ratio" gorm:"default:0"`
  41. AudioCompletionRatio float64 `json:"audio_completion_ratio" gorm:"default:0"`
  42. }
  43. // setCache 写穿透缓存
  44. func setCache(key string, cp *ChannelPricing) {
  45. channelPricingCacheLock.Lock()
  46. channelPricingCache[key] = cp
  47. channelPricingCacheLock.Unlock()
  48. }
  49. func removeCache(key string) {
  50. channelPricingCacheLock.Lock()
  51. delete(channelPricingCache, key)
  52. channelPricingCacheLock.Unlock()
  53. }
  54. // ApplyFields 批量设置定价字段(消除 controller 层的重复赋值)
  55. func (cp *ChannelPricing) ApplyFields(quotaType int, modelRatio, completionRatio, modelPrice float64, tagIds string, cacheRatio, cacheCreationRatio, imageRatio, audioRatio, audioCompletionRatio float64) {
  56. cp.QuotaType = quotaType
  57. cp.ModelRatio = modelRatio
  58. cp.CompletionRatio = completionRatio
  59. cp.ModelPrice = modelPrice
  60. cp.TagIds = tagIds
  61. cp.CacheRatio = cacheRatio
  62. cp.CacheCreationRatio = cacheCreationRatio
  63. cp.ImageRatio = imageRatio
  64. cp.AudioRatio = audioRatio
  65. cp.AudioCompletionRatio = audioCompletionRatio
  66. }
  67. func (cp *ChannelPricing) Insert() error {
  68. now := common.GetTimestamp()
  69. cp.CreatedTime = now
  70. cp.UpdatedTime = now
  71. err := DB.Create(cp).Error
  72. if err == nil {
  73. setCache(getChannelPricingCacheKey(cp.ModelName, cp.ChannelId), cp)
  74. }
  75. return err
  76. }
  77. func (cp *ChannelPricing) Update() error {
  78. cp.UpdatedTime = common.GetTimestamp()
  79. err := DB.Model(&ChannelPricing{}).Where("id = ?", cp.Id).
  80. Select("quota_type", "model_ratio", "completion_ratio", "model_price",
  81. "tag_ids", "cache_ratio", "cache_creation_ratio", "image_ratio",
  82. "audio_ratio", "audio_completion_ratio", "updated_time").
  83. Updates(cp).Error
  84. if err == nil {
  85. setCache(getChannelPricingCacheKey(cp.ModelName, cp.ChannelId), cp)
  86. }
  87. return err
  88. }
  89. func (cp *ChannelPricing) Delete() error {
  90. var existing ChannelPricing
  91. if err := DB.First(&existing, cp.Id).Error; err != nil {
  92. return err
  93. }
  94. err := DB.Delete(cp).Error
  95. if err == nil {
  96. removeCache(getChannelPricingCacheKey(existing.ModelName, existing.ChannelId))
  97. }
  98. return err
  99. }
  100. // GetChannelPricing 获取指定模型在指定渠道的定价
  101. func GetChannelPricing(modelName string, channelId int) (*ChannelPricing, error) {
  102. var cp ChannelPricing
  103. err := DB.Where("model_name = ? AND channel_id = ?", modelName, channelId).First(&cp).Error
  104. if err != nil {
  105. return nil, err
  106. }
  107. return &cp, nil
  108. }
  109. // GetChannelPricingByModel 获取指定模型的所有渠道定价
  110. func GetChannelPricingByModel(modelName string) ([]*ChannelPricing, error) {
  111. var list []*ChannelPricing
  112. err := DB.Where("model_name = ?", modelName).Find(&list).Error
  113. return list, err
  114. }
  115. // GetAllChannelPricing 获取所有渠道定价(分页)
  116. func GetAllChannelPricing(offset int, limit int) ([]*ChannelPricing, int64, error) {
  117. var list []*ChannelPricing
  118. var total int64
  119. if err := DB.Model(&ChannelPricing{}).Count(&total).Error; err != nil {
  120. return nil, 0, err
  121. }
  122. err := DB.Order("id DESC").Offset(offset).Limit(limit).Find(&list).Error
  123. return list, total, err
  124. }
  125. // BatchUpsertChannelPricing 批量创建或更新渠道定价
  126. func BatchUpsertChannelPricing(pricings []*ChannelPricing) error {
  127. if len(pricings) == 0 {
  128. return nil
  129. }
  130. now := common.GetTimestamp()
  131. for _, cp := range pricings {
  132. cp.UpdatedTime = now
  133. // 仅在 CreatedTime 为空时设置(新记录)
  134. if cp.CreatedTime == 0 {
  135. cp.CreatedTime = now
  136. }
  137. }
  138. // 使用 GORM 的 OnConflict 实现 upsert
  139. // 唯一索引为 idx_model_channel (model_name, channel_id)
  140. err := DB.Clauses(clause.OnConflict{
  141. Columns: []clause.Column{
  142. {Name: "model_name"},
  143. {Name: "channel_id"},
  144. },
  145. DoUpdates: clause.AssignmentColumns([]string{
  146. "quota_type",
  147. "model_ratio",
  148. "completion_ratio",
  149. "model_price",
  150. "tag_ids",
  151. "cache_ratio",
  152. "cache_creation_ratio",
  153. "image_ratio",
  154. "audio_ratio",
  155. "audio_completion_ratio",
  156. "updated_time",
  157. }),
  158. }).Create(&pricings).Error
  159. if err == nil {
  160. for _, cp := range pricings {
  161. setCache(getChannelPricingCacheKey(cp.ModelName, cp.ChannelId), cp)
  162. }
  163. }
  164. return err
  165. }
  166. // getChannelPricingCacheKey 生成缓存键
  167. func getChannelPricingCacheKey(modelName string, channelId int) string {
  168. return fmt.Sprintf("%s:%d", modelName, channelId)
  169. }
  170. // GetEffectivePricing 获取有效定价(纯内存查找)
  171. func GetEffectivePricing(modelName string, channelId int) (*ChannelPricing, bool) {
  172. key := getChannelPricingCacheKey(modelName, channelId)
  173. channelPricingCacheLock.RLock()
  174. cp, ok := channelPricingCache[key]
  175. channelPricingCacheLock.RUnlock()
  176. if !ok {
  177. return nil, false
  178. }
  179. return cp, true
  180. }
  181. // ParseTagIds 解析逗号分隔的标签ID字符串为 PricingTag 切片
  182. func ParseTagIds(tagIds string, tagMap map[int]*PricingTag) []*PricingTag {
  183. if tagIds == "" {
  184. return nil
  185. }
  186. tags := make([]*PricingTag, 0)
  187. for _, idStr := range strings.Split(tagIds, ",") {
  188. if id, err := strconv.Atoi(strings.TrimSpace(idStr)); err == nil {
  189. if tag, ok := tagMap[id]; ok {
  190. tags = append(tags, tag)
  191. }
  192. }
  193. }
  194. return tags
  195. }
  196. // LoadChannelPricingCache 全量加载渠道定价到内存(启动时调用)
  197. func LoadChannelPricingCache() {
  198. var pricings []*ChannelPricing
  199. if err := DB.Find(&pricings).Error; err != nil {
  200. common.SysError("[ChannelPricing] LoadChannelPricingCache failed: " + err.Error())
  201. return
  202. }
  203. channelPricingCacheLock.Lock()
  204. channelPricingCache = make(map[string]*ChannelPricing, len(pricings))
  205. for _, cp := range pricings {
  206. key := getChannelPricingCacheKey(cp.ModelName, cp.ChannelId)
  207. channelPricingCache[key] = cp
  208. }
  209. channelPricingCacheLock.Unlock()
  210. common.SysLog(fmt.Sprintf("[ChannelPricing] cache loaded %d records", len(pricings)))
  211. }
  212. // ChannelPricingWithChannel 带渠道信息的定价响应
  213. type ChannelPricingWithChannel struct {
  214. Id int `json:"id"`
  215. ChannelId int `json:"channel_id"`
  216. ChannelName string `json:"channel_name"`
  217. ChannelType int `json:"channel_type"`
  218. TagIds string `json:"tag_ids" gorm:"column:tag_ids"` // 渠道定价的标签ID列表(逗号分隔)
  219. Tags []*PricingTag `json:"tags" gorm:"-"` // 渠道定价的标签详情(不参与数据库扫描)
  220. QuotaType int `json:"quota_type"` // 0=按量, 1=按次
  221. ModelRatio float64 `json:"model_ratio"`
  222. CompletionRatio float64 `json:"completion_ratio"`
  223. ModelPrice float64 `json:"model_price"`
  224. HasCustomPricing bool `json:"has_custom_pricing"` // 是否有自定义定价
  225. CacheRatio float64 `json:"cache_ratio"`
  226. CacheCreationRatio float64 `json:"cache_creation_ratio"`
  227. ImageRatio float64 `json:"image_ratio"`
  228. AudioRatio float64 `json:"audio_ratio"`
  229. AudioCompletionRatio float64 `json:"audio_completion_ratio"`
  230. }
  231. // GetChannelPricingByModelWithChannelInfo 获取指定模型的渠道定价(带渠道信息)
  232. // 返回所有支持该模型的渠道,对于没有渠道定价的渠道使用全局默认价格
  233. func GetChannelPricingByModelWithChannelInfo(modelName string) ([]*ChannelPricingWithChannel, error) {
  234. var results []*ChannelPricingWithChannel
  235. // 获取全局默认价格
  236. globalModelRatio, hasRatio, _ := ratio_setting.GetModelRatio(modelName)
  237. globalModelPrice, hasPrice := ratio_setting.GetModelPrice(modelName, false)
  238. globalCompletionRatio := ratio_setting.GetCompletionRatio(modelName)
  239. // 确定默认计费类型
  240. var defaultQuotaType int
  241. if hasPrice {
  242. defaultQuotaType = QuotaTypeByCall
  243. } else {
  244. defaultQuotaType = QuotaTypeByTokens
  245. }
  246. // 如果没有全局价格,设置默认值
  247. if !hasRatio {
  248. globalModelRatio = 0
  249. }
  250. if !hasPrice {
  251. globalModelPrice = 0
  252. }
  253. // 获取全局扩展比率(用于 CASE WHEN 回退)
  254. globalCacheRatio, _ := ratio_setting.GetCacheRatio(modelName)
  255. globalCacheCreationRatio, _ := ratio_setting.GetCreateCacheRatio(modelName)
  256. globalImageRatio, _ := ratio_setting.GetImageRatio(modelName)
  257. globalAudioRatio := ratio_setting.GetAudioRatio(modelName)
  258. globalAudioCompletionRatio := ratio_setting.GetAudioCompletionRatio(modelName)
  259. // 查询所有支持该模型的渠道,左连接渠道定价表
  260. err := DB.Table("abilities").
  261. Select(`abilities.channel_id, channels.name as channel_name, channels.type as channel_type,
  262. COALESCE(channel_pricings.quota_type, ?) as quota_type,
  263. COALESCE(channel_pricings.model_ratio, ?) as model_ratio,
  264. COALESCE(channel_pricings.completion_ratio, ?) as completion_ratio,
  265. COALESCE(channel_pricings.model_price, ?) as model_price,
  266. channel_pricings.id as id,
  267. channel_pricings.tag_ids as tag_ids,
  268. CASE WHEN channel_pricings.cache_ratio > 0 THEN channel_pricings.cache_ratio ELSE ? END as cache_ratio,
  269. CASE WHEN channel_pricings.cache_creation_ratio > 0 THEN channel_pricings.cache_creation_ratio ELSE ? END as cache_creation_ratio,
  270. CASE WHEN channel_pricings.image_ratio > 0 THEN channel_pricings.image_ratio ELSE ? END as image_ratio,
  271. CASE WHEN channel_pricings.audio_ratio > 0 THEN channel_pricings.audio_ratio ELSE ? END as audio_ratio,
  272. CASE WHEN channel_pricings.audio_completion_ratio > 0 THEN channel_pricings.audio_completion_ratio ELSE ? END as audio_completion_ratio,
  273. (channel_pricings.id IS NOT NULL) as has_custom_pricing`,
  274. defaultQuotaType, globalModelRatio, globalCompletionRatio, globalModelPrice, globalCacheRatio, globalCacheCreationRatio, globalImageRatio, globalAudioRatio, globalAudioCompletionRatio).
  275. Joins("LEFT JOIN channels ON abilities.channel_id = channels.id").
  276. 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).
  277. Where("abilities.model = ?", modelName).
  278. Where("abilities.enabled = ?", true).
  279. Where("channels.status = ?", 1). // 只显示启用的渠道
  280. 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, channel_pricings.cache_ratio, channel_pricings.cache_creation_ratio, channel_pricings.image_ratio, channel_pricings.audio_ratio, channel_pricings.audio_completion_ratio").
  281. Scan(&results).Error
  282. if err != nil {
  283. return nil, err
  284. }
  285. // 获取所有定价标签
  286. allTags, err := GetAllPricingTags()
  287. if err != nil {
  288. return results, nil // 如果获取标签失败,仍然返回基础结果
  289. }
  290. // 建立标签ID到标签的映射
  291. tagMap := make(map[int]*PricingTag)
  292. for _, tag := range allTags {
  293. tagMap[tag.Id] = tag
  294. }
  295. // 为每个渠道定价填充标签
  296. for _, result := range results {
  297. result.Tags = ParseTagIds(result.TagIds, tagMap)
  298. }
  299. return results, nil
  300. }