channel.go 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911
  1. package model
  2. import (
  3. "database/sql/driver"
  4. "encoding/json"
  5. "errors"
  6. "fmt"
  7. "math/rand"
  8. "one-api/common"
  9. "one-api/constant"
  10. "one-api/dto"
  11. "one-api/types"
  12. "strings"
  13. "sync"
  14. "gorm.io/gorm"
  15. )
  16. type Channel struct {
  17. Id int `json:"id"`
  18. Type int `json:"type" gorm:"default:0"`
  19. Key string `json:"key" gorm:"not null"`
  20. OpenAIOrganization *string `json:"openai_organization"`
  21. TestModel *string `json:"test_model"`
  22. Status int `json:"status" gorm:"default:1"`
  23. Name string `json:"name" gorm:"index"`
  24. Weight *uint `json:"weight" gorm:"default:0"`
  25. CreatedTime int64 `json:"created_time" gorm:"bigint"`
  26. TestTime int64 `json:"test_time" gorm:"bigint"`
  27. ResponseTime int `json:"response_time"` // in milliseconds
  28. BaseURL *string `json:"base_url" gorm:"column:base_url;default:''"`
  29. Other string `json:"other"`
  30. Balance float64 `json:"balance"` // in USD
  31. BalanceUpdatedTime int64 `json:"balance_updated_time" gorm:"bigint"`
  32. Models string `json:"models"`
  33. Group string `json:"group" gorm:"type:varchar(64);default:'default'"`
  34. UsedQuota int64 `json:"used_quota" gorm:"bigint;default:0"`
  35. ModelMapping *string `json:"model_mapping" gorm:"type:text"`
  36. //MaxInputTokens *int `json:"max_input_tokens" gorm:"default:0"`
  37. StatusCodeMapping *string `json:"status_code_mapping" gorm:"type:varchar(1024);default:''"`
  38. Priority *int64 `json:"priority" gorm:"bigint;default:0"`
  39. AutoBan *int `json:"auto_ban" gorm:"default:1"`
  40. OtherInfo string `json:"other_info"`
  41. Tag *string `json:"tag" gorm:"index"`
  42. Setting *string `json:"setting" gorm:"type:text"` // 渠道额外设置
  43. ParamOverride *string `json:"param_override" gorm:"type:text"`
  44. // add after v0.8.5
  45. ChannelInfo ChannelInfo `json:"channel_info" gorm:"type:json"`
  46. // cache info
  47. Keys []string `json:"-" gorm:"-"`
  48. }
  49. type ChannelInfo struct {
  50. IsMultiKey bool `json:"is_multi_key"` // 是否多Key模式
  51. MultiKeySize int `json:"multi_key_size"` // 多Key模式下的Key数量
  52. MultiKeyStatusList map[int]int `json:"multi_key_status_list"` // key状态列表,key index -> status
  53. MultiKeyPollingIndex int `json:"multi_key_polling_index"` // 多Key模式下轮询的key索引
  54. MultiKeyMode constant.MultiKeyMode `json:"multi_key_mode"`
  55. }
  56. // Value implements driver.Valuer interface
  57. func (c ChannelInfo) Value() (driver.Value, error) {
  58. return common.Marshal(&c)
  59. }
  60. // Scan implements sql.Scanner interface
  61. func (c *ChannelInfo) Scan(value interface{}) error {
  62. bytesValue, _ := value.([]byte)
  63. return common.Unmarshal(bytesValue, c)
  64. }
  65. func (channel *Channel) getKeys() []string {
  66. if channel.Key == "" {
  67. return []string{}
  68. }
  69. if len(channel.Keys) > 0 {
  70. return channel.Keys
  71. }
  72. trimmed := strings.TrimSpace(channel.Key)
  73. // If the key starts with '[', try to parse it as a JSON array (e.g., for Vertex AI scenarios)
  74. if strings.HasPrefix(trimmed, "[") {
  75. var arr []json.RawMessage
  76. if err := common.Unmarshal([]byte(trimmed), &arr); err == nil {
  77. res := make([]string, len(arr))
  78. for i, v := range arr {
  79. res[i] = string(v)
  80. }
  81. return res
  82. }
  83. }
  84. // Otherwise, fall back to splitting by newline
  85. keys := strings.Split(strings.Trim(channel.Key, "\n"), "\n")
  86. return keys
  87. }
  88. func (channel *Channel) GetNextEnabledKey() (string, int, *types.NewAPIError) {
  89. // If not in multi-key mode, return the original key string directly.
  90. if !channel.ChannelInfo.IsMultiKey {
  91. return channel.Key, 0, nil
  92. }
  93. // Obtain all keys (split by \n)
  94. keys := channel.getKeys()
  95. if len(keys) == 0 {
  96. // No keys available, return error, should disable the channel
  97. return "", 0, types.NewError(errors.New("no keys available"), types.ErrorCodeChannelNoAvailableKey)
  98. }
  99. statusList := channel.ChannelInfo.MultiKeyStatusList
  100. // helper to get key status, default to enabled when missing
  101. getStatus := func(idx int) int {
  102. if statusList == nil {
  103. return common.ChannelStatusEnabled
  104. }
  105. if status, ok := statusList[idx]; ok {
  106. return status
  107. }
  108. return common.ChannelStatusEnabled
  109. }
  110. // Collect indexes of enabled keys
  111. enabledIdx := make([]int, 0, len(keys))
  112. for i := range keys {
  113. if getStatus(i) == common.ChannelStatusEnabled {
  114. enabledIdx = append(enabledIdx, i)
  115. }
  116. }
  117. // If no specific status list or none enabled, fall back to first key
  118. if len(enabledIdx) == 0 {
  119. return keys[0], 0, nil
  120. }
  121. switch channel.ChannelInfo.MultiKeyMode {
  122. case constant.MultiKeyModeRandom:
  123. // Randomly pick one enabled key
  124. selectedIdx := enabledIdx[rand.Intn(len(enabledIdx))]
  125. return keys[selectedIdx], selectedIdx, nil
  126. case constant.MultiKeyModePolling:
  127. // Use channel-specific lock to ensure thread-safe polling
  128. lock := getChannelPollingLock(channel.Id)
  129. lock.Lock()
  130. defer lock.Unlock()
  131. channelInfo, err := CacheGetChannelInfo(channel.Id)
  132. if err != nil {
  133. return "", 0, types.NewError(err, types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
  134. }
  135. //println("before polling index:", channel.ChannelInfo.MultiKeyPollingIndex)
  136. defer func() {
  137. if common.DebugEnabled {
  138. println(fmt.Sprintf("channel %d polling index: %d", channel.Id, channel.ChannelInfo.MultiKeyPollingIndex))
  139. }
  140. if !common.MemoryCacheEnabled {
  141. _ = channel.SaveChannelInfo()
  142. } else {
  143. // CacheUpdateChannel(channel)
  144. }
  145. }()
  146. // Start from the saved polling index and look for the next enabled key
  147. start := channelInfo.MultiKeyPollingIndex
  148. if start < 0 || start >= len(keys) {
  149. start = 0
  150. }
  151. for i := 0; i < len(keys); i++ {
  152. idx := (start + i) % len(keys)
  153. if getStatus(idx) == common.ChannelStatusEnabled {
  154. // update polling index for next call (point to the next position)
  155. channel.ChannelInfo.MultiKeyPollingIndex = (idx + 1) % len(keys)
  156. return keys[idx], idx, nil
  157. }
  158. }
  159. // Fallback – should not happen, but return first enabled key
  160. return keys[enabledIdx[0]], enabledIdx[0], nil
  161. default:
  162. // Unknown mode, default to first enabled key (or original key string)
  163. return keys[enabledIdx[0]], enabledIdx[0], nil
  164. }
  165. }
  166. func (channel *Channel) SaveChannelInfo() error {
  167. return DB.Model(channel).Update("channel_info", channel.ChannelInfo).Error
  168. }
  169. func (channel *Channel) GetModels() []string {
  170. if channel.Models == "" {
  171. return []string{}
  172. }
  173. return strings.Split(strings.Trim(channel.Models, ","), ",")
  174. }
  175. func (channel *Channel) GetGroups() []string {
  176. if channel.Group == "" {
  177. return []string{}
  178. }
  179. groups := strings.Split(strings.Trim(channel.Group, ","), ",")
  180. for i, group := range groups {
  181. groups[i] = strings.TrimSpace(group)
  182. }
  183. return groups
  184. }
  185. func (channel *Channel) GetOtherInfo() map[string]interface{} {
  186. otherInfo := make(map[string]interface{})
  187. if channel.OtherInfo != "" {
  188. err := common.Unmarshal([]byte(channel.OtherInfo), &otherInfo)
  189. if err != nil {
  190. common.SysError("failed to unmarshal other info: " + err.Error())
  191. }
  192. }
  193. return otherInfo
  194. }
  195. func (channel *Channel) SetOtherInfo(otherInfo map[string]interface{}) {
  196. otherInfoBytes, err := json.Marshal(otherInfo)
  197. if err != nil {
  198. common.SysError("failed to marshal other info: " + err.Error())
  199. return
  200. }
  201. channel.OtherInfo = string(otherInfoBytes)
  202. }
  203. func (channel *Channel) GetTag() string {
  204. if channel.Tag == nil {
  205. return ""
  206. }
  207. return *channel.Tag
  208. }
  209. func (channel *Channel) SetTag(tag string) {
  210. channel.Tag = &tag
  211. }
  212. func (channel *Channel) GetAutoBan() bool {
  213. if channel.AutoBan == nil {
  214. return false
  215. }
  216. return *channel.AutoBan == 1
  217. }
  218. func (channel *Channel) Save() error {
  219. return DB.Save(channel).Error
  220. }
  221. func GetAllChannels(startIdx int, num int, selectAll bool, idSort bool) ([]*Channel, error) {
  222. var channels []*Channel
  223. var err error
  224. order := "priority desc"
  225. if idSort {
  226. order = "id desc"
  227. }
  228. if selectAll {
  229. err = DB.Order(order).Find(&channels).Error
  230. } else {
  231. err = DB.Order(order).Limit(num).Offset(startIdx).Omit("key").Find(&channels).Error
  232. }
  233. return channels, err
  234. }
  235. func GetChannelsByTag(tag string, idSort bool) ([]*Channel, error) {
  236. var channels []*Channel
  237. order := "priority desc"
  238. if idSort {
  239. order = "id desc"
  240. }
  241. err := DB.Where("tag = ?", tag).Order(order).Find(&channels).Error
  242. return channels, err
  243. }
  244. func SearchChannels(keyword string, group string, model string, idSort bool) ([]*Channel, error) {
  245. var channels []*Channel
  246. modelsCol := "`models`"
  247. // 如果是 PostgreSQL,使用双引号
  248. if common.UsingPostgreSQL {
  249. modelsCol = `"models"`
  250. }
  251. baseURLCol := "`base_url`"
  252. // 如果是 PostgreSQL,使用双引号
  253. if common.UsingPostgreSQL {
  254. baseURLCol = `"base_url"`
  255. }
  256. order := "priority desc"
  257. if idSort {
  258. order = "id desc"
  259. }
  260. // 构造基础查询
  261. baseQuery := DB.Model(&Channel{}).Omit("key")
  262. // 构造WHERE子句
  263. var whereClause string
  264. var args []interface{}
  265. if group != "" && group != "null" {
  266. var groupCondition string
  267. if common.UsingMySQL {
  268. groupCondition = `CONCAT(',', ` + commonGroupCol + `, ',') LIKE ?`
  269. } else {
  270. // sqlite, PostgreSQL
  271. groupCondition = `(',' || ` + commonGroupCol + ` || ',') LIKE ?`
  272. }
  273. whereClause = "(id = ? OR name LIKE ? OR " + commonKeyCol + " = ? OR " + baseURLCol + " LIKE ?) AND " + modelsCol + ` LIKE ? AND ` + groupCondition
  274. args = append(args, common.String2Int(keyword), "%"+keyword+"%", keyword, "%"+keyword+"%", "%"+model+"%", "%,"+group+",%")
  275. } else {
  276. whereClause = "(id = ? OR name LIKE ? OR " + commonKeyCol + " = ? OR " + baseURLCol + " LIKE ?) AND " + modelsCol + " LIKE ?"
  277. args = append(args, common.String2Int(keyword), "%"+keyword+"%", keyword, "%"+keyword+"%", "%"+model+"%")
  278. }
  279. // 执行查询
  280. err := baseQuery.Where(whereClause, args...).Order(order).Find(&channels).Error
  281. if err != nil {
  282. return nil, err
  283. }
  284. return channels, nil
  285. }
  286. func GetChannelById(id int, selectAll bool) (*Channel, error) {
  287. channel := &Channel{Id: id}
  288. var err error = nil
  289. if selectAll {
  290. err = DB.First(channel, "id = ?", id).Error
  291. } else {
  292. err = DB.Omit("key").First(channel, "id = ?", id).Error
  293. }
  294. if err != nil {
  295. return nil, err
  296. }
  297. if channel == nil {
  298. return nil, errors.New("channel not found")
  299. }
  300. return channel, nil
  301. }
  302. func BatchInsertChannels(channels []Channel) error {
  303. var err error
  304. err = DB.Create(&channels).Error
  305. if err != nil {
  306. return err
  307. }
  308. for _, channel_ := range channels {
  309. err = channel_.AddAbilities()
  310. if err != nil {
  311. return err
  312. }
  313. }
  314. return nil
  315. }
  316. func BatchDeleteChannels(ids []int) error {
  317. //使用事务 删除channel表和channel_ability表
  318. tx := DB.Begin()
  319. err := tx.Where("id in (?)", ids).Delete(&Channel{}).Error
  320. if err != nil {
  321. // 回滚事务
  322. tx.Rollback()
  323. return err
  324. }
  325. err = tx.Where("channel_id in (?)", ids).Delete(&Ability{}).Error
  326. if err != nil {
  327. // 回滚事务
  328. tx.Rollback()
  329. return err
  330. }
  331. // 提交事务
  332. tx.Commit()
  333. return err
  334. }
  335. func (channel *Channel) GetPriority() int64 {
  336. if channel.Priority == nil {
  337. return 0
  338. }
  339. return *channel.Priority
  340. }
  341. func (channel *Channel) GetWeight() int {
  342. if channel.Weight == nil {
  343. return 0
  344. }
  345. return int(*channel.Weight)
  346. }
  347. func (channel *Channel) GetBaseURL() string {
  348. if channel.BaseURL == nil {
  349. return ""
  350. }
  351. return *channel.BaseURL
  352. }
  353. func (channel *Channel) GetModelMapping() string {
  354. if channel.ModelMapping == nil {
  355. return ""
  356. }
  357. return *channel.ModelMapping
  358. }
  359. func (channel *Channel) GetStatusCodeMapping() string {
  360. if channel.StatusCodeMapping == nil {
  361. return ""
  362. }
  363. return *channel.StatusCodeMapping
  364. }
  365. func (channel *Channel) Insert() error {
  366. var err error
  367. err = DB.Create(channel).Error
  368. if err != nil {
  369. return err
  370. }
  371. err = channel.AddAbilities()
  372. return err
  373. }
  374. func (channel *Channel) Update() error {
  375. // If this is a multi-key channel, recalculate MultiKeySize based on the current key list to avoid inconsistency after editing keys
  376. if channel.ChannelInfo.IsMultiKey {
  377. var keyStr string
  378. if channel.Key != "" {
  379. keyStr = channel.Key
  380. } else {
  381. // If key is not provided, read the existing key from the database
  382. if existing, err := GetChannelById(channel.Id, true); err == nil {
  383. keyStr = existing.Key
  384. }
  385. }
  386. // Parse the key list (supports newline separation or JSON array)
  387. keys := []string{}
  388. if keyStr != "" {
  389. trimmed := strings.TrimSpace(keyStr)
  390. if strings.HasPrefix(trimmed, "[") {
  391. var arr []json.RawMessage
  392. if err := common.Unmarshal([]byte(trimmed), &arr); err == nil {
  393. keys = make([]string, len(arr))
  394. for i, v := range arr {
  395. keys[i] = string(v)
  396. }
  397. }
  398. }
  399. if len(keys) == 0 { // fallback to newline split
  400. keys = strings.Split(strings.Trim(keyStr, "\n"), "\n")
  401. }
  402. }
  403. channel.ChannelInfo.MultiKeySize = len(keys)
  404. // Clean up status data that exceeds the new key count to prevent index out of range
  405. if channel.ChannelInfo.MultiKeyStatusList != nil {
  406. for idx := range channel.ChannelInfo.MultiKeyStatusList {
  407. if idx >= channel.ChannelInfo.MultiKeySize {
  408. delete(channel.ChannelInfo.MultiKeyStatusList, idx)
  409. }
  410. }
  411. }
  412. }
  413. var err error
  414. err = DB.Model(channel).Updates(channel).Error
  415. if err != nil {
  416. return err
  417. }
  418. DB.Model(channel).First(channel, "id = ?", channel.Id)
  419. err = channel.UpdateAbilities(nil)
  420. return err
  421. }
  422. func (channel *Channel) UpdateResponseTime(responseTime int64) {
  423. err := DB.Model(channel).Select("response_time", "test_time").Updates(Channel{
  424. TestTime: common.GetTimestamp(),
  425. ResponseTime: int(responseTime),
  426. }).Error
  427. if err != nil {
  428. common.SysError("failed to update response time: " + err.Error())
  429. }
  430. }
  431. func (channel *Channel) UpdateBalance(balance float64) {
  432. err := DB.Model(channel).Select("balance_updated_time", "balance").Updates(Channel{
  433. BalanceUpdatedTime: common.GetTimestamp(),
  434. Balance: balance,
  435. }).Error
  436. if err != nil {
  437. common.SysError("failed to update balance: " + err.Error())
  438. }
  439. }
  440. func (channel *Channel) Delete() error {
  441. var err error
  442. err = DB.Delete(channel).Error
  443. if err != nil {
  444. return err
  445. }
  446. err = channel.DeleteAbilities()
  447. return err
  448. }
  449. var channelStatusLock sync.Mutex
  450. // channelPollingLocks stores locks for each channel.id to ensure thread-safe polling
  451. var channelPollingLocks sync.Map
  452. // getChannelPollingLock returns or creates a mutex for the given channel ID
  453. func getChannelPollingLock(channelId int) *sync.Mutex {
  454. if lock, exists := channelPollingLocks.Load(channelId); exists {
  455. return lock.(*sync.Mutex)
  456. }
  457. // Create new lock for this channel
  458. newLock := &sync.Mutex{}
  459. actual, _ := channelPollingLocks.LoadOrStore(channelId, newLock)
  460. return actual.(*sync.Mutex)
  461. }
  462. // CleanupChannelPollingLocks removes locks for channels that no longer exist
  463. // This is optional and can be called periodically to prevent memory leaks
  464. func CleanupChannelPollingLocks() {
  465. var activeChannelIds []int
  466. DB.Model(&Channel{}).Pluck("id", &activeChannelIds)
  467. activeChannelSet := make(map[int]bool)
  468. for _, id := range activeChannelIds {
  469. activeChannelSet[id] = true
  470. }
  471. channelPollingLocks.Range(func(key, value interface{}) bool {
  472. channelId := key.(int)
  473. if !activeChannelSet[channelId] {
  474. channelPollingLocks.Delete(channelId)
  475. }
  476. return true
  477. })
  478. }
  479. func handlerMultiKeyUpdate(channel *Channel, usingKey string, status int) {
  480. keys := channel.getKeys()
  481. if len(keys) == 0 {
  482. channel.Status = status
  483. } else {
  484. var keyIndex int
  485. for i, key := range keys {
  486. if key == usingKey {
  487. keyIndex = i
  488. break
  489. }
  490. }
  491. if channel.ChannelInfo.MultiKeyStatusList == nil {
  492. channel.ChannelInfo.MultiKeyStatusList = make(map[int]int)
  493. }
  494. if status == common.ChannelStatusEnabled {
  495. delete(channel.ChannelInfo.MultiKeyStatusList, keyIndex)
  496. } else {
  497. channel.ChannelInfo.MultiKeyStatusList[keyIndex] = status
  498. }
  499. if len(channel.ChannelInfo.MultiKeyStatusList) >= channel.ChannelInfo.MultiKeySize {
  500. channel.Status = common.ChannelStatusAutoDisabled
  501. info := channel.GetOtherInfo()
  502. info["status_reason"] = "All keys are disabled"
  503. info["status_time"] = common.GetTimestamp()
  504. channel.SetOtherInfo(info)
  505. }
  506. }
  507. }
  508. func UpdateChannelStatus(channelId int, usingKey string, status int, reason string) bool {
  509. if common.MemoryCacheEnabled {
  510. channelStatusLock.Lock()
  511. defer channelStatusLock.Unlock()
  512. channelCache, _ := CacheGetChannel(channelId)
  513. if channelCache == nil {
  514. return false
  515. }
  516. if channelCache.ChannelInfo.IsMultiKey {
  517. // 如果是多Key模式,更新缓存中的状态
  518. handlerMultiKeyUpdate(channelCache, usingKey, status)
  519. //CacheUpdateChannel(channelCache)
  520. //return true
  521. } else {
  522. // 如果缓存渠道存在,且状态已是目标状态,直接返回
  523. if channelCache.Status == status {
  524. return false
  525. }
  526. CacheUpdateChannelStatus(channelId, status)
  527. }
  528. }
  529. shouldUpdateAbilities := false
  530. defer func() {
  531. if shouldUpdateAbilities {
  532. err := UpdateAbilityStatus(channelId, status == common.ChannelStatusEnabled)
  533. if err != nil {
  534. common.SysError("failed to update ability status: " + err.Error())
  535. }
  536. }
  537. }()
  538. channel, err := GetChannelById(channelId, true)
  539. if err != nil {
  540. return false
  541. } else {
  542. if channel.Status == status {
  543. return false
  544. }
  545. if channel.ChannelInfo.IsMultiKey {
  546. beforeStatus := channel.Status
  547. handlerMultiKeyUpdate(channel, usingKey, status)
  548. if beforeStatus != channel.Status {
  549. shouldUpdateAbilities = true
  550. }
  551. } else {
  552. info := channel.GetOtherInfo()
  553. info["status_reason"] = reason
  554. info["status_time"] = common.GetTimestamp()
  555. channel.SetOtherInfo(info)
  556. channel.Status = status
  557. shouldUpdateAbilities = true
  558. }
  559. err = channel.Save()
  560. if err != nil {
  561. common.SysError("failed to update channel status: " + err.Error())
  562. return false
  563. }
  564. }
  565. return true
  566. }
  567. func EnableChannelByTag(tag string) error {
  568. err := DB.Model(&Channel{}).Where("tag = ?", tag).Update("status", common.ChannelStatusEnabled).Error
  569. if err != nil {
  570. return err
  571. }
  572. err = UpdateAbilityStatusByTag(tag, true)
  573. return err
  574. }
  575. func DisableChannelByTag(tag string) error {
  576. err := DB.Model(&Channel{}).Where("tag = ?", tag).Update("status", common.ChannelStatusManuallyDisabled).Error
  577. if err != nil {
  578. return err
  579. }
  580. err = UpdateAbilityStatusByTag(tag, false)
  581. return err
  582. }
  583. func EditChannelByTag(tag string, newTag *string, modelMapping *string, models *string, group *string, priority *int64, weight *uint) error {
  584. updateData := Channel{}
  585. shouldReCreateAbilities := false
  586. updatedTag := tag
  587. // 如果 newTag 不为空且不等于 tag,则更新 tag
  588. if newTag != nil && *newTag != tag {
  589. updateData.Tag = newTag
  590. updatedTag = *newTag
  591. }
  592. if modelMapping != nil && *modelMapping != "" {
  593. updateData.ModelMapping = modelMapping
  594. }
  595. if models != nil && *models != "" {
  596. shouldReCreateAbilities = true
  597. updateData.Models = *models
  598. }
  599. if group != nil && *group != "" {
  600. shouldReCreateAbilities = true
  601. updateData.Group = *group
  602. }
  603. if priority != nil {
  604. updateData.Priority = priority
  605. }
  606. if weight != nil {
  607. updateData.Weight = weight
  608. }
  609. err := DB.Model(&Channel{}).Where("tag = ?", tag).Updates(updateData).Error
  610. if err != nil {
  611. return err
  612. }
  613. if shouldReCreateAbilities {
  614. channels, err := GetChannelsByTag(updatedTag, false)
  615. if err == nil {
  616. for _, channel := range channels {
  617. err = channel.UpdateAbilities(nil)
  618. if err != nil {
  619. common.SysError("failed to update abilities: " + err.Error())
  620. }
  621. }
  622. }
  623. } else {
  624. err := UpdateAbilityByTag(tag, newTag, priority, weight)
  625. if err != nil {
  626. return err
  627. }
  628. }
  629. return nil
  630. }
  631. func UpdateChannelUsedQuota(id int, quota int) {
  632. if common.BatchUpdateEnabled {
  633. addNewRecord(BatchUpdateTypeChannelUsedQuota, id, quota)
  634. return
  635. }
  636. updateChannelUsedQuota(id, quota)
  637. }
  638. func updateChannelUsedQuota(id int, quota int) {
  639. err := DB.Model(&Channel{}).Where("id = ?", id).Update("used_quota", gorm.Expr("used_quota + ?", quota)).Error
  640. if err != nil {
  641. common.SysError("failed to update channel used quota: " + err.Error())
  642. }
  643. }
  644. func DeleteChannelByStatus(status int64) (int64, error) {
  645. result := DB.Where("status = ?", status).Delete(&Channel{})
  646. return result.RowsAffected, result.Error
  647. }
  648. func DeleteDisabledChannel() (int64, error) {
  649. result := DB.Where("status = ? or status = ?", common.ChannelStatusAutoDisabled, common.ChannelStatusManuallyDisabled).Delete(&Channel{})
  650. return result.RowsAffected, result.Error
  651. }
  652. func GetPaginatedTags(offset int, limit int) ([]*string, error) {
  653. var tags []*string
  654. err := DB.Model(&Channel{}).Select("DISTINCT tag").Where("tag != ''").Offset(offset).Limit(limit).Find(&tags).Error
  655. return tags, err
  656. }
  657. func SearchTags(keyword string, group string, model string, idSort bool) ([]*string, error) {
  658. var tags []*string
  659. modelsCol := "`models`"
  660. // 如果是 PostgreSQL,使用双引号
  661. if common.UsingPostgreSQL {
  662. modelsCol = `"models"`
  663. }
  664. baseURLCol := "`base_url`"
  665. // 如果是 PostgreSQL,使用双引号
  666. if common.UsingPostgreSQL {
  667. baseURLCol = `"base_url"`
  668. }
  669. order := "priority desc"
  670. if idSort {
  671. order = "id desc"
  672. }
  673. // 构造基础查询
  674. baseQuery := DB.Model(&Channel{}).Omit("key")
  675. // 构造WHERE子句
  676. var whereClause string
  677. var args []interface{}
  678. if group != "" && group != "null" {
  679. var groupCondition string
  680. if common.UsingMySQL {
  681. groupCondition = `CONCAT(',', ` + commonGroupCol + `, ',') LIKE ?`
  682. } else {
  683. // sqlite, PostgreSQL
  684. groupCondition = `(',' || ` + commonGroupCol + ` || ',') LIKE ?`
  685. }
  686. whereClause = "(id = ? OR name LIKE ? OR " + commonKeyCol + " = ? OR " + baseURLCol + " LIKE ?) AND " + modelsCol + ` LIKE ? AND ` + groupCondition
  687. args = append(args, common.String2Int(keyword), "%"+keyword+"%", keyword, "%"+keyword+"%", "%"+model+"%", "%,"+group+",%")
  688. } else {
  689. whereClause = "(id = ? OR name LIKE ? OR " + commonKeyCol + " = ? OR " + baseURLCol + " LIKE ?) AND " + modelsCol + " LIKE ?"
  690. args = append(args, common.String2Int(keyword), "%"+keyword+"%", keyword, "%"+keyword+"%", "%"+model+"%")
  691. }
  692. subQuery := baseQuery.Where(whereClause, args...).
  693. Select("tag").
  694. Where("tag != ''").
  695. Order(order)
  696. err := DB.Table("(?) as sub", subQuery).
  697. Select("DISTINCT tag").
  698. Find(&tags).Error
  699. if err != nil {
  700. return nil, err
  701. }
  702. return tags, nil
  703. }
  704. func (channel *Channel) ValidateSettings() error {
  705. channelParams := &dto.ChannelSettings{}
  706. if channel.Setting != nil && *channel.Setting != "" {
  707. err := common.Unmarshal([]byte(*channel.Setting), channelParams)
  708. if err != nil {
  709. return err
  710. }
  711. }
  712. return nil
  713. }
  714. func (channel *Channel) GetSetting() dto.ChannelSettings {
  715. setting := dto.ChannelSettings{}
  716. if channel.Setting != nil && *channel.Setting != "" {
  717. err := common.Unmarshal([]byte(*channel.Setting), &setting)
  718. if err != nil {
  719. common.SysError("failed to unmarshal setting: " + err.Error())
  720. channel.Setting = nil // 清空设置以避免后续错误
  721. _ = channel.Save() // 保存修改
  722. }
  723. }
  724. return setting
  725. }
  726. func (channel *Channel) SetSetting(setting dto.ChannelSettings) {
  727. settingBytes, err := common.Marshal(setting)
  728. if err != nil {
  729. common.SysError("failed to marshal setting: " + err.Error())
  730. return
  731. }
  732. channel.Setting = common.GetPointer[string](string(settingBytes))
  733. }
  734. func (channel *Channel) GetParamOverride() map[string]interface{} {
  735. paramOverride := make(map[string]interface{})
  736. if channel.ParamOverride != nil && *channel.ParamOverride != "" {
  737. err := common.Unmarshal([]byte(*channel.ParamOverride), &paramOverride)
  738. if err != nil {
  739. common.SysError("failed to unmarshal param override: " + err.Error())
  740. }
  741. }
  742. return paramOverride
  743. }
  744. func GetChannelsByIds(ids []int) ([]*Channel, error) {
  745. var channels []*Channel
  746. err := DB.Where("id in (?)", ids).Find(&channels).Error
  747. return channels, err
  748. }
  749. func BatchSetChannelTag(ids []int, tag *string) error {
  750. // 开启事务
  751. tx := DB.Begin()
  752. if tx.Error != nil {
  753. return tx.Error
  754. }
  755. // 更新标签
  756. err := tx.Model(&Channel{}).Where("id in (?)", ids).Update("tag", tag).Error
  757. if err != nil {
  758. tx.Rollback()
  759. return err
  760. }
  761. // update ability status
  762. channels, err := GetChannelsByIds(ids)
  763. if err != nil {
  764. tx.Rollback()
  765. return err
  766. }
  767. for _, channel := range channels {
  768. err = channel.UpdateAbilities(tx)
  769. if err != nil {
  770. tx.Rollback()
  771. return err
  772. }
  773. }
  774. // 提交事务
  775. return tx.Commit().Error
  776. }
  777. // CountAllChannels returns total channels in DB
  778. func CountAllChannels() (int64, error) {
  779. var total int64
  780. err := DB.Model(&Channel{}).Count(&total).Error
  781. return total, err
  782. }
  783. // CountAllTags returns number of non-empty distinct tags
  784. func CountAllTags() (int64, error) {
  785. var total int64
  786. err := DB.Model(&Channel{}).Where("tag is not null AND tag != ''").Distinct("tag").Count(&total).Error
  787. return total, err
  788. }
  789. // Get channels of specified type with pagination
  790. func GetChannelsByType(startIdx int, num int, idSort bool, channelType int) ([]*Channel, error) {
  791. var channels []*Channel
  792. order := "priority desc"
  793. if idSort {
  794. order = "id desc"
  795. }
  796. err := DB.Where("type = ?", channelType).Order(order).Limit(num).Offset(startIdx).Omit("key").Find(&channels).Error
  797. return channels, err
  798. }
  799. // Count channels of specific type
  800. func CountChannelsByType(channelType int) (int64, error) {
  801. var count int64
  802. err := DB.Model(&Channel{}).Where("type = ?", channelType).Count(&count).Error
  803. return count, err
  804. }
  805. // Return map[type]count for all channels
  806. func CountChannelsGroupByType() (map[int64]int64, error) {
  807. type result struct {
  808. Type int64 `gorm:"column:type"`
  809. Count int64 `gorm:"column:count"`
  810. }
  811. var results []result
  812. err := DB.Model(&Channel{}).Select("type, count(*) as count").Group("type").Find(&results).Error
  813. if err != nil {
  814. return nil, err
  815. }
  816. counts := make(map[int64]int64)
  817. for _, r := range results {
  818. counts[r.Type] = r.Count
  819. }
  820. return counts, nil
  821. }