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.
 
 
 

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