channel.go 26 KB

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