channel.go 30 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001
  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. return DB.Omit("key").Save(channel).Error
  231. }
  232. func GetAllChannels(startIdx int, num int, selectAll bool, idSort bool) ([]*Channel, error) {
  233. var channels []*Channel
  234. var err error
  235. order := "priority desc"
  236. if idSort {
  237. order = "id desc"
  238. }
  239. if selectAll {
  240. err = DB.Order(order).Find(&channels).Error
  241. } else {
  242. err = DB.Order(order).Limit(num).Offset(startIdx).Omit("key").Find(&channels).Error
  243. }
  244. return channels, err
  245. }
  246. func GetChannelsByTag(tag string, idSort bool) ([]*Channel, error) {
  247. var channels []*Channel
  248. order := "priority desc"
  249. if idSort {
  250. order = "id desc"
  251. }
  252. err := DB.Where("tag = ?", tag).Order(order).Find(&channels).Error
  253. return channels, err
  254. }
  255. func SearchChannels(keyword string, group string, model string, idSort bool) ([]*Channel, error) {
  256. var channels []*Channel
  257. modelsCol := "`models`"
  258. // 如果是 PostgreSQL,使用双引号
  259. if common.UsingPostgreSQL {
  260. modelsCol = `"models"`
  261. }
  262. baseURLCol := "`base_url`"
  263. // 如果是 PostgreSQL,使用双引号
  264. if common.UsingPostgreSQL {
  265. baseURLCol = `"base_url"`
  266. }
  267. order := "priority desc"
  268. if idSort {
  269. order = "id desc"
  270. }
  271. // 构造基础查询
  272. baseQuery := DB.Model(&Channel{}).Omit("key")
  273. // 构造WHERE子句
  274. var whereClause string
  275. var args []interface{}
  276. if group != "" && group != "null" {
  277. var groupCondition string
  278. if common.UsingMySQL {
  279. groupCondition = `CONCAT(',', ` + commonGroupCol + `, ',') LIKE ?`
  280. } else {
  281. // sqlite, PostgreSQL
  282. groupCondition = `(',' || ` + commonGroupCol + ` || ',') LIKE ?`
  283. }
  284. whereClause = "(id = ? OR name LIKE ? OR " + commonKeyCol + " = ? OR " + baseURLCol + " LIKE ?) AND " + modelsCol + ` LIKE ? AND ` + groupCondition
  285. args = append(args, common.String2Int(keyword), "%"+keyword+"%", keyword, "%"+keyword+"%", "%"+model+"%", "%,"+group+",%")
  286. } else {
  287. whereClause = "(id = ? OR name LIKE ? OR " + commonKeyCol + " = ? OR " + baseURLCol + " LIKE ?) AND " + modelsCol + " LIKE ?"
  288. args = append(args, common.String2Int(keyword), "%"+keyword+"%", keyword, "%"+keyword+"%", "%"+model+"%")
  289. }
  290. // 执行查询
  291. err := baseQuery.Where(whereClause, args...).Order(order).Find(&channels).Error
  292. if err != nil {
  293. return nil, err
  294. }
  295. return channels, nil
  296. }
  297. func GetChannelById(id int, selectAll bool) (*Channel, error) {
  298. channel := &Channel{Id: id}
  299. var err error = nil
  300. if selectAll {
  301. err = DB.First(channel, "id = ?", id).Error
  302. } else {
  303. err = DB.Omit("key").First(channel, "id = ?", id).Error
  304. }
  305. if err != nil {
  306. return nil, err
  307. }
  308. if channel == nil {
  309. return nil, errors.New("channel not found")
  310. }
  311. return channel, nil
  312. }
  313. func BatchInsertChannels(channels []Channel) error {
  314. if len(channels) == 0 {
  315. return nil
  316. }
  317. tx := DB.Begin()
  318. if tx.Error != nil {
  319. return tx.Error
  320. }
  321. defer func() {
  322. if r := recover(); r != nil {
  323. tx.Rollback()
  324. }
  325. }()
  326. for _, chunk := range lo.Chunk(channels, 50) {
  327. if err := tx.Create(&chunk).Error; err != nil {
  328. tx.Rollback()
  329. return err
  330. }
  331. for _, channel_ := range chunk {
  332. if err := channel_.AddAbilities(tx); err != nil {
  333. tx.Rollback()
  334. return err
  335. }
  336. }
  337. }
  338. return tx.Commit().Error
  339. }
  340. func BatchDeleteChannels(ids []int) error {
  341. if len(ids) == 0 {
  342. return nil
  343. }
  344. // 使用事务 分批删除channel表和abilities表
  345. tx := DB.Begin()
  346. if tx.Error != nil {
  347. return tx.Error
  348. }
  349. for _, chunk := range lo.Chunk(ids, 200) {
  350. if err := tx.Where("id in (?)", chunk).Delete(&Channel{}).Error; err != nil {
  351. tx.Rollback()
  352. return err
  353. }
  354. if err := tx.Where("channel_id in (?)", chunk).Delete(&Ability{}).Error; err != nil {
  355. tx.Rollback()
  356. return err
  357. }
  358. }
  359. return tx.Commit().Error
  360. }
  361. func (channel *Channel) GetPriority() int64 {
  362. if channel.Priority == nil {
  363. return 0
  364. }
  365. return *channel.Priority
  366. }
  367. func (channel *Channel) GetWeight() int {
  368. if channel.Weight == nil {
  369. return 0
  370. }
  371. return int(*channel.Weight)
  372. }
  373. func (channel *Channel) GetBaseURL() string {
  374. if channel.BaseURL == nil {
  375. return ""
  376. }
  377. url := *channel.BaseURL
  378. if url == "" {
  379. url = constant.ChannelBaseURLs[channel.Type]
  380. }
  381. return url
  382. }
  383. func (channel *Channel) GetModelMapping() string {
  384. if channel.ModelMapping == nil {
  385. return ""
  386. }
  387. return *channel.ModelMapping
  388. }
  389. func (channel *Channel) GetStatusCodeMapping() string {
  390. if channel.StatusCodeMapping == nil {
  391. return ""
  392. }
  393. return *channel.StatusCodeMapping
  394. }
  395. func (channel *Channel) Insert() error {
  396. var err error
  397. err = DB.Create(channel).Error
  398. if err != nil {
  399. return err
  400. }
  401. err = channel.AddAbilities(nil)
  402. return err
  403. }
  404. func (channel *Channel) Update() error {
  405. // If this is a multi-key channel, recalculate MultiKeySize based on the current key list to avoid inconsistency after editing keys
  406. if channel.ChannelInfo.IsMultiKey {
  407. var keyStr string
  408. if channel.Key != "" {
  409. keyStr = channel.Key
  410. } else {
  411. // If key is not provided, read the existing key from the database
  412. if existing, err := GetChannelById(channel.Id, true); err == nil {
  413. keyStr = existing.Key
  414. }
  415. }
  416. // Parse the key list (supports newline separation or JSON array)
  417. keys := []string{}
  418. if keyStr != "" {
  419. trimmed := strings.TrimSpace(keyStr)
  420. if strings.HasPrefix(trimmed, "[") {
  421. var arr []json.RawMessage
  422. if err := common.Unmarshal([]byte(trimmed), &arr); err == nil {
  423. keys = make([]string, len(arr))
  424. for i, v := range arr {
  425. keys[i] = string(v)
  426. }
  427. }
  428. }
  429. if len(keys) == 0 { // fallback to newline split
  430. keys = strings.Split(strings.Trim(keyStr, "\n"), "\n")
  431. }
  432. }
  433. channel.ChannelInfo.MultiKeySize = len(keys)
  434. // Clean up status data that exceeds the new key count to prevent index out of range
  435. if channel.ChannelInfo.MultiKeyStatusList != nil {
  436. for idx := range channel.ChannelInfo.MultiKeyStatusList {
  437. if idx >= channel.ChannelInfo.MultiKeySize {
  438. delete(channel.ChannelInfo.MultiKeyStatusList, idx)
  439. }
  440. }
  441. }
  442. }
  443. var err error
  444. err = DB.Model(channel).Updates(channel).Error
  445. if err != nil {
  446. return err
  447. }
  448. DB.Model(channel).First(channel, "id = ?", channel.Id)
  449. err = channel.UpdateAbilities(nil)
  450. return err
  451. }
  452. func (channel *Channel) UpdateResponseTime(responseTime int64) {
  453. err := DB.Model(channel).Select("response_time", "test_time").Updates(Channel{
  454. TestTime: common.GetTimestamp(),
  455. ResponseTime: int(responseTime),
  456. }).Error
  457. if err != nil {
  458. common.SysLog(fmt.Sprintf("failed to update response time: channel_id=%d, error=%v", channel.Id, err))
  459. }
  460. }
  461. func (channel *Channel) UpdateBalance(balance float64) {
  462. err := DB.Model(channel).Select("balance_updated_time", "balance").Updates(Channel{
  463. BalanceUpdatedTime: common.GetTimestamp(),
  464. Balance: balance,
  465. }).Error
  466. if err != nil {
  467. common.SysLog(fmt.Sprintf("failed to update balance: channel_id=%d, error=%v", channel.Id, err))
  468. }
  469. }
  470. func (channel *Channel) Delete() error {
  471. var err error
  472. err = DB.Delete(channel).Error
  473. if err != nil {
  474. return err
  475. }
  476. err = channel.DeleteAbilities()
  477. return err
  478. }
  479. var channelStatusLock sync.Mutex
  480. // channelPollingLocks stores locks for each channel.id to ensure thread-safe polling
  481. var channelPollingLocks sync.Map
  482. // GetChannelPollingLock returns or creates a mutex for the given channel ID
  483. func GetChannelPollingLock(channelId int) *sync.Mutex {
  484. if lock, exists := channelPollingLocks.Load(channelId); exists {
  485. return lock.(*sync.Mutex)
  486. }
  487. // Create new lock for this channel
  488. newLock := &sync.Mutex{}
  489. actual, _ := channelPollingLocks.LoadOrStore(channelId, newLock)
  490. return actual.(*sync.Mutex)
  491. }
  492. // CleanupChannelPollingLocks removes locks for channels that no longer exist
  493. // This is optional and can be called periodically to prevent memory leaks
  494. func CleanupChannelPollingLocks() {
  495. var activeChannelIds []int
  496. DB.Model(&Channel{}).Pluck("id", &activeChannelIds)
  497. activeChannelSet := make(map[int]bool)
  498. for _, id := range activeChannelIds {
  499. activeChannelSet[id] = true
  500. }
  501. channelPollingLocks.Range(func(key, value interface{}) bool {
  502. channelId := key.(int)
  503. if !activeChannelSet[channelId] {
  504. channelPollingLocks.Delete(channelId)
  505. }
  506. return true
  507. })
  508. }
  509. func handlerMultiKeyUpdate(channel *Channel, usingKey string, status int, reason string) {
  510. keys := channel.GetKeys()
  511. if len(keys) == 0 {
  512. channel.Status = status
  513. } else {
  514. var keyIndex int
  515. for i, key := range keys {
  516. if key == usingKey {
  517. keyIndex = i
  518. break
  519. }
  520. }
  521. if channel.ChannelInfo.MultiKeyStatusList == nil {
  522. channel.ChannelInfo.MultiKeyStatusList = make(map[int]int)
  523. }
  524. if status == common.ChannelStatusEnabled {
  525. delete(channel.ChannelInfo.MultiKeyStatusList, keyIndex)
  526. } else {
  527. channel.ChannelInfo.MultiKeyStatusList[keyIndex] = status
  528. if channel.ChannelInfo.MultiKeyDisabledReason == nil {
  529. channel.ChannelInfo.MultiKeyDisabledReason = make(map[int]string)
  530. }
  531. if channel.ChannelInfo.MultiKeyDisabledTime == nil {
  532. channel.ChannelInfo.MultiKeyDisabledTime = make(map[int]int64)
  533. }
  534. channel.ChannelInfo.MultiKeyDisabledReason[keyIndex] = reason
  535. channel.ChannelInfo.MultiKeyDisabledTime[keyIndex] = common.GetTimestamp()
  536. }
  537. if len(channel.ChannelInfo.MultiKeyStatusList) >= channel.ChannelInfo.MultiKeySize {
  538. channel.Status = common.ChannelStatusAutoDisabled
  539. info := channel.GetOtherInfo()
  540. info["status_reason"] = "All keys are disabled"
  541. info["status_time"] = common.GetTimestamp()
  542. channel.SetOtherInfo(info)
  543. }
  544. }
  545. }
  546. func UpdateChannelStatus(channelId int, usingKey string, status int, reason string) bool {
  547. if common.MemoryCacheEnabled {
  548. channelStatusLock.Lock()
  549. defer channelStatusLock.Unlock()
  550. channelCache, _ := CacheGetChannel(channelId)
  551. if channelCache == nil {
  552. return false
  553. }
  554. if channelCache.ChannelInfo.IsMultiKey {
  555. // Use per-channel lock to prevent concurrent map read/write with GetNextEnabledKey
  556. pollingLock := GetChannelPollingLock(channelId)
  557. pollingLock.Lock()
  558. // 如果是多Key模式,更新缓存中的状态
  559. handlerMultiKeyUpdate(channelCache, usingKey, status, reason)
  560. pollingLock.Unlock()
  561. //CacheUpdateChannel(channelCache)
  562. //return true
  563. } else {
  564. // 如果缓存渠道存在,且状态已是目标状态,直接返回
  565. if channelCache.Status == status {
  566. return false
  567. }
  568. CacheUpdateChannelStatus(channelId, status)
  569. }
  570. }
  571. shouldUpdateAbilities := false
  572. defer func() {
  573. if shouldUpdateAbilities {
  574. err := UpdateAbilityStatus(channelId, status == common.ChannelStatusEnabled)
  575. if err != nil {
  576. common.SysLog(fmt.Sprintf("failed to update ability status: channel_id=%d, error=%v", channelId, err))
  577. }
  578. }
  579. }()
  580. channel, err := GetChannelById(channelId, true)
  581. if err != nil {
  582. return false
  583. } else {
  584. if channel.Status == status {
  585. return false
  586. }
  587. if channel.ChannelInfo.IsMultiKey {
  588. beforeStatus := channel.Status
  589. // Protect map writes with the same per-channel lock used by readers
  590. pollingLock := GetChannelPollingLock(channelId)
  591. pollingLock.Lock()
  592. handlerMultiKeyUpdate(channel, usingKey, status, reason)
  593. pollingLock.Unlock()
  594. if beforeStatus != channel.Status {
  595. shouldUpdateAbilities = true
  596. }
  597. } else {
  598. info := channel.GetOtherInfo()
  599. info["status_reason"] = reason
  600. info["status_time"] = common.GetTimestamp()
  601. channel.SetOtherInfo(info)
  602. channel.Status = status
  603. shouldUpdateAbilities = true
  604. }
  605. err = channel.SaveWithoutKey()
  606. if err != nil {
  607. common.SysLog(fmt.Sprintf("failed to update channel status: channel_id=%d, status=%d, error=%v", channel.Id, status, err))
  608. return false
  609. }
  610. }
  611. return true
  612. }
  613. func EnableChannelByTag(tag string) error {
  614. err := DB.Model(&Channel{}).Where("tag = ?", tag).Update("status", common.ChannelStatusEnabled).Error
  615. if err != nil {
  616. return err
  617. }
  618. err = UpdateAbilityStatusByTag(tag, true)
  619. return err
  620. }
  621. func DisableChannelByTag(tag string) error {
  622. err := DB.Model(&Channel{}).Where("tag = ?", tag).Update("status", common.ChannelStatusManuallyDisabled).Error
  623. if err != nil {
  624. return err
  625. }
  626. err = UpdateAbilityStatusByTag(tag, false)
  627. return err
  628. }
  629. func EditChannelByTag(tag string, newTag *string, modelMapping *string, models *string, group *string, priority *int64, weight *uint, paramOverride *string, headerOverride *string) error {
  630. updateData := Channel{}
  631. shouldReCreateAbilities := false
  632. updatedTag := tag
  633. // 如果 newTag 不为空且不等于 tag,则更新 tag
  634. if newTag != nil && *newTag != tag {
  635. updateData.Tag = newTag
  636. updatedTag = *newTag
  637. }
  638. if modelMapping != nil && *modelMapping != "" {
  639. updateData.ModelMapping = modelMapping
  640. }
  641. if models != nil && *models != "" {
  642. shouldReCreateAbilities = true
  643. updateData.Models = *models
  644. }
  645. if group != nil && *group != "" {
  646. shouldReCreateAbilities = true
  647. updateData.Group = *group
  648. }
  649. if priority != nil {
  650. updateData.Priority = priority
  651. }
  652. if weight != nil {
  653. updateData.Weight = weight
  654. }
  655. if paramOverride != nil {
  656. updateData.ParamOverride = paramOverride
  657. }
  658. if headerOverride != nil {
  659. updateData.HeaderOverride = headerOverride
  660. }
  661. err := DB.Model(&Channel{}).Where("tag = ?", tag).Updates(updateData).Error
  662. if err != nil {
  663. return err
  664. }
  665. if shouldReCreateAbilities {
  666. channels, err := GetChannelsByTag(updatedTag, false)
  667. if err == nil {
  668. for _, channel := range channels {
  669. err = channel.UpdateAbilities(nil)
  670. if err != nil {
  671. common.SysLog(fmt.Sprintf("failed to update abilities: channel_id=%d, tag=%s, error=%v", channel.Id, channel.GetTag(), err))
  672. }
  673. }
  674. }
  675. } else {
  676. err := UpdateAbilityByTag(tag, newTag, priority, weight)
  677. if err != nil {
  678. return err
  679. }
  680. }
  681. return nil
  682. }
  683. func UpdateChannelUsedQuota(id int, quota int) {
  684. if common.BatchUpdateEnabled {
  685. addNewRecord(BatchUpdateTypeChannelUsedQuota, id, quota)
  686. return
  687. }
  688. updateChannelUsedQuota(id, quota)
  689. }
  690. func updateChannelUsedQuota(id int, quota int) {
  691. err := DB.Model(&Channel{}).Where("id = ?", id).Update("used_quota", gorm.Expr("used_quota + ?", quota)).Error
  692. if err != nil {
  693. common.SysLog(fmt.Sprintf("failed to update channel used quota: channel_id=%d, delta_quota=%d, error=%v", id, quota, err))
  694. }
  695. }
  696. func DeleteChannelByStatus(status int64) (int64, error) {
  697. result := DB.Where("status = ?", status).Delete(&Channel{})
  698. return result.RowsAffected, result.Error
  699. }
  700. func DeleteDisabledChannel() (int64, error) {
  701. result := DB.Where("status = ? or status = ?", common.ChannelStatusAutoDisabled, common.ChannelStatusManuallyDisabled).Delete(&Channel{})
  702. return result.RowsAffected, result.Error
  703. }
  704. func GetPaginatedTags(offset int, limit int) ([]*string, error) {
  705. var tags []*string
  706. err := DB.Model(&Channel{}).Select("DISTINCT tag").Where("tag != ''").Offset(offset).Limit(limit).Find(&tags).Error
  707. return tags, err
  708. }
  709. func SearchTags(keyword string, group string, model string, idSort bool) ([]*string, error) {
  710. var tags []*string
  711. modelsCol := "`models`"
  712. // 如果是 PostgreSQL,使用双引号
  713. if common.UsingPostgreSQL {
  714. modelsCol = `"models"`
  715. }
  716. baseURLCol := "`base_url`"
  717. // 如果是 PostgreSQL,使用双引号
  718. if common.UsingPostgreSQL {
  719. baseURLCol = `"base_url"`
  720. }
  721. order := "priority desc"
  722. if idSort {
  723. order = "id desc"
  724. }
  725. // 构造基础查询
  726. baseQuery := DB.Model(&Channel{}).Omit("key")
  727. // 构造WHERE子句
  728. var whereClause string
  729. var args []interface{}
  730. if group != "" && group != "null" {
  731. var groupCondition string
  732. if common.UsingMySQL {
  733. groupCondition = `CONCAT(',', ` + commonGroupCol + `, ',') LIKE ?`
  734. } else {
  735. // sqlite, PostgreSQL
  736. groupCondition = `(',' || ` + commonGroupCol + ` || ',') LIKE ?`
  737. }
  738. whereClause = "(id = ? OR name LIKE ? OR " + commonKeyCol + " = ? OR " + baseURLCol + " LIKE ?) AND " + modelsCol + ` LIKE ? AND ` + groupCondition
  739. args = append(args, common.String2Int(keyword), "%"+keyword+"%", keyword, "%"+keyword+"%", "%"+model+"%", "%,"+group+",%")
  740. } else {
  741. whereClause = "(id = ? OR name LIKE ? OR " + commonKeyCol + " = ? OR " + baseURLCol + " LIKE ?) AND " + modelsCol + " LIKE ?"
  742. args = append(args, common.String2Int(keyword), "%"+keyword+"%", keyword, "%"+keyword+"%", "%"+model+"%")
  743. }
  744. subQuery := baseQuery.Where(whereClause, args...).
  745. Select("tag").
  746. Where("tag != ''").
  747. Order(order)
  748. err := DB.Table("(?) as sub", subQuery).
  749. Select("DISTINCT tag").
  750. Find(&tags).Error
  751. if err != nil {
  752. return nil, err
  753. }
  754. return tags, nil
  755. }
  756. func (channel *Channel) ValidateSettings() error {
  757. channelParams := &dto.ChannelSettings{}
  758. if channel.Setting != nil && *channel.Setting != "" {
  759. err := common.Unmarshal([]byte(*channel.Setting), channelParams)
  760. if err != nil {
  761. return err
  762. }
  763. }
  764. return nil
  765. }
  766. func (channel *Channel) GetSetting() dto.ChannelSettings {
  767. setting := dto.ChannelSettings{}
  768. if channel.Setting != nil && *channel.Setting != "" {
  769. err := common.Unmarshal([]byte(*channel.Setting), &setting)
  770. if err != nil {
  771. common.SysLog(fmt.Sprintf("failed to unmarshal setting: channel_id=%d, error=%v", channel.Id, err))
  772. channel.Setting = nil // 清空设置以避免后续错误
  773. _ = channel.Save() // 保存修改
  774. }
  775. }
  776. return setting
  777. }
  778. func (channel *Channel) SetSetting(setting dto.ChannelSettings) {
  779. settingBytes, err := common.Marshal(setting)
  780. if err != nil {
  781. common.SysLog(fmt.Sprintf("failed to marshal setting: channel_id=%d, error=%v", channel.Id, err))
  782. return
  783. }
  784. channel.Setting = common.GetPointer[string](string(settingBytes))
  785. }
  786. func (channel *Channel) GetOtherSettings() dto.ChannelOtherSettings {
  787. setting := dto.ChannelOtherSettings{}
  788. if channel.OtherSettings != "" {
  789. err := common.UnmarshalJsonStr(channel.OtherSettings, &setting)
  790. if err != nil {
  791. common.SysLog(fmt.Sprintf("failed to unmarshal setting: channel_id=%d, error=%v", channel.Id, err))
  792. channel.OtherSettings = "{}" // 清空设置以避免后续错误
  793. _ = channel.Save() // 保存修改
  794. }
  795. }
  796. return setting
  797. }
  798. func (channel *Channel) SetOtherSettings(setting dto.ChannelOtherSettings) {
  799. settingBytes, err := common.Marshal(setting)
  800. if err != nil {
  801. common.SysLog(fmt.Sprintf("failed to marshal setting: channel_id=%d, error=%v", channel.Id, err))
  802. return
  803. }
  804. channel.OtherSettings = string(settingBytes)
  805. }
  806. func (channel *Channel) GetParamOverride() map[string]interface{} {
  807. paramOverride := make(map[string]interface{})
  808. if channel.ParamOverride != nil && *channel.ParamOverride != "" {
  809. err := common.Unmarshal([]byte(*channel.ParamOverride), &paramOverride)
  810. if err != nil {
  811. common.SysLog(fmt.Sprintf("failed to unmarshal param override: channel_id=%d, error=%v", channel.Id, err))
  812. }
  813. }
  814. return paramOverride
  815. }
  816. func (channel *Channel) GetHeaderOverride() map[string]interface{} {
  817. headerOverride := make(map[string]interface{})
  818. if channel.HeaderOverride != nil && *channel.HeaderOverride != "" {
  819. err := common.Unmarshal([]byte(*channel.HeaderOverride), &headerOverride)
  820. if err != nil {
  821. common.SysLog(fmt.Sprintf("failed to unmarshal header override: channel_id=%d, error=%v", channel.Id, err))
  822. }
  823. }
  824. return headerOverride
  825. }
  826. func GetChannelsByIds(ids []int) ([]*Channel, error) {
  827. var channels []*Channel
  828. err := DB.Where("id in (?)", ids).Find(&channels).Error
  829. return channels, err
  830. }
  831. func BatchSetChannelTag(ids []int, tag *string) error {
  832. // 开启事务
  833. tx := DB.Begin()
  834. if tx.Error != nil {
  835. return tx.Error
  836. }
  837. // 更新标签
  838. err := tx.Model(&Channel{}).Where("id in (?)", ids).Update("tag", tag).Error
  839. if err != nil {
  840. tx.Rollback()
  841. return err
  842. }
  843. // update ability status
  844. channels, err := GetChannelsByIds(ids)
  845. if err != nil {
  846. tx.Rollback()
  847. return err
  848. }
  849. for _, channel := range channels {
  850. err = channel.UpdateAbilities(tx)
  851. if err != nil {
  852. tx.Rollback()
  853. return err
  854. }
  855. }
  856. // 提交事务
  857. return tx.Commit().Error
  858. }
  859. // CountAllChannels returns total channels in DB
  860. func CountAllChannels() (int64, error) {
  861. var total int64
  862. err := DB.Model(&Channel{}).Count(&total).Error
  863. return total, err
  864. }
  865. // CountAllTags returns number of non-empty distinct tags
  866. func CountAllTags() (int64, error) {
  867. var total int64
  868. err := DB.Model(&Channel{}).Where("tag is not null AND tag != ''").Distinct("tag").Count(&total).Error
  869. return total, err
  870. }
  871. // Get channels of specified type with pagination
  872. func GetChannelsByType(startIdx int, num int, idSort bool, channelType int) ([]*Channel, error) {
  873. var channels []*Channel
  874. order := "priority desc"
  875. if idSort {
  876. order = "id desc"
  877. }
  878. err := DB.Where("type = ?", channelType).Order(order).Limit(num).Offset(startIdx).Omit("key").Find(&channels).Error
  879. return channels, err
  880. }
  881. // Count channels of specific type
  882. func CountChannelsByType(channelType int) (int64, error) {
  883. var count int64
  884. err := DB.Model(&Channel{}).Where("type = ?", channelType).Count(&count).Error
  885. return count, err
  886. }
  887. // Return map[type]count for all channels
  888. func CountChannelsGroupByType() (map[int64]int64, error) {
  889. type result struct {
  890. Type int64 `gorm:"column:type"`
  891. Count int64 `gorm:"column:count"`
  892. }
  893. var results []result
  894. err := DB.Model(&Channel{}).Select("type, count(*) as count").Group("type").Find(&results).Error
  895. if err != nil {
  896. return nil, err
  897. }
  898. counts := make(map[int64]int64)
  899. for _, r := range results {
  900. counts[r.Type] = r.Count
  901. }
  902. return counts, nil
  903. }