Ви не можете вибрати більше 25 тем Теми мають розпочинатися з літери або цифри, можуть містити дефіси (-) і не повинні перевищувати 35 символів.
 
 
 

450 рядки
12 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. // GetAbilitiesByChannelId 获取指定渠道的所有能力
  53. func GetAbilitiesByChannelId(channelId int) ([]*Ability, error) {
  54. var abilities []*Ability
  55. err := DB.Where("channel_id = ?", channelId).Find(&abilities).Error
  56. return abilities, err
  57. }
  58. // GetModelChannelsForGroup 返回指定模型在指定分组下的可用渠道列表及默认渠道ID
  59. func GetModelChannelsForGroup(modelName string, group string) ([]map[string]any, int, error) {
  60. var channelIds []int
  61. err := DB.Model(&Ability{}).
  62. Where("model = ?", modelName).
  63. Where("enabled = ?", true).
  64. Where(commonGroupCol+" = ?", group).
  65. Distinct("channel_id").
  66. Pluck("channel_id", &channelIds).Error
  67. if err != nil {
  68. return nil, 0, err
  69. }
  70. if len(channelIds) == 0 {
  71. return []map[string]any{}, 0, nil
  72. }
  73. type channelInfo struct {
  74. Id int `json:"id"`
  75. Name string `json:"name"`
  76. PublicName string `json:"public_name"`
  77. }
  78. var channels []channelInfo
  79. err = DB.Table("channels").
  80. Where("id IN ? AND status = ?", channelIds, common.ChannelStatusEnabled).
  81. Select("id, name, public_name").
  82. Find(&channels).Error
  83. if err != nil {
  84. return nil, 0, err
  85. }
  86. defaultChannelId := 0
  87. if defaultChId, ok := GetDefaultChannelId(modelName); ok {
  88. for _, id := range channelIds {
  89. if id == defaultChId {
  90. defaultChannelId = defaultChId
  91. break
  92. }
  93. }
  94. }
  95. result := make([]map[string]any, 0, len(channels))
  96. for _, ch := range channels {
  97. result = append(result, map[string]any{
  98. "id": ch.Id,
  99. "name": ch.Name,
  100. "public_name": ch.PublicName,
  101. })
  102. }
  103. return result, defaultChannelId, nil
  104. }
  105. func getPriority(group string, model string, retry int) (int, error) {
  106. var priorities []int
  107. err := DB.Model(&Ability{}).
  108. Select("DISTINCT(priority)").
  109. Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, model, true).
  110. Order("priority DESC"). // 按优先级降序排序
  111. Pluck("priority", &priorities).Error // Pluck用于将查询的结果直接扫描到一个切片中
  112. if err != nil {
  113. // 处理错误
  114. return 0, err
  115. }
  116. if len(priorities) == 0 {
  117. // 如果没有查询到优先级,则返回错误
  118. return 0, errors.New("数据库一致性被破坏")
  119. }
  120. // 确定要使用的优先级
  121. var priorityToUse int
  122. if retry >= len(priorities) {
  123. // 如果重试次数大于优先级数,则使用最小的优先级
  124. priorityToUse = priorities[len(priorities)-1]
  125. } else {
  126. priorityToUse = priorities[retry]
  127. }
  128. return priorityToUse, nil
  129. }
  130. func getChannelQuery(group string, model string, retry int) (*gorm.DB, error) {
  131. maxPrioritySubQuery := DB.Model(&Ability{}).Select("MAX(priority)").Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, model, true)
  132. channelQuery := DB.Where(commonGroupCol+" = ? and model = ? and enabled = ? and priority = (?)", group, model, true, maxPrioritySubQuery)
  133. if retry != 0 {
  134. priority, err := getPriority(group, model, retry)
  135. if err != nil {
  136. return nil, err
  137. } else {
  138. channelQuery = DB.Where(commonGroupCol+" = ? and model = ? and enabled = ? and priority = ?", group, model, true, priority)
  139. }
  140. }
  141. return channelQuery, nil
  142. }
  143. func GetChannel(group string, model string, retry int) (*Channel, error) {
  144. var abilities []Ability
  145. var err error = nil
  146. channelQuery, err := getChannelQuery(group, model, retry)
  147. if err != nil {
  148. return nil, err
  149. }
  150. if common.UsingSQLite || common.UsingPostgreSQL {
  151. err = channelQuery.Order("weight DESC").Find(&abilities).Error
  152. } else {
  153. err = channelQuery.Order("weight DESC").Find(&abilities).Error
  154. }
  155. if err != nil {
  156. return nil, err
  157. }
  158. channel := Channel{}
  159. if len(abilities) > 0 {
  160. // Randomly choose one
  161. weightSum := uint(0)
  162. for _, ability_ := range abilities {
  163. weightSum += ability_.Weight + 10
  164. }
  165. // Randomly choose one
  166. weight := common.GetRandomInt(int(weightSum))
  167. for _, ability_ := range abilities {
  168. weight -= int(ability_.Weight) + 10
  169. //log.Printf("weight: %d, ability weight: %d", weight, *ability_.Weight)
  170. if weight <= 0 {
  171. channel.Id = ability_.ChannelId
  172. break
  173. }
  174. }
  175. } else {
  176. return nil, nil
  177. }
  178. err = DB.First(&channel, "id = ?", channel.Id).Error
  179. return &channel, err
  180. }
  181. func (channel *Channel) AddAbilities(tx *gorm.DB) error {
  182. models_ := strings.Split(channel.Models, ",")
  183. // 统一使用 useDB,避免重复的数据库连接选择
  184. useDB := DB
  185. if tx != nil {
  186. useDB = tx
  187. }
  188. // === 自动同步模型到 models 表(批量处理)===
  189. // 收集需要同步的模型名称
  190. modelNames := make([]string, 0, len(models_))
  191. for _, modelName := range models_ {
  192. modelName = strings.TrimSpace(modelName)
  193. if modelName != "" {
  194. modelNames = append(modelNames, modelName)
  195. }
  196. }
  197. // 批量查询已存在的模型
  198. existingModels := make(map[string]bool)
  199. if len(modelNames) > 0 {
  200. var existingModelList []Model
  201. useDB.Where("model_name IN ?", modelNames).Select("model_name").Find(&existingModelList)
  202. for _, m := range existingModelList {
  203. existingModels[m.ModelName] = true
  204. }
  205. }
  206. // 收集需要创建的模型
  207. now := common.GetTimestamp()
  208. newModels := make([]Model, 0, len(modelNames))
  209. for _, modelName := range modelNames {
  210. if !existingModels[modelName] {
  211. newModels = append(newModels, Model{
  212. ModelName: modelName,
  213. Status: 1, // 启用状态
  214. SyncOfficial: 0, // 标记为用户模型,不被官方同步覆盖
  215. CreatedTime: now,
  216. UpdatedTime: now,
  217. })
  218. }
  219. }
  220. // 批量插入新模型(使用 OnConflict 忽略重复)
  221. if len(newModels) > 0 {
  222. for _, chunk := range lo.Chunk(newModels, 50) {
  223. if err := useDB.Clauses(clause.OnConflict{DoNothing: true}).Create(&chunk).Error; err != nil {
  224. return err
  225. }
  226. }
  227. }
  228. // === 自动同步结束 ===
  229. groups_ := strings.Split(channel.Group, ",")
  230. abilitySet := make(map[string]struct{})
  231. abilities := make([]Ability, 0, len(models_))
  232. for _, model := range models_ {
  233. for _, group := range groups_ {
  234. key := group + "|" + model
  235. if _, exists := abilitySet[key]; exists {
  236. continue
  237. }
  238. abilitySet[key] = struct{}{}
  239. ability := Ability{
  240. Group: group,
  241. Model: model,
  242. ChannelId: channel.Id,
  243. Enabled: channel.Status == common.ChannelStatusEnabled,
  244. Priority: channel.Priority,
  245. Weight: uint(channel.GetWeight()),
  246. Tag: channel.Tag,
  247. }
  248. abilities = append(abilities, ability)
  249. }
  250. }
  251. if len(abilities) == 0 {
  252. return nil
  253. }
  254. for _, chunk := range lo.Chunk(abilities, 50) {
  255. err := useDB.Clauses(clause.OnConflict{DoNothing: true}).Create(&chunk).Error
  256. if err != nil {
  257. return err
  258. }
  259. }
  260. return nil
  261. }
  262. func (channel *Channel) DeleteAbilities() error {
  263. return DB.Where("channel_id = ?", channel.Id).Delete(&Ability{}).Error
  264. }
  265. // UpdateAbilities updates abilities of this channel.
  266. // Make sure the channel is completed before calling this function.
  267. func (channel *Channel) UpdateAbilities(tx *gorm.DB) error {
  268. isNewTx := false
  269. // 如果没有传入事务,创建新的事务
  270. if tx == nil {
  271. tx = DB.Begin()
  272. if tx.Error != nil {
  273. return tx.Error
  274. }
  275. isNewTx = true
  276. defer func() {
  277. if r := recover(); r != nil {
  278. tx.Rollback()
  279. }
  280. }()
  281. }
  282. // First delete all abilities of this channel
  283. err := tx.Where("channel_id = ?", channel.Id).Delete(&Ability{}).Error
  284. if err != nil {
  285. if isNewTx {
  286. tx.Rollback()
  287. }
  288. return err
  289. }
  290. // Then add new abilities
  291. models_ := strings.Split(channel.Models, ",")
  292. groups_ := strings.Split(channel.Group, ",")
  293. abilitySet := make(map[string]struct{})
  294. abilities := make([]Ability, 0, len(models_))
  295. for _, model := range models_ {
  296. for _, group := range groups_ {
  297. key := group + "|" + model
  298. if _, exists := abilitySet[key]; exists {
  299. continue
  300. }
  301. abilitySet[key] = struct{}{}
  302. ability := Ability{
  303. Group: group,
  304. Model: model,
  305. ChannelId: channel.Id,
  306. Enabled: channel.Status == common.ChannelStatusEnabled,
  307. Priority: channel.Priority,
  308. Weight: uint(channel.GetWeight()),
  309. Tag: channel.Tag,
  310. }
  311. abilities = append(abilities, ability)
  312. }
  313. }
  314. if len(abilities) > 0 {
  315. for _, chunk := range lo.Chunk(abilities, 50) {
  316. err = tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&chunk).Error
  317. if err != nil {
  318. if isNewTx {
  319. tx.Rollback()
  320. }
  321. return err
  322. }
  323. }
  324. }
  325. // 如果是新创建的事务,需要提交
  326. if isNewTx {
  327. return tx.Commit().Error
  328. }
  329. return nil
  330. }
  331. func UpdateAbilityStatus(channelId int, status bool) error {
  332. return DB.Model(&Ability{}).Where("channel_id = ?", channelId).Select("enabled").Update("enabled", status).Error
  333. }
  334. func UpdateAbilityStatusByTag(tag string, status bool) error {
  335. return DB.Model(&Ability{}).Where("tag = ?", tag).Select("enabled").Update("enabled", status).Error
  336. }
  337. func UpdateAbilityByTag(tag string, newTag *string, priority *int64, weight *uint) error {
  338. ability := Ability{}
  339. if newTag != nil {
  340. ability.Tag = newTag
  341. }
  342. if priority != nil {
  343. ability.Priority = priority
  344. }
  345. if weight != nil {
  346. ability.Weight = *weight
  347. }
  348. return DB.Model(&Ability{}).Where("tag = ?", tag).Updates(ability).Error
  349. }
  350. var fixLock = sync.Mutex{}
  351. func FixAbility() (int, int, error) {
  352. lock := fixLock.TryLock()
  353. if !lock {
  354. return 0, 0, errors.New("已经有一个修复任务在运行中,请稍后再试")
  355. }
  356. defer fixLock.Unlock()
  357. // truncate abilities table
  358. if common.UsingSQLite {
  359. err := DB.Exec("DELETE FROM abilities").Error
  360. if err != nil {
  361. common.SysLog(fmt.Sprintf("Delete abilities failed: %s", err.Error()))
  362. return 0, 0, err
  363. }
  364. } else {
  365. err := DB.Exec("TRUNCATE TABLE abilities").Error
  366. if err != nil {
  367. common.SysLog(fmt.Sprintf("Truncate abilities failed: %s", err.Error()))
  368. return 0, 0, err
  369. }
  370. }
  371. var channels []*Channel
  372. // Find all channels
  373. err := DB.Model(&Channel{}).Find(&channels).Error
  374. if err != nil {
  375. return 0, 0, err
  376. }
  377. if len(channels) == 0 {
  378. return 0, 0, nil
  379. }
  380. successCount := 0
  381. failCount := 0
  382. for _, chunk := range lo.Chunk(channels, 50) {
  383. ids := lo.Map(chunk, func(c *Channel, _ int) int { return c.Id })
  384. // Delete all abilities of this channel
  385. err = DB.Where("channel_id IN ?", ids).Delete(&Ability{}).Error
  386. if err != nil {
  387. common.SysLog(fmt.Sprintf("Delete abilities failed: %s", err.Error()))
  388. failCount += len(chunk)
  389. continue
  390. }
  391. // Then add new abilities
  392. for _, channel := range chunk {
  393. err = channel.AddAbilities(nil)
  394. if err != nil {
  395. common.SysLog(fmt.Sprintf("Add abilities for channel %d failed: %s", channel.Id, err.Error()))
  396. failCount++
  397. } else {
  398. successCount++
  399. }
  400. }
  401. }
  402. InitChannelCache()
  403. return successCount, failCount, nil
  404. }