Nie możesz wybrać więcej, niż 25 tematów Tematy muszą się zaczynać od litery lub cyfry, mogą zawierać myślniki ('-') i mogą mieć do 35 znaków.
 
 
 

189 wiersze
5.6 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. func GetVendorModelCounts() (map[int64]int64, error) {
  85. var stats []struct {
  86. VendorID int64
  87. Count int64
  88. }
  89. if err := DB.Model(&Model{}).
  90. Select("vendor_id as vendor_id, count(*) as count").
  91. Group("vendor_id").
  92. Scan(&stats).Error; err != nil {
  93. return nil, err
  94. }
  95. m := make(map[int64]int64, len(stats))
  96. for _, s := range stats {
  97. m[s.VendorID] = s.Count
  98. }
  99. return m, nil
  100. }
  101. func GetAllModels(offset int, limit int) ([]*Model, error) {
  102. var models []*Model
  103. err := DB.Order("sort_order ASC, id ASC").Offset(offset).Limit(limit).Find(&models).Error
  104. return models, err
  105. }
  106. func GetBoundChannelsByModelsMap(modelNames []string) (map[string][]BoundChannel, error) {
  107. result := make(map[string][]BoundChannel)
  108. if len(modelNames) == 0 {
  109. return result, nil
  110. }
  111. type row struct {
  112. Model string
  113. Name string
  114. Type int
  115. }
  116. var rows []row
  117. err := DB.Table("channels").
  118. Select("abilities.model as model, channels.name as name, channels.type as type").
  119. Joins("JOIN abilities ON abilities.channel_id = channels.id").
  120. Where("abilities.model IN ? AND abilities.enabled = ?", modelNames, true).
  121. Distinct().
  122. Scan(&rows).Error
  123. if err != nil {
  124. return nil, err
  125. }
  126. for _, r := range rows {
  127. result[r.Model] = append(result[r.Model], BoundChannel{Name: r.Name, Type: r.Type})
  128. }
  129. return result, nil
  130. }
  131. func SearchModels(keyword string, vendor string, offset int, limit int) ([]*Model, int64, error) {
  132. var models []*Model
  133. db := DB.Model(&Model{})
  134. if keyword != "" {
  135. like := "%" + keyword + "%"
  136. db = db.Where("model_name LIKE ? OR description LIKE ? OR tags LIKE ?", like, like, like)
  137. }
  138. if vendor != "" {
  139. if vid, err := strconv.Atoi(vendor); err == nil {
  140. db = db.Where("models.vendor_id = ?", vid)
  141. } else {
  142. db = db.Joins("JOIN vendors ON vendors.id = models.vendor_id").Where("vendors.name LIKE ?", "%"+vendor+"%")
  143. }
  144. }
  145. var total int64
  146. if err := db.Count(&total).Error; err != nil {
  147. return nil, 0, err
  148. }
  149. if err := db.Order("sort_order ASC, models.id ASC").Offset(offset).Limit(limit).Find(&models).Error; err != nil {
  150. return nil, 0, err
  151. }
  152. return models, total, nil
  153. }
  154. // ReorderModels 批量更新模型排序值
  155. func ReorderModels(items []struct {
  156. Id int `json:"id"`
  157. SortOrder int `json:"sort_order"`
  158. }) error {
  159. return DB.Transaction(func(tx *gorm.DB) error {
  160. for _, item := range items {
  161. if err := tx.Model(&Model{}).Where("id = ?", item.Id).Update("sort_order", item.SortOrder).Error; err != nil {
  162. return err
  163. }
  164. }
  165. return nil
  166. })
  167. }