Du kannst nicht mehr als 25 Themen auswählen Themen müssen entweder mit einem Buchstaben oder einer Ziffer beginnen. Sie können Bindestriche („-“) enthalten und bis zu 35 Zeichen lang sein.
 
 
 

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