25'ten fazla konu seçemezsiniz Konular bir harf veya rakamla başlamalı, kısa çizgiler ('-') içerebilir ve en fazla 35 karakter uzunluğunda olabilir.
 
 
 

390 satır
10 KiB

  1. package model
  2. import (
  3. "errors"
  4. "fmt"
  5. "strings"
  6. "sync"
  7. "github.com/QuantumNous/new-api/common"
  8. "github.com/samber/lo"
  9. "gorm.io/gorm"
  10. "gorm.io/gorm/clause"
  11. )
  12. type Ability struct {
  13. Group string `json:"group" gorm:"type:varchar(64);primaryKey;autoIncrement:false"`
  14. Model string `json:"model" gorm:"type:varchar(255);primaryKey;autoIncrement:false"`
  15. ChannelId int `json:"channel_id" gorm:"primaryKey;autoIncrement:false;index"`
  16. Enabled bool `json:"enabled"`
  17. Priority *int64 `json:"priority" gorm:"bigint;default:0;index"`
  18. Weight uint `json:"weight" gorm:"default:0;index"`
  19. Tag *string `json:"tag" gorm:"index"`
  20. }
  21. type AbilityWithChannel struct {
  22. Ability
  23. ChannelType int `json:"channel_type"`
  24. }
  25. func GetAllEnableAbilityWithChannels() ([]AbilityWithChannel, error) {
  26. var abilities []AbilityWithChannel
  27. err := DB.Table("abilities").
  28. Select("abilities.*, channels.type as channel_type").
  29. Joins("left join channels on abilities.channel_id = channels.id").
  30. Joins("left join models on abilities.model = models.model_name").
  31. Where("abilities.enabled = ? and models.model_name is NOT null", true).
  32. Scan(&abilities).Error
  33. return abilities, err
  34. }
  35. func GetGroupEnabledModels(group string) []string {
  36. var models []string
  37. // Find distinct models
  38. DB.Table("abilities").Where(commonGroupCol+" = ? and enabled = ?", group, true).Distinct("model").Pluck("model", &models)
  39. return models
  40. }
  41. func GetEnabledModels() []string {
  42. var models []string
  43. // Find distinct models
  44. DB.Table("abilities").Where("enabled = ?", true).Distinct("model").Pluck("model", &models)
  45. return models
  46. }
  47. func GetAllEnableAbilities() []Ability {
  48. var abilities []Ability
  49. DB.Find(&abilities, "enabled = ?", true)
  50. return abilities
  51. }
  52. func getPriority(group string, model string, retry int) (int, error) {
  53. var priorities []int
  54. err := DB.Model(&Ability{}).
  55. Select("DISTINCT(priority)").
  56. Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, model, true).
  57. Order("priority DESC"). // 按优先级降序排序
  58. Pluck("priority", &priorities).Error // Pluck用于将查询的结果直接扫描到一个切片中
  59. if err != nil {
  60. // 处理错误
  61. return 0, err
  62. }
  63. if len(priorities) == 0 {
  64. // 如果没有查询到优先级,则返回错误
  65. return 0, errors.New("数据库一致性被破坏")
  66. }
  67. // 确定要使用的优先级
  68. var priorityToUse int
  69. if retry >= len(priorities) {
  70. // 如果重试次数大于优先级数,则使用最小的优先级
  71. priorityToUse = priorities[len(priorities)-1]
  72. } else {
  73. priorityToUse = priorities[retry]
  74. }
  75. return priorityToUse, nil
  76. }
  77. func getChannelQuery(group string, model string, retry int) (*gorm.DB, error) {
  78. maxPrioritySubQuery := DB.Model(&Ability{}).Select("MAX(priority)").Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, model, true)
  79. channelQuery := DB.Where(commonGroupCol+" = ? and model = ? and enabled = ? and priority = (?)", group, model, true, maxPrioritySubQuery)
  80. if retry != 0 {
  81. priority, err := getPriority(group, model, retry)
  82. if err != nil {
  83. return nil, err
  84. } else {
  85. channelQuery = DB.Where(commonGroupCol+" = ? and model = ? and enabled = ? and priority = ?", group, model, true, priority)
  86. }
  87. }
  88. return channelQuery, nil
  89. }
  90. func GetChannel(group string, model string, retry int) (*Channel, error) {
  91. var abilities []Ability
  92. var err error = nil
  93. channelQuery, err := getChannelQuery(group, model, retry)
  94. if err != nil {
  95. return nil, err
  96. }
  97. if common.UsingSQLite || common.UsingPostgreSQL {
  98. err = channelQuery.Order("weight DESC").Find(&abilities).Error
  99. } else {
  100. err = channelQuery.Order("weight DESC").Find(&abilities).Error
  101. }
  102. if err != nil {
  103. return nil, err
  104. }
  105. channel := Channel{}
  106. if len(abilities) > 0 {
  107. // Randomly choose one
  108. weightSum := uint(0)
  109. for _, ability_ := range abilities {
  110. weightSum += ability_.Weight + 10
  111. }
  112. // Randomly choose one
  113. weight := common.GetRandomInt(int(weightSum))
  114. for _, ability_ := range abilities {
  115. weight -= int(ability_.Weight) + 10
  116. //log.Printf("weight: %d, ability weight: %d", weight, *ability_.Weight)
  117. if weight <= 0 {
  118. channel.Id = ability_.ChannelId
  119. break
  120. }
  121. }
  122. } else {
  123. return nil, nil
  124. }
  125. err = DB.First(&channel, "id = ?", channel.Id).Error
  126. return &channel, err
  127. }
  128. func (channel *Channel) AddAbilities(tx *gorm.DB) error {
  129. models_ := strings.Split(channel.Models, ",")
  130. // 统一使用 useDB,避免重复的数据库连接选择
  131. useDB := DB
  132. if tx != nil {
  133. useDB = tx
  134. }
  135. // === 自动同步模型到 models 表(批量处理)===
  136. // 收集需要同步的模型名称
  137. modelNames := make([]string, 0, len(models_))
  138. for _, modelName := range models_ {
  139. modelName = strings.TrimSpace(modelName)
  140. if modelName != "" {
  141. modelNames = append(modelNames, modelName)
  142. }
  143. }
  144. // 批量查询已存在的模型
  145. existingModels := make(map[string]bool)
  146. if len(modelNames) > 0 {
  147. var existingModelList []Model
  148. useDB.Where("model_name IN ?", modelNames).Select("model_name").Find(&existingModelList)
  149. for _, m := range existingModelList {
  150. existingModels[m.ModelName] = true
  151. }
  152. }
  153. // 收集需要创建的模型
  154. now := common.GetTimestamp()
  155. newModels := make([]Model, 0, len(modelNames))
  156. for _, modelName := range modelNames {
  157. if !existingModels[modelName] {
  158. newModels = append(newModels, Model{
  159. ModelName: modelName,
  160. Status: 1, // 启用状态
  161. SyncOfficial: 0, // 标记为用户模型,不被官方同步覆盖
  162. CreatedTime: now,
  163. UpdatedTime: now,
  164. })
  165. }
  166. }
  167. // 批量插入新模型(使用 OnConflict 忽略重复)
  168. if len(newModels) > 0 {
  169. for _, chunk := range lo.Chunk(newModels, 50) {
  170. if err := useDB.Clauses(clause.OnConflict{DoNothing: true}).Create(&chunk).Error; err != nil {
  171. return err
  172. }
  173. }
  174. }
  175. // === 自动同步结束 ===
  176. groups_ := strings.Split(channel.Group, ",")
  177. abilitySet := make(map[string]struct{})
  178. abilities := make([]Ability, 0, len(models_))
  179. for _, model := range models_ {
  180. for _, group := range groups_ {
  181. key := group + "|" + model
  182. if _, exists := abilitySet[key]; exists {
  183. continue
  184. }
  185. abilitySet[key] = struct{}{}
  186. ability := Ability{
  187. Group: group,
  188. Model: model,
  189. ChannelId: channel.Id,
  190. Enabled: channel.Status == common.ChannelStatusEnabled,
  191. Priority: channel.Priority,
  192. Weight: uint(channel.GetWeight()),
  193. Tag: channel.Tag,
  194. }
  195. abilities = append(abilities, ability)
  196. }
  197. }
  198. if len(abilities) == 0 {
  199. return nil
  200. }
  201. for _, chunk := range lo.Chunk(abilities, 50) {
  202. err := useDB.Clauses(clause.OnConflict{DoNothing: true}).Create(&chunk).Error
  203. if err != nil {
  204. return err
  205. }
  206. }
  207. return nil
  208. }
  209. func (channel *Channel) DeleteAbilities() error {
  210. return DB.Where("channel_id = ?", channel.Id).Delete(&Ability{}).Error
  211. }
  212. // UpdateAbilities updates abilities of this channel.
  213. // Make sure the channel is completed before calling this function.
  214. func (channel *Channel) UpdateAbilities(tx *gorm.DB) error {
  215. isNewTx := false
  216. // 如果没有传入事务,创建新的事务
  217. if tx == nil {
  218. tx = DB.Begin()
  219. if tx.Error != nil {
  220. return tx.Error
  221. }
  222. isNewTx = true
  223. defer func() {
  224. if r := recover(); r != nil {
  225. tx.Rollback()
  226. }
  227. }()
  228. }
  229. // First delete all abilities of this channel
  230. err := tx.Where("channel_id = ?", channel.Id).Delete(&Ability{}).Error
  231. if err != nil {
  232. if isNewTx {
  233. tx.Rollback()
  234. }
  235. return err
  236. }
  237. // Then add new abilities
  238. models_ := strings.Split(channel.Models, ",")
  239. groups_ := strings.Split(channel.Group, ",")
  240. abilitySet := make(map[string]struct{})
  241. abilities := make([]Ability, 0, len(models_))
  242. for _, model := range models_ {
  243. for _, group := range groups_ {
  244. key := group + "|" + model
  245. if _, exists := abilitySet[key]; exists {
  246. continue
  247. }
  248. abilitySet[key] = struct{}{}
  249. ability := Ability{
  250. Group: group,
  251. Model: model,
  252. ChannelId: channel.Id,
  253. Enabled: channel.Status == common.ChannelStatusEnabled,
  254. Priority: channel.Priority,
  255. Weight: uint(channel.GetWeight()),
  256. Tag: channel.Tag,
  257. }
  258. abilities = append(abilities, ability)
  259. }
  260. }
  261. if len(abilities) > 0 {
  262. for _, chunk := range lo.Chunk(abilities, 50) {
  263. err = tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&chunk).Error
  264. if err != nil {
  265. if isNewTx {
  266. tx.Rollback()
  267. }
  268. return err
  269. }
  270. }
  271. }
  272. // 如果是新创建的事务,需要提交
  273. if isNewTx {
  274. return tx.Commit().Error
  275. }
  276. return nil
  277. }
  278. func UpdateAbilityStatus(channelId int, status bool) error {
  279. return DB.Model(&Ability{}).Where("channel_id = ?", channelId).Select("enabled").Update("enabled", status).Error
  280. }
  281. func UpdateAbilityStatusByTag(tag string, status bool) error {
  282. return DB.Model(&Ability{}).Where("tag = ?", tag).Select("enabled").Update("enabled", status).Error
  283. }
  284. func UpdateAbilityByTag(tag string, newTag *string, priority *int64, weight *uint) error {
  285. ability := Ability{}
  286. if newTag != nil {
  287. ability.Tag = newTag
  288. }
  289. if priority != nil {
  290. ability.Priority = priority
  291. }
  292. if weight != nil {
  293. ability.Weight = *weight
  294. }
  295. return DB.Model(&Ability{}).Where("tag = ?", tag).Updates(ability).Error
  296. }
  297. var fixLock = sync.Mutex{}
  298. func FixAbility() (int, int, error) {
  299. lock := fixLock.TryLock()
  300. if !lock {
  301. return 0, 0, errors.New("已经有一个修复任务在运行中,请稍后再试")
  302. }
  303. defer fixLock.Unlock()
  304. // truncate abilities table
  305. if common.UsingSQLite {
  306. err := DB.Exec("DELETE FROM abilities").Error
  307. if err != nil {
  308. common.SysLog(fmt.Sprintf("Delete abilities failed: %s", err.Error()))
  309. return 0, 0, err
  310. }
  311. } else {
  312. err := DB.Exec("TRUNCATE TABLE abilities").Error
  313. if err != nil {
  314. common.SysLog(fmt.Sprintf("Truncate abilities failed: %s", err.Error()))
  315. return 0, 0, err
  316. }
  317. }
  318. var channels []*Channel
  319. // Find all channels
  320. err := DB.Model(&Channel{}).Find(&channels).Error
  321. if err != nil {
  322. return 0, 0, err
  323. }
  324. if len(channels) == 0 {
  325. return 0, 0, nil
  326. }
  327. successCount := 0
  328. failCount := 0
  329. for _, chunk := range lo.Chunk(channels, 50) {
  330. ids := lo.Map(chunk, func(c *Channel, _ int) int { return c.Id })
  331. // Delete all abilities of this channel
  332. err = DB.Where("channel_id IN ?", ids).Delete(&Ability{}).Error
  333. if err != nil {
  334. common.SysLog(fmt.Sprintf("Delete abilities failed: %s", err.Error()))
  335. failCount += len(chunk)
  336. continue
  337. }
  338. // Then add new abilities
  339. for _, channel := range chunk {
  340. err = channel.AddAbilities(nil)
  341. if err != nil {
  342. common.SysLog(fmt.Sprintf("Add abilities for channel %d failed: %s", channel.Id, err.Error()))
  343. failCount++
  344. } else {
  345. successCount++
  346. }
  347. }
  348. }
  349. InitChannelCache()
  350. return successCount, failCount, nil
  351. }