channel.go 30 KB

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