選択できるのは25トピックまでです。 トピックは、先頭が英数字で、英数字とダッシュ('-')を使用した35文字以内のものにしてください。
 
 
 

451 行
16 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. // 默认通道缓存:modelName → channelId
  17. defaultChannelCache = make(map[string]int)
  18. defaultChannelCacheLock sync.RWMutex
  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. // === 新增字段(0 = 未设置,回退全局值) ===
  40. CacheRatio float64 `json:"cache_ratio" gorm:"default:0"`
  41. CacheCreationRatio float64 `json:"cache_creation_ratio" gorm:"default:0"`
  42. ImageRatio float64 `json:"image_ratio" gorm:"default:0"`
  43. AudioRatio float64 `json:"audio_ratio" gorm:"default:0"`
  44. AudioCompletionRatio float64 `json:"audio_completion_ratio" gorm:"default:0"`
  45. IsDefault bool `json:"is_default" gorm:"default:false;index"`
  46. }
  47. // setCache 写穿透缓存
  48. func setCache(key string, cp *ChannelPricing) {
  49. channelPricingCacheLock.Lock()
  50. channelPricingCache[key] = cp
  51. channelPricingCacheLock.Unlock()
  52. }
  53. func removeCache(key string) {
  54. channelPricingCacheLock.Lock()
  55. delete(channelPricingCache, key)
  56. channelPricingCacheLock.Unlock()
  57. }
  58. // ApplyFields 批量设置定价字段(消除 controller 层的重复赋值)
  59. func (cp *ChannelPricing) ApplyFields(quotaType int, modelRatio, completionRatio, modelPrice float64, tagIds string, cacheRatio, cacheCreationRatio, imageRatio, audioRatio, audioCompletionRatio float64) {
  60. cp.QuotaType = quotaType
  61. cp.ModelRatio = modelRatio
  62. cp.CompletionRatio = completionRatio
  63. cp.ModelPrice = modelPrice
  64. cp.TagIds = tagIds
  65. cp.CacheRatio = cacheRatio
  66. cp.CacheCreationRatio = cacheCreationRatio
  67. cp.ImageRatio = imageRatio
  68. cp.AudioRatio = audioRatio
  69. cp.AudioCompletionRatio = audioCompletionRatio
  70. }
  71. func (cp *ChannelPricing) Insert() error {
  72. now := common.GetTimestamp()
  73. cp.CreatedTime = now
  74. cp.UpdatedTime = now
  75. err := DB.Create(cp).Error
  76. if err == nil {
  77. setCache(getChannelPricingCacheKey(cp.ModelName, cp.ChannelId), cp)
  78. if cp.IsDefault {
  79. setDefaultChannelCache(cp.ModelName, cp.ChannelId)
  80. }
  81. }
  82. return err
  83. }
  84. func (cp *ChannelPricing) Update() error {
  85. cp.UpdatedTime = common.GetTimestamp()
  86. err := DB.Model(&ChannelPricing{}).Where("id = ?", cp.Id).
  87. Select("quota_type", "model_ratio", "completion_ratio", "model_price",
  88. "tag_ids", "cache_ratio", "cache_creation_ratio", "image_ratio",
  89. "audio_ratio", "audio_completion_ratio", "is_default", "updated_time").
  90. Updates(cp).Error
  91. if err == nil {
  92. setCache(getChannelPricingCacheKey(cp.ModelName, cp.ChannelId), cp)
  93. if cp.IsDefault {
  94. setDefaultChannelCache(cp.ModelName, cp.ChannelId)
  95. } else {
  96. clearDefaultChannelCacheIfMatch(cp.ModelName, cp.Id)
  97. }
  98. }
  99. return err
  100. }
  101. func (cp *ChannelPricing) Delete() error {
  102. var existing ChannelPricing
  103. if err := DB.First(&existing, cp.Id).Error; err != nil {
  104. return err
  105. }
  106. err := DB.Delete(cp).Error
  107. if err == nil {
  108. removeCache(getChannelPricingCacheKey(existing.ModelName, existing.ChannelId))
  109. if existing.IsDefault {
  110. clearDefaultChannelCache(existing.ModelName)
  111. }
  112. }
  113. return err
  114. }
  115. // GetChannelPricing 获取指定模型在指定渠道的定价
  116. func GetChannelPricing(modelName string, channelId int) (*ChannelPricing, error) {
  117. var cp ChannelPricing
  118. err := DB.Where("model_name = ? AND channel_id = ?", modelName, channelId).First(&cp).Error
  119. if err != nil {
  120. return nil, err
  121. }
  122. return &cp, nil
  123. }
  124. // GetChannelPricingByModel 获取指定模型的所有渠道定价
  125. func GetChannelPricingByModel(modelName string) ([]*ChannelPricing, error) {
  126. var list []*ChannelPricing
  127. err := DB.Where("model_name = ?", modelName).Find(&list).Error
  128. return list, err
  129. }
  130. // GetAllChannelPricing 获取所有渠道定价(分页)
  131. func GetAllChannelPricing(offset int, limit int) ([]*ChannelPricing, int64, error) {
  132. var list []*ChannelPricing
  133. var total int64
  134. if err := DB.Model(&ChannelPricing{}).Count(&total).Error; err != nil {
  135. return nil, 0, err
  136. }
  137. err := DB.Order("id DESC").Offset(offset).Limit(limit).Find(&list).Error
  138. return list, total, err
  139. }
  140. // BatchUpsertChannelPricing 批量创建或更新渠道定价
  141. func BatchUpsertChannelPricing(pricings []*ChannelPricing) error {
  142. if len(pricings) == 0 {
  143. return nil
  144. }
  145. now := common.GetTimestamp()
  146. for _, cp := range pricings {
  147. cp.UpdatedTime = now
  148. // 仅在 CreatedTime 为空时设置(新记录)
  149. if cp.CreatedTime == 0 {
  150. cp.CreatedTime = now
  151. }
  152. }
  153. // 使用 GORM 的 OnConflict 实现 upsert
  154. // 唯一索引为 idx_model_channel (model_name, channel_id)
  155. err := DB.Clauses(clause.OnConflict{
  156. Columns: []clause.Column{
  157. {Name: "model_name"},
  158. {Name: "channel_id"},
  159. },
  160. DoUpdates: clause.AssignmentColumns([]string{
  161. "quota_type",
  162. "model_ratio",
  163. "completion_ratio",
  164. "model_price",
  165. "tag_ids",
  166. "cache_ratio",
  167. "cache_creation_ratio",
  168. "image_ratio",
  169. "audio_ratio",
  170. "audio_completion_ratio",
  171. "updated_time",
  172. }),
  173. }).Create(&pricings).Error
  174. if err == nil {
  175. for _, cp := range pricings {
  176. setCache(getChannelPricingCacheKey(cp.ModelName, cp.ChannelId), cp)
  177. }
  178. }
  179. return err
  180. }
  181. // getChannelPricingCacheKey 生成缓存键
  182. func getChannelPricingCacheKey(modelName string, channelId int) string {
  183. return fmt.Sprintf("%s:%d", modelName, channelId)
  184. }
  185. // GetEffectivePricing 获取有效定价(纯内存查找)
  186. func GetEffectivePricing(modelName string, channelId int) (*ChannelPricing, bool) {
  187. key := getChannelPricingCacheKey(modelName, channelId)
  188. channelPricingCacheLock.RLock()
  189. cp, ok := channelPricingCache[key]
  190. channelPricingCacheLock.RUnlock()
  191. if !ok {
  192. return nil, false
  193. }
  194. return cp, true
  195. }
  196. // ParseTagIds 解析逗号分隔的标签ID字符串为 PricingTag 切片
  197. func ParseTagIds(tagIds string, tagMap map[int]*PricingTag) []*PricingTag {
  198. if tagIds == "" {
  199. return nil
  200. }
  201. tags := make([]*PricingTag, 0)
  202. for _, idStr := range strings.Split(tagIds, ",") {
  203. if id, err := strconv.Atoi(strings.TrimSpace(idStr)); err == nil {
  204. if tag, ok := tagMap[id]; ok {
  205. tags = append(tags, tag)
  206. }
  207. }
  208. }
  209. return tags
  210. }
  211. // LoadChannelPricingCache 全量加载渠道定价到内存(启动时调用)
  212. func LoadChannelPricingCache() {
  213. var pricings []*ChannelPricing
  214. if err := DB.Find(&pricings).Error; err != nil {
  215. common.SysError("[ChannelPricing] LoadChannelPricingCache failed: " + err.Error())
  216. return
  217. }
  218. channelPricingCacheLock.Lock()
  219. channelPricingCache = make(map[string]*ChannelPricing, len(pricings))
  220. for _, cp := range pricings {
  221. key := getChannelPricingCacheKey(cp.ModelName, cp.ChannelId)
  222. channelPricingCache[key] = cp
  223. }
  224. channelPricingCacheLock.Unlock()
  225. rebuildDefaultChannelCache(pricings)
  226. common.SysLog(fmt.Sprintf("[ChannelPricing] cache loaded %d records", len(pricings)))
  227. }
  228. // ChannelPricingWithChannel 带渠道信息的定价响应
  229. type ChannelPricingWithChannel struct {
  230. Id int `json:"id"`
  231. ChannelId int `json:"channel_id"`
  232. ChannelName string `json:"channel_name"`
  233. ChannelPublicName string `json:"channel_public_name"`
  234. ChannelType int `json:"channel_type"`
  235. TagIds string `json:"tag_ids" gorm:"column:tag_ids"` // 渠道定价的标签ID列表(逗号分隔)
  236. Tags []*PricingTag `json:"tags" gorm:"-"` // 渠道定价的标签详情(不参与数据库扫描)
  237. QuotaType int `json:"quota_type"` // 0=按量, 1=按次
  238. ModelRatio float64 `json:"model_ratio"`
  239. CompletionRatio float64 `json:"completion_ratio"`
  240. ModelPrice float64 `json:"model_price"`
  241. HasCustomPricing bool `json:"has_custom_pricing"` // 是否有自定义定价
  242. CacheRatio float64 `json:"cache_ratio"`
  243. CacheCreationRatio float64 `json:"cache_creation_ratio"`
  244. ImageRatio float64 `json:"image_ratio"`
  245. AudioRatio float64 `json:"audio_ratio"`
  246. AudioCompletionRatio float64 `json:"audio_completion_ratio"`
  247. IsDefault bool `json:"is_default"`
  248. UserRatio float64 `json:"user_ratio"`
  249. }
  250. // GetChannelPricingByModelWithChannelInfo 获取指定模型的渠道定价(带渠道信息)
  251. // 返回所有支持该模型的渠道,对于没有渠道定价的渠道使用全局默认价格
  252. func GetChannelPricingByModelWithChannelInfo(modelName string) ([]*ChannelPricingWithChannel, error) {
  253. var results []*ChannelPricingWithChannel
  254. // 获取全局默认价格
  255. globalModelRatio, hasRatio, _ := ratio_setting.GetModelRatio(modelName)
  256. globalModelPrice, hasPrice := ratio_setting.GetModelPrice(modelName, false)
  257. globalCompletionRatio := ratio_setting.GetCompletionRatio(modelName)
  258. // 确定默认计费类型
  259. var defaultQuotaType int
  260. if hasPrice {
  261. defaultQuotaType = QuotaTypeByCall
  262. } else {
  263. defaultQuotaType = QuotaTypeByTokens
  264. }
  265. // 如果没有全局价格,设置默认值
  266. if !hasRatio {
  267. globalModelRatio = 0
  268. }
  269. if !hasPrice {
  270. globalModelPrice = 0
  271. }
  272. // 查询所有支持该模型的渠道,左连接渠道定价表
  273. // 高级字段(cache/image/audio)不回退全局值,直接返回 0
  274. err := DB.Table("abilities").
  275. Select(`abilities.channel_id, channels.name as channel_name, channels.public_name as channel_public_name, channels.type as channel_type,
  276. COALESCE(channel_pricings.quota_type, ?) as quota_type,
  277. COALESCE(channel_pricings.model_ratio, ?) as model_ratio,
  278. COALESCE(channel_pricings.completion_ratio, ?) as completion_ratio,
  279. COALESCE(channel_pricings.model_price, ?) as model_price,
  280. channel_pricings.id as id,
  281. channel_pricings.tag_ids as tag_ids,
  282. COALESCE(channel_pricings.cache_ratio, 0) as cache_ratio,
  283. COALESCE(channel_pricings.cache_creation_ratio, 0) as cache_creation_ratio,
  284. COALESCE(channel_pricings.image_ratio, 0) as image_ratio,
  285. COALESCE(channel_pricings.audio_ratio, 0) as audio_ratio,
  286. COALESCE(channel_pricings.audio_completion_ratio, 0) as audio_completion_ratio,
  287. COALESCE(channel_pricings.is_default, false) as is_default,
  288. (channel_pricings.id IS NOT NULL) as has_custom_pricing`,
  289. defaultQuotaType, globalModelRatio, globalCompletionRatio, globalModelPrice).
  290. Joins("LEFT JOIN channels ON abilities.channel_id = channels.id").
  291. 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).
  292. Where("abilities.model = ?", modelName).
  293. Where("abilities.enabled = ?", true).
  294. Where("channels.status = ?", 1). // 只显示启用的渠道
  295. Group("abilities.channel_id, channels.name, channels.public_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, channel_pricings.is_default").
  296. Scan(&results).Error
  297. if err != nil {
  298. return nil, err
  299. }
  300. // 获取所有定价标签
  301. allTags, err := GetAllPricingTags()
  302. if err != nil {
  303. return results, nil // 如果获取标签失败,仍然返回基础结果
  304. }
  305. // 建立标签ID到标签的映射
  306. tagMap := make(map[int]*PricingTag)
  307. for _, tag := range allTags {
  308. tagMap[tag.Id] = tag
  309. }
  310. // 为每个渠道定价填充标签
  311. for _, result := range results {
  312. result.Tags = ParseTagIds(result.TagIds, tagMap)
  313. }
  314. return results, nil
  315. }
  316. // === 默认通道缓存 ===
  317. // syncIsDefaultToCache 将默认标记的变更同步到 channelPricingCache
  318. func syncIsDefaultToCache(modelName string, channelId int, isDefault bool) {
  319. channelPricingCacheLock.Lock()
  320. for key, cp := range channelPricingCache {
  321. if cp.ModelName == modelName {
  322. if isDefault {
  323. cp.IsDefault = cp.ChannelId == channelId
  324. } else {
  325. cp.IsDefault = false
  326. }
  327. }
  328. channelPricingCache[key] = cp
  329. }
  330. channelPricingCacheLock.Unlock()
  331. }
  332. // setDefaultChannelCache 设置默认通道缓存
  333. func setDefaultChannelCache(modelName string, channelId int) {
  334. defaultChannelCacheLock.Lock()
  335. defaultChannelCache[modelName] = channelId
  336. defaultChannelCacheLock.Unlock()
  337. }
  338. // clearDefaultChannelCache 清除指定模型的默认通道缓存
  339. func clearDefaultChannelCache(modelName string) {
  340. defaultChannelCacheLock.Lock()
  341. delete(defaultChannelCache, modelName)
  342. defaultChannelCacheLock.Unlock()
  343. }
  344. // clearDefaultChannelCacheIfMatch 如果默认通道的定价记录 ID 匹配则清除
  345. func clearDefaultChannelCacheIfMatch(modelName string, pricingId int) {
  346. defaultChannelCacheLock.RLock()
  347. cachedId, ok := defaultChannelCache[modelName]
  348. defaultChannelCacheLock.RUnlock()
  349. if !ok {
  350. return
  351. }
  352. // 需要通过缓存找到对应的 pricing 来比对
  353. key := getChannelPricingCacheKey(modelName, cachedId)
  354. channelPricingCacheLock.RLock()
  355. cp, exists := channelPricingCache[key]
  356. channelPricingCacheLock.RUnlock()
  357. if exists && cp.Id == pricingId {
  358. clearDefaultChannelCache(modelName)
  359. }
  360. }
  361. // rebuildDefaultChannelCache 从全量数据构建默认通道缓存(启动时调用)
  362. func rebuildDefaultChannelCache(pricings []*ChannelPricing) {
  363. defaultChannelCacheLock.Lock()
  364. defaultChannelCache = make(map[string]int)
  365. for _, cp := range pricings {
  366. if cp.IsDefault {
  367. defaultChannelCache[cp.ModelName] = cp.ChannelId
  368. }
  369. }
  370. defaultChannelCacheLock.Unlock()
  371. common.SysLog(fmt.Sprintf("[ChannelPricing] default channel cache loaded %d records", len(defaultChannelCache)))
  372. }
  373. // GetDefaultChannelId 获取指定模型的默认通道 ID(纯内存读)
  374. func GetDefaultChannelId(modelName string) (int, bool) {
  375. defaultChannelCacheLock.RLock()
  376. id, ok := defaultChannelCache[modelName]
  377. defaultChannelCacheLock.RUnlock()
  378. return id, ok
  379. }
  380. // SetDefaultChannel 设置指定模型的默认通道(事务保证互斥)
  381. func SetDefaultChannel(modelName string, channelId int) error {
  382. return DB.Transaction(func(tx *gorm.DB) error {
  383. // 清除该模型所有现有的默认标记
  384. if err := tx.Model(&ChannelPricing{}).
  385. Where("model_name = ? AND is_default = ?", modelName, true).
  386. Update("is_default", false).Error; err != nil {
  387. return err
  388. }
  389. // 设置新的默认
  390. if err := tx.Model(&ChannelPricing{}).
  391. Where("model_name = ? AND channel_id = ?", modelName, channelId).
  392. Update("is_default", true).Error; err != nil {
  393. return err
  394. }
  395. // 更新缓存
  396. setDefaultChannelCache(modelName, channelId)
  397. syncIsDefaultToCache(modelName, channelId, true)
  398. return nil
  399. })
  400. }
  401. // ClearDefaultChannel 清除指定模型的默认通道标记
  402. func ClearDefaultChannel(modelName string) error {
  403. err := DB.Model(&ChannelPricing{}).
  404. Where("model_name = ? AND is_default = ?", modelName, true).
  405. Update("is_default", false).Error
  406. if err == nil {
  407. clearDefaultChannelCache(modelName)
  408. syncIsDefaultToCache(modelName, 0, false)
  409. }
  410. return err
  411. }