channel.go 28 KB

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