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.
 
 
 

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