channel.go 25 KB

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