選択できるのは25トピックまでです。 トピックは、先頭が英数字で、英数字とダッシュ('-')を使用した35文字以内のものにしてください。
 
 
 

376 行
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. // === 自动同步模型到 models 表 ===
  131. // 选择数据库连接(优先使用事务)
  132. syncDB := DB
  133. if tx != nil {
  134. syncDB = tx
  135. }
  136. for _, modelName := range models_ {
  137. modelName = strings.TrimSpace(modelName)
  138. if modelName == "" {
  139. continue
  140. }
  141. // 检查 models 表是否存在该模型
  142. var existingModel Model
  143. err := syncDB.Where("model_name = ?", modelName).First(&existingModel).Error
  144. if err != nil {
  145. if errors.Is(err, gorm.ErrRecordNotFound) {
  146. // 自动创建模型记录
  147. newModel := &Model{
  148. ModelName: modelName,
  149. Status: 1, // 启用状态
  150. SyncOfficial: 0, // 标记为用户模型,不被官方同步覆盖
  151. }
  152. if err := newModel.Insert(); err != nil {
  153. common.SysLog(fmt.Sprintf("failed to auto-create model %s: %v", modelName, err))
  154. }
  155. } else {
  156. common.SysLog(fmt.Sprintf("failed to check model %s existence: %v", modelName, err))
  157. }
  158. }
  159. }
  160. // === 自动同步结束 ===
  161. groups_ := strings.Split(channel.Group, ",")
  162. abilitySet := make(map[string]struct{})
  163. abilities := make([]Ability, 0, len(models_))
  164. for _, model := range models_ {
  165. for _, group := range groups_ {
  166. key := group + "|" + model
  167. if _, exists := abilitySet[key]; exists {
  168. continue
  169. }
  170. abilitySet[key] = struct{}{}
  171. ability := Ability{
  172. Group: group,
  173. Model: model,
  174. ChannelId: channel.Id,
  175. Enabled: channel.Status == common.ChannelStatusEnabled,
  176. Priority: channel.Priority,
  177. Weight: uint(channel.GetWeight()),
  178. Tag: channel.Tag,
  179. }
  180. abilities = append(abilities, ability)
  181. }
  182. }
  183. if len(abilities) == 0 {
  184. return nil
  185. }
  186. // choose DB or provided tx
  187. useDB := DB
  188. if tx != nil {
  189. useDB = tx
  190. }
  191. for _, chunk := range lo.Chunk(abilities, 50) {
  192. err := useDB.Clauses(clause.OnConflict{DoNothing: true}).Create(&chunk).Error
  193. if err != nil {
  194. return err
  195. }
  196. }
  197. return nil
  198. }
  199. func (channel *Channel) DeleteAbilities() error {
  200. return DB.Where("channel_id = ?", channel.Id).Delete(&Ability{}).Error
  201. }
  202. // UpdateAbilities updates abilities of this channel.
  203. // Make sure the channel is completed before calling this function.
  204. func (channel *Channel) UpdateAbilities(tx *gorm.DB) error {
  205. isNewTx := false
  206. // 如果没有传入事务,创建新的事务
  207. if tx == nil {
  208. tx = DB.Begin()
  209. if tx.Error != nil {
  210. return tx.Error
  211. }
  212. isNewTx = true
  213. defer func() {
  214. if r := recover(); r != nil {
  215. tx.Rollback()
  216. }
  217. }()
  218. }
  219. // First delete all abilities of this channel
  220. err := tx.Where("channel_id = ?", channel.Id).Delete(&Ability{}).Error
  221. if err != nil {
  222. if isNewTx {
  223. tx.Rollback()
  224. }
  225. return err
  226. }
  227. // Then add new abilities
  228. models_ := strings.Split(channel.Models, ",")
  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. for _, chunk := range lo.Chunk(abilities, 50) {
  253. err = tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&chunk).Error
  254. if err != nil {
  255. if isNewTx {
  256. tx.Rollback()
  257. }
  258. return err
  259. }
  260. }
  261. }
  262. // 如果是新创建的事务,需要提交
  263. if isNewTx {
  264. return tx.Commit().Error
  265. }
  266. return nil
  267. }
  268. func UpdateAbilityStatus(channelId int, status bool) error {
  269. return DB.Model(&Ability{}).Where("channel_id = ?", channelId).Select("enabled").Update("enabled", status).Error
  270. }
  271. func UpdateAbilityStatusByTag(tag string, status bool) error {
  272. return DB.Model(&Ability{}).Where("tag = ?", tag).Select("enabled").Update("enabled", status).Error
  273. }
  274. func UpdateAbilityByTag(tag string, newTag *string, priority *int64, weight *uint) error {
  275. ability := Ability{}
  276. if newTag != nil {
  277. ability.Tag = newTag
  278. }
  279. if priority != nil {
  280. ability.Priority = priority
  281. }
  282. if weight != nil {
  283. ability.Weight = *weight
  284. }
  285. return DB.Model(&Ability{}).Where("tag = ?", tag).Updates(ability).Error
  286. }
  287. var fixLock = sync.Mutex{}
  288. func FixAbility() (int, int, error) {
  289. lock := fixLock.TryLock()
  290. if !lock {
  291. return 0, 0, errors.New("已经有一个修复任务在运行中,请稍后再试")
  292. }
  293. defer fixLock.Unlock()
  294. // truncate abilities table
  295. if common.UsingSQLite {
  296. err := DB.Exec("DELETE FROM abilities").Error
  297. if err != nil {
  298. common.SysLog(fmt.Sprintf("Delete abilities failed: %s", err.Error()))
  299. return 0, 0, err
  300. }
  301. } else {
  302. err := DB.Exec("TRUNCATE TABLE abilities").Error
  303. if err != nil {
  304. common.SysLog(fmt.Sprintf("Truncate abilities failed: %s", err.Error()))
  305. return 0, 0, err
  306. }
  307. }
  308. var channels []*Channel
  309. // Find all channels
  310. err := DB.Model(&Channel{}).Find(&channels).Error
  311. if err != nil {
  312. return 0, 0, err
  313. }
  314. if len(channels) == 0 {
  315. return 0, 0, nil
  316. }
  317. successCount := 0
  318. failCount := 0
  319. for _, chunk := range lo.Chunk(channels, 50) {
  320. ids := lo.Map(chunk, func(c *Channel, _ int) int { return c.Id })
  321. // Delete all abilities of this channel
  322. err = DB.Where("channel_id IN ?", ids).Delete(&Ability{}).Error
  323. if err != nil {
  324. common.SysLog(fmt.Sprintf("Delete abilities failed: %s", err.Error()))
  325. failCount += len(chunk)
  326. continue
  327. }
  328. // Then add new abilities
  329. for _, channel := range chunk {
  330. err = channel.AddAbilities(nil)
  331. if err != nil {
  332. common.SysLog(fmt.Sprintf("Add abilities for channel %d failed: %s", channel.Id, err.Error()))
  333. failCount++
  334. } else {
  335. successCount++
  336. }
  337. }
  338. }
  339. InitChannelCache()
  340. return successCount, failCount, nil
  341. }