You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

199 lines
5.9 KiB

  1. package model
  2. import (
  3. "strconv"
  4. "github.com/QuantumNous/new-api/common"
  5. "gorm.io/gorm"
  6. )
  7. const (
  8. NameRuleExact = iota
  9. NameRulePrefix
  10. NameRuleContains
  11. NameRuleSuffix
  12. )
  13. const (
  14. ModelTypeChat = 1 + iota
  15. ModelTypeImage
  16. ModelTypeAudio
  17. ModelTypeVideo
  18. ModelTypeEmbedding
  19. ModelTypeRerank
  20. ModelTypeVision
  21. ModelTypeOther
  22. )
  23. type BoundChannel struct {
  24. Name string `json:"name"`
  25. Type int `json:"type"`
  26. }
  27. type Model struct {
  28. Id int `json:"id"`
  29. ModelName string `json:"model_name" gorm:"size:128;not null;uniqueIndex:uk_model_name_delete_at,priority:1"`
  30. Description string `json:"description,omitempty" gorm:"type:text"`
  31. Icon string `json:"icon,omitempty" gorm:"type:varchar(128)"`
  32. Tags string `json:"tags,omitempty" gorm:"type:varchar(255)"`
  33. Type int `json:"type" gorm:"default:0"`
  34. VendorID int `json:"vendor_id,omitempty" gorm:"index"`
  35. Endpoints string `json:"endpoints,omitempty" gorm:"type:text"`
  36. Status int `json:"status" gorm:"default:1"`
  37. SyncOfficial int `json:"sync_official" gorm:"default:1"`
  38. SortOrder int `json:"sort_order" gorm:"default:999999"`
  39. CreatedTime int64 `json:"created_time" gorm:"bigint"`
  40. UpdatedTime int64 `json:"updated_time" gorm:"bigint"`
  41. DeletedAt gorm.DeletedAt `json:"-" gorm:"index;uniqueIndex:uk_model_name_delete_at,priority:2"`
  42. BoundChannels []BoundChannel `json:"bound_channels,omitempty" gorm:"-"`
  43. EnableGroups []string `json:"enable_groups,omitempty" gorm:"-"`
  44. QuotaTypes []int `json:"quota_types,omitempty" gorm:"-"`
  45. NameRule int `json:"name_rule" gorm:"default:0"`
  46. MatchedModels []string `json:"matched_models,omitempty" gorm:"-"`
  47. MatchedCount int `json:"matched_count,omitempty" gorm:"-"`
  48. }
  49. func (mi *Model) Insert() error {
  50. now := common.GetTimestamp()
  51. mi.CreatedTime = now
  52. mi.UpdatedTime = now
  53. // 保存原始值(因为 Create 后可能被 GORM 的 default 标签覆盖为 1)
  54. originalStatus := mi.Status
  55. originalSyncOfficial := mi.SyncOfficial
  56. // 先创建记录(GORM 会对零值字段应用默认值)
  57. if err := DB.Create(mi).Error; err != nil {
  58. return err
  59. }
  60. // 使用保存的原始值进行更新,确保零值能正确保存
  61. return DB.Model(&Model{}).Where("id = ?", mi.Id).Updates(map[string]interface{}{
  62. "status": originalStatus,
  63. "sync_official": originalSyncOfficial,
  64. }).Error
  65. }
  66. func IsModelNameDuplicated(id int, name string) (bool, error) {
  67. if name == "" {
  68. return false, nil
  69. }
  70. var cnt int64
  71. err := DB.Model(&Model{}).Where("model_name = ? AND id <> ?", name, id).Count(&cnt).Error
  72. return cnt > 0, err
  73. }
  74. func (mi *Model) Update() error {
  75. mi.UpdatedTime = common.GetTimestamp()
  76. // 使用 Select 强制更新所有字段,包括零值
  77. return DB.Model(&Model{}).Where("id = ?", mi.Id).
  78. Select("model_name", "description", "icon", "tags", "type", "vendor_id", "endpoints", "status", "sync_official", "name_rule", "sort_order", "updated_time").
  79. Updates(mi).Error
  80. }
  81. func (mi *Model) Delete() error {
  82. return DB.Delete(mi).Error
  83. }
  84. // GetDisabledModelNames returns model names that are disabled (status = 0) from the given list.
  85. func GetDisabledModelNames(names []string) []string {
  86. if len(names) == 0 {
  87. return nil
  88. }
  89. var disabled []string
  90. DB.Table("models").Where("model_name IN ? AND status = 0", names).Pluck("model_name", &disabled)
  91. return disabled
  92. }
  93. func GetVendorModelCounts() (map[int64]int64, error) {
  94. var stats []struct {
  95. VendorID int64
  96. Count int64
  97. }
  98. if err := DB.Model(&Model{}).
  99. Select("vendor_id as vendor_id, count(*) as count").
  100. Group("vendor_id").
  101. Scan(&stats).Error; err != nil {
  102. return nil, err
  103. }
  104. m := make(map[int64]int64, len(stats))
  105. for _, s := range stats {
  106. m[s.VendorID] = s.Count
  107. }
  108. return m, nil
  109. }
  110. func GetAllModels(offset int, limit int) ([]*Model, error) {
  111. var models []*Model
  112. err := DB.Order("sort_order ASC, id ASC").Offset(offset).Limit(limit).Find(&models).Error
  113. return models, err
  114. }
  115. func GetBoundChannelsByModelsMap(modelNames []string) (map[string][]BoundChannel, error) {
  116. result := make(map[string][]BoundChannel)
  117. if len(modelNames) == 0 {
  118. return result, nil
  119. }
  120. type row struct {
  121. Model string
  122. Name string
  123. Type int
  124. }
  125. var rows []row
  126. err := DB.Table("channels").
  127. Select("abilities.model as model, channels.name as name, channels.type as type").
  128. Joins("JOIN abilities ON abilities.channel_id = channels.id").
  129. Where("abilities.model IN ? AND abilities.enabled = ?", modelNames, true).
  130. Distinct().
  131. Scan(&rows).Error
  132. if err != nil {
  133. return nil, err
  134. }
  135. for _, r := range rows {
  136. result[r.Model] = append(result[r.Model], BoundChannel{Name: r.Name, Type: r.Type})
  137. }
  138. return result, nil
  139. }
  140. func SearchModels(keyword string, vendor string, offset int, limit int) ([]*Model, int64, error) {
  141. var models []*Model
  142. db := DB.Model(&Model{})
  143. if keyword != "" {
  144. like := "%" + keyword + "%"
  145. db = db.Where("model_name LIKE ? OR description LIKE ? OR tags LIKE ?", like, like, like)
  146. }
  147. if vendor != "" {
  148. if vid, err := strconv.Atoi(vendor); err == nil {
  149. db = db.Where("models.vendor_id = ?", vid)
  150. } else {
  151. db = db.Joins("JOIN vendors ON vendors.id = models.vendor_id").Where("vendors.name LIKE ?", "%"+vendor+"%")
  152. }
  153. }
  154. var total int64
  155. if err := db.Count(&total).Error; err != nil {
  156. return nil, 0, err
  157. }
  158. if err := db.Order("sort_order ASC, models.id ASC").Offset(offset).Limit(limit).Find(&models).Error; err != nil {
  159. return nil, 0, err
  160. }
  161. return models, total, nil
  162. }
  163. // ReorderModels 批量更新模型排序值
  164. func ReorderModels(items []struct {
  165. Id int `json:"id"`
  166. SortOrder int `json:"sort_order"`
  167. }) error {
  168. return DB.Transaction(func(tx *gorm.DB) error {
  169. for _, item := range items {
  170. if err := tx.Model(&Model{}).Where("id = ?", item.Id).Update("sort_order", item.SortOrder).Error; err != nil {
  171. return err
  172. }
  173. }
  174. return nil
  175. })
  176. }