You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

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