channel.go 28 KB

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