25'ten fazla konu seçemezsiniz Konular bir harf veya rakamla başlamalı, kısa çizgiler ('-') içerebilir ve en fazla 35 karakter uzunluğunda olabilir.
 
 
 

1021 satır
30 KiB

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