Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

173 строки
5.1 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. CreatedTime int64 `json:"created_time" gorm:"bigint"`
  39. UpdatedTime int64 `json:"updated_time" gorm:"bigint"`
  40. DeletedAt gorm.DeletedAt `json:"-" gorm:"index;uniqueIndex:uk_model_name_delete_at,priority:2"`
  41. BoundChannels []BoundChannel `json:"bound_channels,omitempty" gorm:"-"`
  42. EnableGroups []string `json:"enable_groups,omitempty" gorm:"-"`
  43. QuotaTypes []int `json:"quota_types,omitempty" gorm:"-"`
  44. NameRule int `json:"name_rule" gorm:"default:0"`
  45. MatchedModels []string `json:"matched_models,omitempty" gorm:"-"`
  46. MatchedCount int `json:"matched_count,omitempty" gorm:"-"`
  47. }
  48. func (mi *Model) Insert() error {
  49. now := common.GetTimestamp()
  50. mi.CreatedTime = now
  51. mi.UpdatedTime = now
  52. // 保存原始值(因为 Create 后可能被 GORM 的 default 标签覆盖为 1)
  53. originalStatus := mi.Status
  54. originalSyncOfficial := mi.SyncOfficial
  55. // 先创建记录(GORM 会对零值字段应用默认值)
  56. if err := DB.Create(mi).Error; err != nil {
  57. return err
  58. }
  59. // 使用保存的原始值进行更新,确保零值能正确保存
  60. return DB.Model(&Model{}).Where("id = ?", mi.Id).Updates(map[string]interface{}{
  61. "status": originalStatus,
  62. "sync_official": originalSyncOfficial,
  63. }).Error
  64. }
  65. func IsModelNameDuplicated(id int, name string) (bool, error) {
  66. if name == "" {
  67. return false, nil
  68. }
  69. var cnt int64
  70. err := DB.Model(&Model{}).Where("model_name = ? AND id <> ?", name, id).Count(&cnt).Error
  71. return cnt > 0, err
  72. }
  73. func (mi *Model) Update() error {
  74. mi.UpdatedTime = common.GetTimestamp()
  75. // 使用 Select 强制更新所有字段,包括零值
  76. return DB.Model(&Model{}).Where("id = ?", mi.Id).
  77. Select("model_name", "description", "icon", "tags", "vendor_id", "endpoints", "status", "sync_official", "name_rule", "updated_time").
  78. Updates(mi).Error
  79. }
  80. func (mi *Model) Delete() error {
  81. return DB.Delete(mi).Error
  82. }
  83. func GetVendorModelCounts() (map[int64]int64, error) {
  84. var stats []struct {
  85. VendorID int64
  86. Count int64
  87. }
  88. if err := DB.Model(&Model{}).
  89. Select("vendor_id as vendor_id, count(*) as count").
  90. Group("vendor_id").
  91. Scan(&stats).Error; err != nil {
  92. return nil, err
  93. }
  94. m := make(map[int64]int64, len(stats))
  95. for _, s := range stats {
  96. m[s.VendorID] = s.Count
  97. }
  98. return m, nil
  99. }
  100. func GetAllModels(offset int, limit int) ([]*Model, error) {
  101. var models []*Model
  102. err := DB.Order("id DESC").Offset(offset).Limit(limit).Find(&models).Error
  103. return models, err
  104. }
  105. func GetBoundChannelsByModelsMap(modelNames []string) (map[string][]BoundChannel, error) {
  106. result := make(map[string][]BoundChannel)
  107. if len(modelNames) == 0 {
  108. return result, nil
  109. }
  110. type row struct {
  111. Model string
  112. Name string
  113. Type int
  114. }
  115. var rows []row
  116. err := DB.Table("channels").
  117. Select("abilities.model as model, channels.name as name, channels.type as type").
  118. Joins("JOIN abilities ON abilities.channel_id = channels.id").
  119. Where("abilities.model IN ? AND abilities.enabled = ?", modelNames, true).
  120. Distinct().
  121. Scan(&rows).Error
  122. if err != nil {
  123. return nil, err
  124. }
  125. for _, r := range rows {
  126. result[r.Model] = append(result[r.Model], BoundChannel{Name: r.Name, Type: r.Type})
  127. }
  128. return result, nil
  129. }
  130. func SearchModels(keyword string, vendor string, offset int, limit int) ([]*Model, int64, error) {
  131. var models []*Model
  132. db := DB.Model(&Model{})
  133. if keyword != "" {
  134. like := "%" + keyword + "%"
  135. db = db.Where("model_name LIKE ? OR description LIKE ? OR tags LIKE ?", like, like, like)
  136. }
  137. if vendor != "" {
  138. if vid, err := strconv.Atoi(vendor); err == nil {
  139. db = db.Where("models.vendor_id = ?", vid)
  140. } else {
  141. db = db.Joins("JOIN vendors ON vendors.id = models.vendor_id").Where("vendors.name LIKE ?", "%"+vendor+"%")
  142. }
  143. }
  144. var total int64
  145. if err := db.Count(&total).Error; err != nil {
  146. return nil, 0, err
  147. }
  148. if err := db.Order("models.id DESC").Offset(offset).Limit(limit).Find(&models).Error; err != nil {
  149. return nil, 0, err
  150. }
  151. return models, total, nil
  152. }