channel.go 50 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954
  1. package controller
  2. import (
  3. "context"
  4. "encoding/json"
  5. "fmt"
  6. "net/http"
  7. "strconv"
  8. "strings"
  9. "time"
  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/model"
  14. relaychannel "github.com/QuantumNous/new-api/relay/channel"
  15. "github.com/QuantumNous/new-api/relay/channel/gemini"
  16. "github.com/QuantumNous/new-api/relay/channel/ollama"
  17. "github.com/QuantumNous/new-api/service"
  18. "github.com/gin-gonic/gin"
  19. )
  20. type OpenAIModel struct {
  21. ID string `json:"id"`
  22. Object string `json:"object"`
  23. Created int64 `json:"created"`
  24. OwnedBy string `json:"owned_by"`
  25. Metadata map[string]any `json:"metadata,omitempty"`
  26. Permission []struct {
  27. ID string `json:"id"`
  28. Object string `json:"object"`
  29. Created int64 `json:"created"`
  30. AllowCreateEngine bool `json:"allow_create_engine"`
  31. AllowSampling bool `json:"allow_sampling"`
  32. AllowLogprobs bool `json:"allow_logprobs"`
  33. AllowSearchIndices bool `json:"allow_search_indices"`
  34. AllowView bool `json:"allow_view"`
  35. AllowFineTuning bool `json:"allow_fine_tuning"`
  36. Organization string `json:"organization"`
  37. Group string `json:"group"`
  38. IsBlocking bool `json:"is_blocking"`
  39. } `json:"permission"`
  40. Root string `json:"root"`
  41. Parent string `json:"parent"`
  42. }
  43. type OpenAIModelsResponse struct {
  44. Data []OpenAIModel `json:"data"`
  45. Success bool `json:"success"`
  46. }
  47. func parseStatusFilter(statusParam string) int {
  48. switch strings.ToLower(statusParam) {
  49. case "enabled", "1":
  50. return common.ChannelStatusEnabled
  51. case "disabled", "0":
  52. return 0
  53. default:
  54. return -1
  55. }
  56. }
  57. func clearChannelInfo(channel *model.Channel) {
  58. if channel.ChannelInfo.IsMultiKey {
  59. channel.ChannelInfo.MultiKeyDisabledReason = nil
  60. channel.ChannelInfo.MultiKeyDisabledTime = nil
  61. }
  62. }
  63. func GetAllChannels(c *gin.Context) {
  64. pageInfo := common.GetPageQuery(c)
  65. channelData := make([]*model.Channel, 0)
  66. idSort, _ := strconv.ParseBool(c.Query("id_sort"))
  67. sortOptions := model.NewChannelSortOptions(c.Query("sort_by"), c.Query("sort_order"), idSort)
  68. enableTagMode, _ := strconv.ParseBool(c.Query("tag_mode"))
  69. statusParam := c.Query("status")
  70. // statusFilter: -1 all, 1 enabled, 0 disabled (include auto & manual)
  71. statusFilter := parseStatusFilter(statusParam)
  72. // type filter
  73. typeStr := c.Query("type")
  74. typeFilter := -1
  75. if typeStr != "" {
  76. if t, err := strconv.Atoi(typeStr); err == nil {
  77. typeFilter = t
  78. }
  79. }
  80. var total int64
  81. if enableTagMode {
  82. tags, err := model.GetPaginatedTags(pageInfo.GetStartIdx(), pageInfo.GetPageSize())
  83. if err != nil {
  84. common.SysError("failed to get paginated tags: " + err.Error())
  85. c.JSON(http.StatusOK, gin.H{"success": false, "message": "获取标签失败,请稍后重试"})
  86. return
  87. }
  88. for _, tag := range tags {
  89. if tag == nil || *tag == "" {
  90. continue
  91. }
  92. tagChannels, err := model.GetChannelsByTag(*tag, idSort, false, sortOptions)
  93. if err != nil {
  94. continue
  95. }
  96. filtered := make([]*model.Channel, 0)
  97. for _, ch := range tagChannels {
  98. if statusFilter == common.ChannelStatusEnabled && ch.Status != common.ChannelStatusEnabled {
  99. continue
  100. }
  101. if statusFilter == 0 && ch.Status == common.ChannelStatusEnabled {
  102. continue
  103. }
  104. if typeFilter >= 0 && ch.Type != typeFilter {
  105. continue
  106. }
  107. filtered = append(filtered, ch)
  108. }
  109. channelData = append(channelData, filtered...)
  110. }
  111. total, _ = model.CountAllTags()
  112. } else {
  113. baseQuery := model.DB.Model(&model.Channel{})
  114. if typeFilter >= 0 {
  115. baseQuery = baseQuery.Where("type = ?", typeFilter)
  116. }
  117. if statusFilter == common.ChannelStatusEnabled {
  118. baseQuery = baseQuery.Where("status = ?", common.ChannelStatusEnabled)
  119. } else if statusFilter == 0 {
  120. baseQuery = baseQuery.Where("status != ?", common.ChannelStatusEnabled)
  121. }
  122. baseQuery.Count(&total)
  123. err := sortOptions.Apply(baseQuery).Limit(pageInfo.GetPageSize()).Offset(pageInfo.GetStartIdx()).Omit("key").Find(&channelData).Error
  124. if err != nil {
  125. common.SysError("failed to get channels: " + err.Error())
  126. c.JSON(http.StatusOK, gin.H{"success": false, "message": "获取渠道列表失败,请稍后重试"})
  127. return
  128. }
  129. }
  130. for _, datum := range channelData {
  131. clearChannelInfo(datum)
  132. }
  133. countQuery := model.DB.Model(&model.Channel{})
  134. if statusFilter == common.ChannelStatusEnabled {
  135. countQuery = countQuery.Where("status = ?", common.ChannelStatusEnabled)
  136. } else if statusFilter == 0 {
  137. countQuery = countQuery.Where("status != ?", common.ChannelStatusEnabled)
  138. }
  139. var results []struct {
  140. Type int64
  141. Count int64
  142. }
  143. _ = countQuery.Select("type, count(*) as count").Group("type").Find(&results).Error
  144. typeCounts := make(map[int64]int64)
  145. for _, r := range results {
  146. typeCounts[r.Type] = r.Count
  147. }
  148. common.ApiSuccess(c, gin.H{
  149. "items": channelData,
  150. "total": total,
  151. "page": pageInfo.GetPage(),
  152. "page_size": pageInfo.GetPageSize(),
  153. "type_counts": typeCounts,
  154. })
  155. return
  156. }
  157. func buildFetchModelsHeaders(channel *model.Channel, key string) (http.Header, error) {
  158. var headers http.Header
  159. switch channel.Type {
  160. case constant.ChannelTypeAnthropic:
  161. headers = GetClaudeAuthHeader(key)
  162. default:
  163. headers = GetAuthHeader(key)
  164. }
  165. headerOverride := channel.GetHeaderOverride()
  166. for k, v := range headerOverride {
  167. if relaychannel.IsHeaderPassthroughRuleKey(k) {
  168. continue
  169. }
  170. str, ok := v.(string)
  171. if !ok {
  172. return nil, fmt.Errorf("invalid header override for key %s", k)
  173. }
  174. if strings.Contains(str, "{api_key}") {
  175. str = strings.ReplaceAll(str, "{api_key}", key)
  176. }
  177. headers.Set(k, str)
  178. }
  179. return headers, nil
  180. }
  181. func FetchUpstreamModels(c *gin.Context) {
  182. id, err := strconv.Atoi(c.Param("id"))
  183. if err != nil {
  184. common.ApiError(c, err)
  185. return
  186. }
  187. channel, err := model.GetChannelById(id, true)
  188. if err != nil {
  189. common.ApiError(c, err)
  190. return
  191. }
  192. ids, err := fetchChannelUpstreamModelIDs(channel)
  193. if err != nil {
  194. c.JSON(http.StatusOK, gin.H{
  195. "success": false,
  196. "message": fmt.Sprintf("获取模型列表失败: %s", err.Error()),
  197. })
  198. return
  199. }
  200. c.JSON(http.StatusOK, gin.H{
  201. "success": true,
  202. "message": "",
  203. "data": ids,
  204. })
  205. }
  206. func FixChannelsAbilities(c *gin.Context) {
  207. success, fails, err := model.FixAbility()
  208. if err != nil {
  209. common.ApiError(c, err)
  210. return
  211. }
  212. c.JSON(http.StatusOK, gin.H{
  213. "success": true,
  214. "message": "",
  215. "data": gin.H{
  216. "success": success,
  217. "fails": fails,
  218. },
  219. })
  220. }
  221. func SearchChannels(c *gin.Context) {
  222. keyword := c.Query("keyword")
  223. group := c.Query("group")
  224. modelKeyword := c.Query("model")
  225. statusParam := c.Query("status")
  226. statusFilter := parseStatusFilter(statusParam)
  227. idSort, _ := strconv.ParseBool(c.Query("id_sort"))
  228. sortOptions := model.NewChannelSortOptions(c.Query("sort_by"), c.Query("sort_order"), idSort)
  229. enableTagMode, _ := strconv.ParseBool(c.Query("tag_mode"))
  230. channelData := make([]*model.Channel, 0)
  231. if enableTagMode {
  232. tags, err := model.SearchTags(keyword, group, modelKeyword, idSort)
  233. if err != nil {
  234. c.JSON(http.StatusOK, gin.H{
  235. "success": false,
  236. "message": err.Error(),
  237. })
  238. return
  239. }
  240. for _, tag := range tags {
  241. if tag != nil && *tag != "" {
  242. tagChannel, err := model.GetChannelsByTag(*tag, idSort, false, sortOptions)
  243. if err == nil {
  244. channelData = append(channelData, tagChannel...)
  245. }
  246. }
  247. }
  248. } else {
  249. channels, err := model.SearchChannels(keyword, group, modelKeyword, idSort, sortOptions)
  250. if err != nil {
  251. c.JSON(http.StatusOK, gin.H{
  252. "success": false,
  253. "message": err.Error(),
  254. })
  255. return
  256. }
  257. channelData = channels
  258. }
  259. if statusFilter == common.ChannelStatusEnabled || statusFilter == 0 {
  260. filtered := make([]*model.Channel, 0, len(channelData))
  261. for _, ch := range channelData {
  262. if statusFilter == common.ChannelStatusEnabled && ch.Status != common.ChannelStatusEnabled {
  263. continue
  264. }
  265. if statusFilter == 0 && ch.Status == common.ChannelStatusEnabled {
  266. continue
  267. }
  268. filtered = append(filtered, ch)
  269. }
  270. channelData = filtered
  271. }
  272. // calculate type counts for search results
  273. typeCounts := make(map[int64]int64)
  274. for _, channel := range channelData {
  275. typeCounts[int64(channel.Type)]++
  276. }
  277. typeParam := c.Query("type")
  278. typeFilter := -1
  279. if typeParam != "" {
  280. if tp, err := strconv.Atoi(typeParam); err == nil {
  281. typeFilter = tp
  282. }
  283. }
  284. if typeFilter >= 0 {
  285. filtered := make([]*model.Channel, 0, len(channelData))
  286. for _, ch := range channelData {
  287. if ch.Type == typeFilter {
  288. filtered = append(filtered, ch)
  289. }
  290. }
  291. channelData = filtered
  292. }
  293. page, _ := strconv.Atoi(c.DefaultQuery("p", "1"))
  294. pageSize, _ := strconv.Atoi(c.DefaultQuery("page_size", "20"))
  295. if page < 1 {
  296. page = 1
  297. }
  298. if pageSize <= 0 {
  299. pageSize = 20
  300. }
  301. total := len(channelData)
  302. startIdx := (page - 1) * pageSize
  303. if startIdx > total {
  304. startIdx = total
  305. }
  306. endIdx := startIdx + pageSize
  307. if endIdx > total {
  308. endIdx = total
  309. }
  310. pagedData := channelData[startIdx:endIdx]
  311. for _, datum := range pagedData {
  312. clearChannelInfo(datum)
  313. }
  314. c.JSON(http.StatusOK, gin.H{
  315. "success": true,
  316. "message": "",
  317. "data": gin.H{
  318. "items": pagedData,
  319. "total": total,
  320. "type_counts": typeCounts,
  321. },
  322. })
  323. return
  324. }
  325. func GetChannel(c *gin.Context) {
  326. id, err := strconv.Atoi(c.Param("id"))
  327. if err != nil {
  328. common.ApiError(c, err)
  329. return
  330. }
  331. channel, err := model.GetChannelById(id, false)
  332. if err != nil {
  333. common.ApiError(c, err)
  334. return
  335. }
  336. if channel != nil {
  337. clearChannelInfo(channel)
  338. }
  339. c.JSON(http.StatusOK, gin.H{
  340. "success": true,
  341. "message": "",
  342. "data": channel,
  343. })
  344. return
  345. }
  346. // GetChannelKey 获取渠道密钥(需要通过安全验证中间件)
  347. // 此函数依赖 SecureVerificationRequired 中间件,确保用户已通过安全验证
  348. func GetChannelKey(c *gin.Context) {
  349. userId := c.GetInt("id")
  350. channelId, err := strconv.Atoi(c.Param("id"))
  351. if err != nil {
  352. common.ApiError(c, fmt.Errorf("渠道ID格式错误: %v", err))
  353. return
  354. }
  355. // 获取渠道信息(包含密钥)
  356. channel, err := model.GetChannelById(channelId, true)
  357. if err != nil {
  358. common.ApiError(c, fmt.Errorf("获取渠道信息失败: %v", err))
  359. return
  360. }
  361. if channel == nil {
  362. common.ApiError(c, fmt.Errorf("渠道不存在"))
  363. return
  364. }
  365. // 记录操作日志
  366. model.RecordLog(userId, model.LogTypeSystem, fmt.Sprintf("查看渠道密钥信息 (渠道ID: %d)", channelId))
  367. // 返回渠道密钥
  368. c.JSON(http.StatusOK, gin.H{
  369. "success": true,
  370. "message": "获取成功",
  371. "data": map[string]interface{}{
  372. "key": channel.Key,
  373. },
  374. })
  375. }
  376. // validateTwoFactorAuth 统一的2FA验证函数
  377. func validateTwoFactorAuth(twoFA *model.TwoFA, code string) bool {
  378. // 尝试验证TOTP
  379. if cleanCode, err := common.ValidateNumericCode(code); err == nil {
  380. if isValid, _ := twoFA.ValidateTOTPAndUpdateUsage(cleanCode); isValid {
  381. return true
  382. }
  383. }
  384. // 尝试验证备用码
  385. if isValid, err := twoFA.ValidateBackupCodeAndUpdateUsage(code); err == nil && isValid {
  386. return true
  387. }
  388. return false
  389. }
  390. // validateChannel 通用的渠道校验函数
  391. func validateChannel(channel *model.Channel, isAdd bool) error {
  392. // 校验 channel settings
  393. if err := channel.ValidateSettings(); err != nil {
  394. return fmt.Errorf("渠道额外设置[channel setting] 格式错误:%s", err.Error())
  395. }
  396. // 如果是添加操作,检查 channel 和 key 是否为空
  397. if isAdd {
  398. if channel == nil || channel.Key == "" {
  399. return fmt.Errorf("channel cannot be empty")
  400. }
  401. // 检查模型名称长度是否超过 255
  402. for _, m := range channel.GetModels() {
  403. if len(m) > 255 {
  404. return fmt.Errorf("模型名称过长: %s", m)
  405. }
  406. }
  407. }
  408. // VertexAI 特殊校验
  409. if channel.Type == constant.ChannelTypeVertexAi {
  410. if channel.Other == "" {
  411. return fmt.Errorf("部署地区不能为空")
  412. }
  413. regionMap, err := common.StrToMap(channel.Other)
  414. if err != nil {
  415. return fmt.Errorf("部署地区必须是标准的Json格式,例如{\"default\": \"us-central1\", \"region2\": \"us-east1\"}")
  416. }
  417. if regionMap["default"] == nil {
  418. return fmt.Errorf("部署地区必须包含default字段")
  419. }
  420. }
  421. // Codex OAuth key validation (optional, only when JSON object is provided)
  422. if channel.Type == constant.ChannelTypeCodex {
  423. trimmedKey := strings.TrimSpace(channel.Key)
  424. if isAdd || trimmedKey != "" {
  425. if !strings.HasPrefix(trimmedKey, "{") {
  426. return fmt.Errorf("Codex key must be a valid JSON object")
  427. }
  428. var keyMap map[string]any
  429. if err := common.Unmarshal([]byte(trimmedKey), &keyMap); err != nil {
  430. return fmt.Errorf("Codex key must be a valid JSON object")
  431. }
  432. if v, ok := keyMap["access_token"]; !ok || v == nil || strings.TrimSpace(fmt.Sprintf("%v", v)) == "" {
  433. return fmt.Errorf("Codex key JSON must include access_token")
  434. }
  435. if v, ok := keyMap["account_id"]; !ok || v == nil || strings.TrimSpace(fmt.Sprintf("%v", v)) == "" {
  436. return fmt.Errorf("Codex key JSON must include account_id")
  437. }
  438. }
  439. }
  440. return nil
  441. }
  442. func RefreshCodexChannelCredential(c *gin.Context) {
  443. channelId, err := strconv.Atoi(c.Param("id"))
  444. if err != nil {
  445. common.ApiError(c, fmt.Errorf("invalid channel id: %w", err))
  446. return
  447. }
  448. ctx, cancel := context.WithTimeout(c.Request.Context(), 10*time.Second)
  449. defer cancel()
  450. oauthKey, ch, err := service.RefreshCodexChannelCredential(ctx, channelId, service.CodexCredentialRefreshOptions{ResetCaches: true})
  451. if err != nil {
  452. common.SysError("failed to refresh codex channel credential: " + err.Error())
  453. c.JSON(http.StatusOK, gin.H{"success": false, "message": "刷新凭证失败,请稍后重试"})
  454. return
  455. }
  456. c.JSON(http.StatusOK, gin.H{
  457. "success": true,
  458. "message": "refreshed",
  459. "data": gin.H{
  460. "expires_at": oauthKey.Expired,
  461. "last_refresh": oauthKey.LastRefresh,
  462. "account_id": oauthKey.AccountID,
  463. "email": oauthKey.Email,
  464. "channel_id": ch.Id,
  465. "channel_type": ch.Type,
  466. "channel_name": ch.Name,
  467. },
  468. })
  469. }
  470. type AddChannelRequest struct {
  471. Mode string `json:"mode"`
  472. MultiKeyMode constant.MultiKeyMode `json:"multi_key_mode"`
  473. BatchAddSetKeyPrefix2Name bool `json:"batch_add_set_key_prefix_2_name"`
  474. Channel *model.Channel `json:"channel"`
  475. }
  476. func getVertexArrayKeys(keys string) ([]string, error) {
  477. if keys == "" {
  478. return nil, nil
  479. }
  480. var keyArray []interface{}
  481. err := common.Unmarshal([]byte(keys), &keyArray)
  482. if err != nil {
  483. return nil, fmt.Errorf("批量添加 Vertex AI 必须使用标准的JsonArray格式,例如[{key1}, {key2}...],请检查输入: %w", err)
  484. }
  485. cleanKeys := make([]string, 0, len(keyArray))
  486. for _, key := range keyArray {
  487. var keyStr string
  488. switch v := key.(type) {
  489. case string:
  490. keyStr = strings.TrimSpace(v)
  491. default:
  492. bytes, err := json.Marshal(v)
  493. if err != nil {
  494. return nil, fmt.Errorf("Vertex AI key JSON 编码失败: %w", err)
  495. }
  496. keyStr = string(bytes)
  497. }
  498. if keyStr != "" {
  499. cleanKeys = append(cleanKeys, keyStr)
  500. }
  501. }
  502. if len(cleanKeys) == 0 {
  503. return nil, fmt.Errorf("批量添加 Vertex AI 的 keys 不能为空")
  504. }
  505. return cleanKeys, nil
  506. }
  507. func AddChannel(c *gin.Context) {
  508. addChannelRequest := AddChannelRequest{}
  509. err := c.ShouldBindJSON(&addChannelRequest)
  510. if err != nil {
  511. common.ApiError(c, err)
  512. return
  513. }
  514. // 使用统一的校验函数
  515. if err := validateChannel(addChannelRequest.Channel, true); err != nil {
  516. c.JSON(http.StatusOK, gin.H{
  517. "success": false,
  518. "message": err.Error(),
  519. })
  520. return
  521. }
  522. addChannelRequest.Channel.CreatedTime = common.GetTimestamp()
  523. keys := make([]string, 0)
  524. switch addChannelRequest.Mode {
  525. case "multi_to_single":
  526. addChannelRequest.Channel.ChannelInfo.IsMultiKey = true
  527. addChannelRequest.Channel.ChannelInfo.MultiKeyMode = addChannelRequest.MultiKeyMode
  528. if addChannelRequest.Channel.Type == constant.ChannelTypeVertexAi && addChannelRequest.Channel.GetOtherSettings().VertexKeyType != dto.VertexKeyTypeAPIKey {
  529. array, err := getVertexArrayKeys(addChannelRequest.Channel.Key)
  530. if err != nil {
  531. c.JSON(http.StatusOK, gin.H{
  532. "success": false,
  533. "message": err.Error(),
  534. })
  535. return
  536. }
  537. addChannelRequest.Channel.ChannelInfo.MultiKeySize = len(array)
  538. addChannelRequest.Channel.Key = strings.Join(array, "\n")
  539. } else {
  540. cleanKeys := make([]string, 0)
  541. for _, key := range strings.Split(addChannelRequest.Channel.Key, "\n") {
  542. if key == "" {
  543. continue
  544. }
  545. key = strings.TrimSpace(key)
  546. cleanKeys = append(cleanKeys, key)
  547. }
  548. addChannelRequest.Channel.ChannelInfo.MultiKeySize = len(cleanKeys)
  549. addChannelRequest.Channel.Key = strings.Join(cleanKeys, "\n")
  550. }
  551. keys = []string{addChannelRequest.Channel.Key}
  552. case "batch":
  553. if addChannelRequest.Channel.Type == constant.ChannelTypeVertexAi && addChannelRequest.Channel.GetOtherSettings().VertexKeyType != dto.VertexKeyTypeAPIKey {
  554. // multi json
  555. keys, err = getVertexArrayKeys(addChannelRequest.Channel.Key)
  556. if err != nil {
  557. c.JSON(http.StatusOK, gin.H{
  558. "success": false,
  559. "message": err.Error(),
  560. })
  561. return
  562. }
  563. } else {
  564. keys = strings.Split(addChannelRequest.Channel.Key, "\n")
  565. }
  566. case "single":
  567. keys = []string{addChannelRequest.Channel.Key}
  568. default:
  569. c.JSON(http.StatusOK, gin.H{
  570. "success": false,
  571. "message": "不支持的添加模式",
  572. })
  573. return
  574. }
  575. channels := make([]model.Channel, 0, len(keys))
  576. for _, key := range keys {
  577. if key == "" {
  578. continue
  579. }
  580. localChannel := addChannelRequest.Channel
  581. localChannel.Key = key
  582. if addChannelRequest.BatchAddSetKeyPrefix2Name && len(keys) > 1 {
  583. keyPrefix := localChannel.Key
  584. if len(localChannel.Key) > 8 {
  585. keyPrefix = localChannel.Key[:8]
  586. }
  587. localChannel.Name = fmt.Sprintf("%s %s", localChannel.Name, keyPrefix)
  588. }
  589. channels = append(channels, *localChannel)
  590. }
  591. err = model.BatchInsertChannels(channels)
  592. if err != nil {
  593. common.ApiError(c, err)
  594. return
  595. }
  596. service.ResetProxyClientCache()
  597. c.JSON(http.StatusOK, gin.H{
  598. "success": true,
  599. "message": "",
  600. })
  601. return
  602. }
  603. func DeleteChannel(c *gin.Context) {
  604. id, _ := strconv.Atoi(c.Param("id"))
  605. channel := model.Channel{Id: id}
  606. err := channel.Delete()
  607. if err != nil {
  608. common.ApiError(c, err)
  609. return
  610. }
  611. model.InitChannelCache()
  612. c.JSON(http.StatusOK, gin.H{
  613. "success": true,
  614. "message": "",
  615. })
  616. return
  617. }
  618. func DeleteDisabledChannel(c *gin.Context) {
  619. rows, err := model.DeleteDisabledChannel()
  620. if err != nil {
  621. common.ApiError(c, err)
  622. return
  623. }
  624. model.InitChannelCache()
  625. c.JSON(http.StatusOK, gin.H{
  626. "success": true,
  627. "message": "",
  628. "data": rows,
  629. })
  630. return
  631. }
  632. type ChannelTag struct {
  633. Tag string `json:"tag"`
  634. NewTag *string `json:"new_tag"`
  635. Priority *int64 `json:"priority"`
  636. Weight *uint `json:"weight"`
  637. ModelMapping *string `json:"model_mapping"`
  638. Models *string `json:"models"`
  639. Groups *string `json:"groups"`
  640. ParamOverride *string `json:"param_override"`
  641. HeaderOverride *string `json:"header_override"`
  642. }
  643. func DisableTagChannels(c *gin.Context) {
  644. channelTag := ChannelTag{}
  645. err := c.ShouldBindJSON(&channelTag)
  646. if err != nil || channelTag.Tag == "" {
  647. c.JSON(http.StatusOK, gin.H{
  648. "success": false,
  649. "message": "参数错误",
  650. })
  651. return
  652. }
  653. err = model.DisableChannelByTag(channelTag.Tag)
  654. if err != nil {
  655. common.ApiError(c, err)
  656. return
  657. }
  658. model.InitChannelCache()
  659. c.JSON(http.StatusOK, gin.H{
  660. "success": true,
  661. "message": "",
  662. })
  663. return
  664. }
  665. func EnableTagChannels(c *gin.Context) {
  666. channelTag := ChannelTag{}
  667. err := c.ShouldBindJSON(&channelTag)
  668. if err != nil || channelTag.Tag == "" {
  669. c.JSON(http.StatusOK, gin.H{
  670. "success": false,
  671. "message": "参数错误",
  672. })
  673. return
  674. }
  675. err = model.EnableChannelByTag(channelTag.Tag)
  676. if err != nil {
  677. common.ApiError(c, err)
  678. return
  679. }
  680. model.InitChannelCache()
  681. c.JSON(http.StatusOK, gin.H{
  682. "success": true,
  683. "message": "",
  684. })
  685. return
  686. }
  687. func EditTagChannels(c *gin.Context) {
  688. channelTag := ChannelTag{}
  689. err := c.ShouldBindJSON(&channelTag)
  690. if err != nil {
  691. c.JSON(http.StatusOK, gin.H{
  692. "success": false,
  693. "message": "参数错误",
  694. })
  695. return
  696. }
  697. if channelTag.Tag == "" {
  698. c.JSON(http.StatusOK, gin.H{
  699. "success": false,
  700. "message": "tag不能为空",
  701. })
  702. return
  703. }
  704. if channelTag.ParamOverride != nil {
  705. trimmed := strings.TrimSpace(*channelTag.ParamOverride)
  706. if trimmed != "" && !json.Valid([]byte(trimmed)) {
  707. c.JSON(http.StatusOK, gin.H{
  708. "success": false,
  709. "message": "参数覆盖必须是合法的 JSON 格式",
  710. })
  711. return
  712. }
  713. channelTag.ParamOverride = common.GetPointer[string](trimmed)
  714. }
  715. if channelTag.HeaderOverride != nil {
  716. trimmed := strings.TrimSpace(*channelTag.HeaderOverride)
  717. if trimmed != "" && !json.Valid([]byte(trimmed)) {
  718. c.JSON(http.StatusOK, gin.H{
  719. "success": false,
  720. "message": "请求头覆盖必须是合法的 JSON 格式",
  721. })
  722. return
  723. }
  724. channelTag.HeaderOverride = common.GetPointer[string](trimmed)
  725. }
  726. err = model.EditChannelByTag(channelTag.Tag, channelTag.NewTag, channelTag.ModelMapping, channelTag.Models, channelTag.Groups, channelTag.Priority, channelTag.Weight, channelTag.ParamOverride, channelTag.HeaderOverride)
  727. if err != nil {
  728. common.ApiError(c, err)
  729. return
  730. }
  731. model.InitChannelCache()
  732. c.JSON(http.StatusOK, gin.H{
  733. "success": true,
  734. "message": "",
  735. })
  736. return
  737. }
  738. type ChannelBatch struct {
  739. Ids []int `json:"ids"`
  740. Tag *string `json:"tag"`
  741. }
  742. func DeleteChannelBatch(c *gin.Context) {
  743. channelBatch := ChannelBatch{}
  744. err := c.ShouldBindJSON(&channelBatch)
  745. if err != nil || len(channelBatch.Ids) == 0 {
  746. c.JSON(http.StatusOK, gin.H{
  747. "success": false,
  748. "message": "参数错误",
  749. })
  750. return
  751. }
  752. err = model.BatchDeleteChannels(channelBatch.Ids)
  753. if err != nil {
  754. common.ApiError(c, err)
  755. return
  756. }
  757. model.InitChannelCache()
  758. c.JSON(http.StatusOK, gin.H{
  759. "success": true,
  760. "message": "",
  761. "data": len(channelBatch.Ids),
  762. })
  763. return
  764. }
  765. type PatchChannel struct {
  766. model.Channel
  767. MultiKeyMode *string `json:"multi_key_mode"`
  768. KeyMode *string `json:"key_mode"` // 多key模式下密钥覆盖或者追加
  769. }
  770. func UpdateChannel(c *gin.Context) {
  771. channel := PatchChannel{}
  772. err := c.ShouldBindJSON(&channel)
  773. if err != nil {
  774. common.ApiError(c, err)
  775. return
  776. }
  777. // 使用统一的校验函数
  778. if err := validateChannel(&channel.Channel, false); err != nil {
  779. c.JSON(http.StatusOK, gin.H{
  780. "success": false,
  781. "message": err.Error(),
  782. })
  783. return
  784. }
  785. // Preserve existing ChannelInfo to ensure multi-key channels keep correct state even if the client does not send ChannelInfo in the request.
  786. originChannel, err := model.GetChannelById(channel.Id, true)
  787. if err != nil {
  788. c.JSON(http.StatusOK, gin.H{
  789. "success": false,
  790. "message": err.Error(),
  791. })
  792. return
  793. }
  794. // Always copy the original ChannelInfo so that fields like IsMultiKey and MultiKeySize are retained.
  795. channel.ChannelInfo = originChannel.ChannelInfo
  796. // If the request explicitly specifies a new MultiKeyMode, apply it on top of the original info.
  797. if channel.MultiKeyMode != nil && *channel.MultiKeyMode != "" {
  798. channel.ChannelInfo.MultiKeyMode = constant.MultiKeyMode(*channel.MultiKeyMode)
  799. }
  800. // 处理多key模式下的密钥追加/覆盖逻辑
  801. if channel.KeyMode != nil && channel.ChannelInfo.IsMultiKey {
  802. switch *channel.KeyMode {
  803. case "append":
  804. // 追加模式:将新密钥添加到现有密钥列表
  805. if originChannel.Key != "" {
  806. var newKeys []string
  807. var existingKeys []string
  808. // 解析现有密钥
  809. if strings.HasPrefix(strings.TrimSpace(originChannel.Key), "[") {
  810. // JSON数组格式
  811. var arr []json.RawMessage
  812. if err := json.Unmarshal([]byte(strings.TrimSpace(originChannel.Key)), &arr); err == nil {
  813. existingKeys = make([]string, len(arr))
  814. for i, v := range arr {
  815. existingKeys[i] = string(v)
  816. }
  817. }
  818. } else {
  819. // 换行分隔格式
  820. existingKeys = strings.Split(strings.Trim(originChannel.Key, "\n"), "\n")
  821. }
  822. // 处理 Vertex AI 的特殊情况
  823. if channel.Type == constant.ChannelTypeVertexAi && channel.GetOtherSettings().VertexKeyType != dto.VertexKeyTypeAPIKey {
  824. // 尝试解析新密钥为JSON数组
  825. if strings.HasPrefix(strings.TrimSpace(channel.Key), "[") {
  826. array, err := getVertexArrayKeys(channel.Key)
  827. if err != nil {
  828. c.JSON(http.StatusOK, gin.H{
  829. "success": false,
  830. "message": "追加密钥解析失败: " + err.Error(),
  831. })
  832. return
  833. }
  834. newKeys = array
  835. } else {
  836. // 单个JSON密钥
  837. newKeys = []string{channel.Key}
  838. }
  839. } else {
  840. // 普通渠道的处理
  841. inputKeys := strings.Split(channel.Key, "\n")
  842. for _, key := range inputKeys {
  843. key = strings.TrimSpace(key)
  844. if key != "" {
  845. newKeys = append(newKeys, key)
  846. }
  847. }
  848. }
  849. seen := make(map[string]struct{}, len(existingKeys)+len(newKeys))
  850. for _, key := range existingKeys {
  851. normalized := strings.TrimSpace(key)
  852. if normalized == "" {
  853. continue
  854. }
  855. seen[normalized] = struct{}{}
  856. }
  857. dedupedNewKeys := make([]string, 0, len(newKeys))
  858. for _, key := range newKeys {
  859. normalized := strings.TrimSpace(key)
  860. if normalized == "" {
  861. continue
  862. }
  863. if _, ok := seen[normalized]; ok {
  864. continue
  865. }
  866. seen[normalized] = struct{}{}
  867. dedupedNewKeys = append(dedupedNewKeys, normalized)
  868. }
  869. allKeys := append(existingKeys, dedupedNewKeys...)
  870. channel.Key = strings.Join(allKeys, "\n")
  871. }
  872. case "replace":
  873. // 覆盖模式:直接使用新密钥(默认行为,不需要特殊处理)
  874. }
  875. }
  876. err = channel.Update()
  877. if err != nil {
  878. common.ApiError(c, err)
  879. return
  880. }
  881. model.InitChannelCache()
  882. service.ResetProxyClientCache()
  883. channel.Key = ""
  884. clearChannelInfo(&channel.Channel)
  885. c.JSON(http.StatusOK, gin.H{
  886. "success": true,
  887. "message": "",
  888. "data": channel,
  889. })
  890. return
  891. }
  892. func FetchModels(c *gin.Context) {
  893. var req struct {
  894. BaseURL string `json:"base_url"`
  895. Type int `json:"type"`
  896. Key string `json:"key"`
  897. }
  898. if err := c.ShouldBindJSON(&req); err != nil {
  899. c.JSON(http.StatusBadRequest, gin.H{
  900. "success": false,
  901. "message": "Invalid request",
  902. })
  903. return
  904. }
  905. baseURL := req.BaseURL
  906. if baseURL == "" {
  907. baseURL = constant.ChannelBaseURLs[req.Type]
  908. }
  909. // remove line breaks and extra spaces.
  910. key := strings.TrimSpace(req.Key)
  911. key = strings.Split(key, "\n")[0]
  912. if req.Type == constant.ChannelTypeOllama {
  913. models, err := ollama.FetchOllamaModels(baseURL, key)
  914. if err != nil {
  915. c.JSON(http.StatusOK, gin.H{
  916. "success": false,
  917. "message": fmt.Sprintf("获取Ollama模型失败: %s", err.Error()),
  918. })
  919. return
  920. }
  921. names := make([]string, 0, len(models))
  922. for _, modelInfo := range models {
  923. names = append(names, modelInfo.Name)
  924. }
  925. c.JSON(http.StatusOK, gin.H{
  926. "success": true,
  927. "data": names,
  928. })
  929. return
  930. }
  931. if req.Type == constant.ChannelTypeGemini {
  932. models, err := gemini.FetchGeminiModels(baseURL, key, "")
  933. if err != nil {
  934. c.JSON(http.StatusOK, gin.H{
  935. "success": false,
  936. "message": fmt.Sprintf("获取Gemini模型失败: %s", err.Error()),
  937. })
  938. return
  939. }
  940. c.JSON(http.StatusOK, gin.H{
  941. "success": true,
  942. "data": models,
  943. })
  944. return
  945. }
  946. client := &http.Client{}
  947. url := fmt.Sprintf("%s/v1/models", baseURL)
  948. request, err := http.NewRequest("GET", url, nil)
  949. if err != nil {
  950. c.JSON(http.StatusInternalServerError, gin.H{
  951. "success": false,
  952. "message": err.Error(),
  953. })
  954. return
  955. }
  956. request.Header.Set("Authorization", "Bearer "+key)
  957. response, err := client.Do(request)
  958. if err != nil {
  959. c.JSON(http.StatusInternalServerError, gin.H{
  960. "success": false,
  961. "message": err.Error(),
  962. })
  963. return
  964. }
  965. //check status code
  966. if response.StatusCode != http.StatusOK {
  967. c.JSON(http.StatusInternalServerError, gin.H{
  968. "success": false,
  969. "message": "Failed to fetch models",
  970. })
  971. return
  972. }
  973. defer response.Body.Close()
  974. var result struct {
  975. Data []struct {
  976. ID string `json:"id"`
  977. } `json:"data"`
  978. }
  979. if err := json.NewDecoder(response.Body).Decode(&result); err != nil {
  980. c.JSON(http.StatusInternalServerError, gin.H{
  981. "success": false,
  982. "message": err.Error(),
  983. })
  984. return
  985. }
  986. var models []string
  987. for _, model := range result.Data {
  988. models = append(models, model.ID)
  989. }
  990. c.JSON(http.StatusOK, gin.H{
  991. "success": true,
  992. "data": models,
  993. })
  994. }
  995. func BatchSetChannelTag(c *gin.Context) {
  996. channelBatch := ChannelBatch{}
  997. err := c.ShouldBindJSON(&channelBatch)
  998. if err != nil || len(channelBatch.Ids) == 0 {
  999. c.JSON(http.StatusOK, gin.H{
  1000. "success": false,
  1001. "message": "参数错误",
  1002. })
  1003. return
  1004. }
  1005. err = model.BatchSetChannelTag(channelBatch.Ids, channelBatch.Tag)
  1006. if err != nil {
  1007. common.ApiError(c, err)
  1008. return
  1009. }
  1010. model.InitChannelCache()
  1011. c.JSON(http.StatusOK, gin.H{
  1012. "success": true,
  1013. "message": "",
  1014. "data": len(channelBatch.Ids),
  1015. })
  1016. return
  1017. }
  1018. func GetTagModels(c *gin.Context) {
  1019. tag := c.Query("tag")
  1020. if tag == "" {
  1021. c.JSON(http.StatusBadRequest, gin.H{
  1022. "success": false,
  1023. "message": "tag不能为空",
  1024. })
  1025. return
  1026. }
  1027. channels, err := model.GetChannelsByTag(tag, false, false) // idSort=false, selectAll=false
  1028. if err != nil {
  1029. c.JSON(http.StatusInternalServerError, gin.H{
  1030. "success": false,
  1031. "message": err.Error(),
  1032. })
  1033. return
  1034. }
  1035. var longestModels string
  1036. maxLength := 0
  1037. // Find the longest models string among all channels with the given tag
  1038. for _, channel := range channels {
  1039. if channel.Models != "" {
  1040. currentModels := strings.Split(channel.Models, ",")
  1041. if len(currentModels) > maxLength {
  1042. maxLength = len(currentModels)
  1043. longestModels = channel.Models
  1044. }
  1045. }
  1046. }
  1047. c.JSON(http.StatusOK, gin.H{
  1048. "success": true,
  1049. "message": "",
  1050. "data": longestModels,
  1051. })
  1052. return
  1053. }
  1054. // CopyChannel handles cloning an existing channel with its key.
  1055. // POST /api/channel/copy/:id
  1056. // Optional query params:
  1057. //
  1058. // suffix - string appended to the original name (default "_复制")
  1059. // reset_balance - bool, when true will reset balance & used_quota to 0 (default true)
  1060. func CopyChannel(c *gin.Context) {
  1061. id, err := strconv.Atoi(c.Param("id"))
  1062. if err != nil {
  1063. c.JSON(http.StatusOK, gin.H{"success": false, "message": "invalid id"})
  1064. return
  1065. }
  1066. suffix := c.DefaultQuery("suffix", "_复制")
  1067. resetBalance := true
  1068. if rbStr := c.DefaultQuery("reset_balance", "true"); rbStr != "" {
  1069. if v, err := strconv.ParseBool(rbStr); err == nil {
  1070. resetBalance = v
  1071. }
  1072. }
  1073. // fetch original channel with key
  1074. origin, err := model.GetChannelById(id, true)
  1075. if err != nil {
  1076. common.SysError("failed to get channel by id: " + err.Error())
  1077. c.JSON(http.StatusOK, gin.H{"success": false, "message": "获取渠道信息失败,请稍后重试"})
  1078. return
  1079. }
  1080. // clone channel
  1081. clone := *origin // shallow copy is sufficient as we will overwrite primitives
  1082. clone.Id = 0 // let DB auto-generate
  1083. clone.CreatedTime = common.GetTimestamp()
  1084. clone.Name = origin.Name + suffix
  1085. clone.TestTime = 0
  1086. clone.ResponseTime = 0
  1087. if resetBalance {
  1088. clone.Balance = 0
  1089. clone.UsedQuota = 0
  1090. }
  1091. // insert
  1092. if err := model.BatchInsertChannels([]model.Channel{clone}); err != nil {
  1093. common.SysError("failed to clone channel: " + err.Error())
  1094. c.JSON(http.StatusOK, gin.H{"success": false, "message": "复制渠道失败,请稍后重试"})
  1095. return
  1096. }
  1097. model.InitChannelCache()
  1098. // success
  1099. c.JSON(http.StatusOK, gin.H{"success": true, "message": "", "data": gin.H{"id": clone.Id}})
  1100. }
  1101. // MultiKeyManageRequest represents the request for multi-key management operations
  1102. type MultiKeyManageRequest struct {
  1103. ChannelId int `json:"channel_id"`
  1104. Action string `json:"action"` // "disable_key", "enable_key", "delete_key", "delete_disabled_keys", "get_key_status"
  1105. KeyIndex *int `json:"key_index,omitempty"` // for disable_key, enable_key, and delete_key actions
  1106. Page int `json:"page,omitempty"` // for get_key_status pagination
  1107. PageSize int `json:"page_size,omitempty"` // for get_key_status pagination
  1108. Status *int `json:"status,omitempty"` // for get_key_status filtering: 1=enabled, 2=manual_disabled, 3=auto_disabled, nil=all
  1109. }
  1110. // MultiKeyStatusResponse represents the response for key status query
  1111. type MultiKeyStatusResponse struct {
  1112. Keys []KeyStatus `json:"keys"`
  1113. Total int `json:"total"`
  1114. Page int `json:"page"`
  1115. PageSize int `json:"page_size"`
  1116. TotalPages int `json:"total_pages"`
  1117. // Statistics
  1118. EnabledCount int `json:"enabled_count"`
  1119. ManualDisabledCount int `json:"manual_disabled_count"`
  1120. AutoDisabledCount int `json:"auto_disabled_count"`
  1121. }
  1122. type KeyStatus struct {
  1123. Index int `json:"index"`
  1124. Status int `json:"status"` // 1: enabled, 2: disabled
  1125. DisabledTime int64 `json:"disabled_time,omitempty"`
  1126. Reason string `json:"reason,omitempty"`
  1127. KeyPreview string `json:"key_preview"` // first 10 chars of key for identification
  1128. }
  1129. // ManageMultiKeys handles multi-key management operations
  1130. func ManageMultiKeys(c *gin.Context) {
  1131. request := MultiKeyManageRequest{}
  1132. err := c.ShouldBindJSON(&request)
  1133. if err != nil {
  1134. common.ApiError(c, err)
  1135. return
  1136. }
  1137. channel, err := model.GetChannelById(request.ChannelId, true)
  1138. if err != nil {
  1139. c.JSON(http.StatusOK, gin.H{
  1140. "success": false,
  1141. "message": "渠道不存在",
  1142. })
  1143. return
  1144. }
  1145. if !channel.ChannelInfo.IsMultiKey {
  1146. c.JSON(http.StatusOK, gin.H{
  1147. "success": false,
  1148. "message": "该渠道不是多密钥模式",
  1149. })
  1150. return
  1151. }
  1152. lock := model.GetChannelPollingLock(channel.Id)
  1153. lock.Lock()
  1154. defer lock.Unlock()
  1155. switch request.Action {
  1156. case "get_key_status":
  1157. keys := channel.GetKeys()
  1158. // Default pagination parameters
  1159. page := request.Page
  1160. pageSize := request.PageSize
  1161. if page <= 0 {
  1162. page = 1
  1163. }
  1164. if pageSize <= 0 {
  1165. pageSize = 50 // Default page size
  1166. }
  1167. // Statistics for all keys (unchanged by filtering)
  1168. var enabledCount, manualDisabledCount, autoDisabledCount int
  1169. // Build all key status data first
  1170. var allKeyStatusList []KeyStatus
  1171. for i, key := range keys {
  1172. status := 1 // default enabled
  1173. var disabledTime int64
  1174. var reason string
  1175. if channel.ChannelInfo.MultiKeyStatusList != nil {
  1176. if s, exists := channel.ChannelInfo.MultiKeyStatusList[i]; exists {
  1177. status = s
  1178. }
  1179. }
  1180. // Count for statistics (all keys)
  1181. switch status {
  1182. case 1:
  1183. enabledCount++
  1184. case 2:
  1185. manualDisabledCount++
  1186. case 3:
  1187. autoDisabledCount++
  1188. }
  1189. if status != 1 {
  1190. if channel.ChannelInfo.MultiKeyDisabledTime != nil {
  1191. disabledTime = channel.ChannelInfo.MultiKeyDisabledTime[i]
  1192. }
  1193. if channel.ChannelInfo.MultiKeyDisabledReason != nil {
  1194. reason = channel.ChannelInfo.MultiKeyDisabledReason[i]
  1195. }
  1196. }
  1197. // Create key preview (first 10 chars)
  1198. keyPreview := key
  1199. if len(key) > 10 {
  1200. keyPreview = key[:10] + "..."
  1201. }
  1202. allKeyStatusList = append(allKeyStatusList, KeyStatus{
  1203. Index: i,
  1204. Status: status,
  1205. DisabledTime: disabledTime,
  1206. Reason: reason,
  1207. KeyPreview: keyPreview,
  1208. })
  1209. }
  1210. // Apply status filter if specified
  1211. var filteredKeyStatusList []KeyStatus
  1212. if request.Status != nil {
  1213. for _, keyStatus := range allKeyStatusList {
  1214. if keyStatus.Status == *request.Status {
  1215. filteredKeyStatusList = append(filteredKeyStatusList, keyStatus)
  1216. }
  1217. }
  1218. } else {
  1219. filteredKeyStatusList = allKeyStatusList
  1220. }
  1221. // Calculate pagination based on filtered results
  1222. filteredTotal := len(filteredKeyStatusList)
  1223. totalPages := (filteredTotal + pageSize - 1) / pageSize
  1224. if totalPages == 0 {
  1225. totalPages = 1
  1226. }
  1227. if page > totalPages {
  1228. page = totalPages
  1229. }
  1230. // Calculate range for current page
  1231. start := (page - 1) * pageSize
  1232. end := start + pageSize
  1233. if end > filteredTotal {
  1234. end = filteredTotal
  1235. }
  1236. // Get the page data
  1237. var pageKeyStatusList []KeyStatus
  1238. if start < filteredTotal {
  1239. pageKeyStatusList = filteredKeyStatusList[start:end]
  1240. }
  1241. c.JSON(http.StatusOK, gin.H{
  1242. "success": true,
  1243. "message": "",
  1244. "data": MultiKeyStatusResponse{
  1245. Keys: pageKeyStatusList,
  1246. Total: filteredTotal, // Total of filtered results
  1247. Page: page,
  1248. PageSize: pageSize,
  1249. TotalPages: totalPages,
  1250. EnabledCount: enabledCount, // Overall statistics
  1251. ManualDisabledCount: manualDisabledCount, // Overall statistics
  1252. AutoDisabledCount: autoDisabledCount, // Overall statistics
  1253. },
  1254. })
  1255. return
  1256. case "disable_key":
  1257. if request.KeyIndex == nil {
  1258. c.JSON(http.StatusOK, gin.H{
  1259. "success": false,
  1260. "message": "未指定要禁用的密钥索引",
  1261. })
  1262. return
  1263. }
  1264. keyIndex := *request.KeyIndex
  1265. if keyIndex < 0 || keyIndex >= channel.ChannelInfo.MultiKeySize {
  1266. c.JSON(http.StatusOK, gin.H{
  1267. "success": false,
  1268. "message": "密钥索引超出范围",
  1269. })
  1270. return
  1271. }
  1272. if channel.ChannelInfo.MultiKeyStatusList == nil {
  1273. channel.ChannelInfo.MultiKeyStatusList = make(map[int]int)
  1274. }
  1275. if channel.ChannelInfo.MultiKeyDisabledTime == nil {
  1276. channel.ChannelInfo.MultiKeyDisabledTime = make(map[int]int64)
  1277. }
  1278. if channel.ChannelInfo.MultiKeyDisabledReason == nil {
  1279. channel.ChannelInfo.MultiKeyDisabledReason = make(map[int]string)
  1280. }
  1281. channel.ChannelInfo.MultiKeyStatusList[keyIndex] = 2 // disabled
  1282. err = channel.Update()
  1283. if err != nil {
  1284. common.ApiError(c, err)
  1285. return
  1286. }
  1287. model.InitChannelCache()
  1288. c.JSON(http.StatusOK, gin.H{
  1289. "success": true,
  1290. "message": "密钥已禁用",
  1291. })
  1292. return
  1293. case "enable_key":
  1294. if request.KeyIndex == nil {
  1295. c.JSON(http.StatusOK, gin.H{
  1296. "success": false,
  1297. "message": "未指定要启用的密钥索引",
  1298. })
  1299. return
  1300. }
  1301. keyIndex := *request.KeyIndex
  1302. if keyIndex < 0 || keyIndex >= channel.ChannelInfo.MultiKeySize {
  1303. c.JSON(http.StatusOK, gin.H{
  1304. "success": false,
  1305. "message": "密钥索引超出范围",
  1306. })
  1307. return
  1308. }
  1309. // 从状态列表中删除该密钥的记录,使其回到默认启用状态
  1310. if channel.ChannelInfo.MultiKeyStatusList != nil {
  1311. delete(channel.ChannelInfo.MultiKeyStatusList, keyIndex)
  1312. }
  1313. if channel.ChannelInfo.MultiKeyDisabledTime != nil {
  1314. delete(channel.ChannelInfo.MultiKeyDisabledTime, keyIndex)
  1315. }
  1316. if channel.ChannelInfo.MultiKeyDisabledReason != nil {
  1317. delete(channel.ChannelInfo.MultiKeyDisabledReason, keyIndex)
  1318. }
  1319. err = channel.Update()
  1320. if err != nil {
  1321. common.ApiError(c, err)
  1322. return
  1323. }
  1324. model.InitChannelCache()
  1325. c.JSON(http.StatusOK, gin.H{
  1326. "success": true,
  1327. "message": "密钥已启用",
  1328. })
  1329. return
  1330. case "enable_all_keys":
  1331. // 清空所有禁用状态,使所有密钥回到默认启用状态
  1332. var enabledCount int
  1333. if channel.ChannelInfo.MultiKeyStatusList != nil {
  1334. enabledCount = len(channel.ChannelInfo.MultiKeyStatusList)
  1335. }
  1336. channel.ChannelInfo.MultiKeyStatusList = make(map[int]int)
  1337. channel.ChannelInfo.MultiKeyDisabledTime = make(map[int]int64)
  1338. channel.ChannelInfo.MultiKeyDisabledReason = make(map[int]string)
  1339. err = channel.Update()
  1340. if err != nil {
  1341. common.ApiError(c, err)
  1342. return
  1343. }
  1344. model.InitChannelCache()
  1345. c.JSON(http.StatusOK, gin.H{
  1346. "success": true,
  1347. "message": fmt.Sprintf("已启用 %d 个密钥", enabledCount),
  1348. })
  1349. return
  1350. case "disable_all_keys":
  1351. // 禁用所有启用的密钥
  1352. if channel.ChannelInfo.MultiKeyStatusList == nil {
  1353. channel.ChannelInfo.MultiKeyStatusList = make(map[int]int)
  1354. }
  1355. if channel.ChannelInfo.MultiKeyDisabledTime == nil {
  1356. channel.ChannelInfo.MultiKeyDisabledTime = make(map[int]int64)
  1357. }
  1358. if channel.ChannelInfo.MultiKeyDisabledReason == nil {
  1359. channel.ChannelInfo.MultiKeyDisabledReason = make(map[int]string)
  1360. }
  1361. var disabledCount int
  1362. for i := 0; i < channel.ChannelInfo.MultiKeySize; i++ {
  1363. status := 1 // default enabled
  1364. if s, exists := channel.ChannelInfo.MultiKeyStatusList[i]; exists {
  1365. status = s
  1366. }
  1367. // 只禁用当前启用的密钥
  1368. if status == 1 {
  1369. channel.ChannelInfo.MultiKeyStatusList[i] = 2 // disabled
  1370. disabledCount++
  1371. }
  1372. }
  1373. if disabledCount == 0 {
  1374. c.JSON(http.StatusOK, gin.H{
  1375. "success": false,
  1376. "message": "没有可禁用的密钥",
  1377. })
  1378. return
  1379. }
  1380. err = channel.Update()
  1381. if err != nil {
  1382. common.ApiError(c, err)
  1383. return
  1384. }
  1385. model.InitChannelCache()
  1386. c.JSON(http.StatusOK, gin.H{
  1387. "success": true,
  1388. "message": fmt.Sprintf("已禁用 %d 个密钥", disabledCount),
  1389. })
  1390. return
  1391. case "delete_key":
  1392. if request.KeyIndex == nil {
  1393. c.JSON(http.StatusOK, gin.H{
  1394. "success": false,
  1395. "message": "未指定要删除的密钥索引",
  1396. })
  1397. return
  1398. }
  1399. keyIndex := *request.KeyIndex
  1400. if keyIndex < 0 || keyIndex >= channel.ChannelInfo.MultiKeySize {
  1401. c.JSON(http.StatusOK, gin.H{
  1402. "success": false,
  1403. "message": "密钥索引超出范围",
  1404. })
  1405. return
  1406. }
  1407. keys := channel.GetKeys()
  1408. var remainingKeys []string
  1409. var newStatusList = make(map[int]int)
  1410. var newDisabledTime = make(map[int]int64)
  1411. var newDisabledReason = make(map[int]string)
  1412. newIndex := 0
  1413. for i, key := range keys {
  1414. // 跳过要删除的密钥
  1415. if i == keyIndex {
  1416. continue
  1417. }
  1418. remainingKeys = append(remainingKeys, key)
  1419. // 保留其他密钥的状态信息,重新索引
  1420. if channel.ChannelInfo.MultiKeyStatusList != nil {
  1421. if status, exists := channel.ChannelInfo.MultiKeyStatusList[i]; exists && status != 1 {
  1422. newStatusList[newIndex] = status
  1423. }
  1424. }
  1425. if channel.ChannelInfo.MultiKeyDisabledTime != nil {
  1426. if t, exists := channel.ChannelInfo.MultiKeyDisabledTime[i]; exists {
  1427. newDisabledTime[newIndex] = t
  1428. }
  1429. }
  1430. if channel.ChannelInfo.MultiKeyDisabledReason != nil {
  1431. if r, exists := channel.ChannelInfo.MultiKeyDisabledReason[i]; exists {
  1432. newDisabledReason[newIndex] = r
  1433. }
  1434. }
  1435. newIndex++
  1436. }
  1437. if len(remainingKeys) == 0 {
  1438. c.JSON(http.StatusOK, gin.H{
  1439. "success": false,
  1440. "message": "不能删除最后一个密钥",
  1441. })
  1442. return
  1443. }
  1444. // Update channel with remaining keys
  1445. channel.Key = strings.Join(remainingKeys, "\n")
  1446. channel.ChannelInfo.MultiKeySize = len(remainingKeys)
  1447. channel.ChannelInfo.MultiKeyStatusList = newStatusList
  1448. channel.ChannelInfo.MultiKeyDisabledTime = newDisabledTime
  1449. channel.ChannelInfo.MultiKeyDisabledReason = newDisabledReason
  1450. err = channel.Update()
  1451. if err != nil {
  1452. common.ApiError(c, err)
  1453. return
  1454. }
  1455. model.InitChannelCache()
  1456. c.JSON(http.StatusOK, gin.H{
  1457. "success": true,
  1458. "message": "密钥已删除",
  1459. })
  1460. return
  1461. case "delete_disabled_keys":
  1462. keys := channel.GetKeys()
  1463. var remainingKeys []string
  1464. var deletedCount int
  1465. var newStatusList = make(map[int]int)
  1466. var newDisabledTime = make(map[int]int64)
  1467. var newDisabledReason = make(map[int]string)
  1468. newIndex := 0
  1469. for i, key := range keys {
  1470. status := 1 // default enabled
  1471. if channel.ChannelInfo.MultiKeyStatusList != nil {
  1472. if s, exists := channel.ChannelInfo.MultiKeyStatusList[i]; exists {
  1473. status = s
  1474. }
  1475. }
  1476. // 只删除自动禁用(status == 3)的密钥,保留启用(status == 1)和手动禁用(status == 2)的密钥
  1477. if status == 3 {
  1478. deletedCount++
  1479. } else {
  1480. remainingKeys = append(remainingKeys, key)
  1481. // 保留非自动禁用密钥的状态信息,重新索引
  1482. if status != 1 {
  1483. newStatusList[newIndex] = status
  1484. if channel.ChannelInfo.MultiKeyDisabledTime != nil {
  1485. if t, exists := channel.ChannelInfo.MultiKeyDisabledTime[i]; exists {
  1486. newDisabledTime[newIndex] = t
  1487. }
  1488. }
  1489. if channel.ChannelInfo.MultiKeyDisabledReason != nil {
  1490. if r, exists := channel.ChannelInfo.MultiKeyDisabledReason[i]; exists {
  1491. newDisabledReason[newIndex] = r
  1492. }
  1493. }
  1494. }
  1495. newIndex++
  1496. }
  1497. }
  1498. if deletedCount == 0 {
  1499. c.JSON(http.StatusOK, gin.H{
  1500. "success": false,
  1501. "message": "没有需要删除的自动禁用密钥",
  1502. })
  1503. return
  1504. }
  1505. // Update channel with remaining keys
  1506. channel.Key = strings.Join(remainingKeys, "\n")
  1507. channel.ChannelInfo.MultiKeySize = len(remainingKeys)
  1508. channel.ChannelInfo.MultiKeyStatusList = newStatusList
  1509. channel.ChannelInfo.MultiKeyDisabledTime = newDisabledTime
  1510. channel.ChannelInfo.MultiKeyDisabledReason = newDisabledReason
  1511. err = channel.Update()
  1512. if err != nil {
  1513. common.ApiError(c, err)
  1514. return
  1515. }
  1516. model.InitChannelCache()
  1517. c.JSON(http.StatusOK, gin.H{
  1518. "success": true,
  1519. "message": fmt.Sprintf("已删除 %d 个自动禁用的密钥", deletedCount),
  1520. "data": deletedCount,
  1521. })
  1522. return
  1523. default:
  1524. c.JSON(http.StatusOK, gin.H{
  1525. "success": false,
  1526. "message": "不支持的操作",
  1527. })
  1528. return
  1529. }
  1530. }
  1531. // OllamaPullModel 拉取 Ollama 模型
  1532. func OllamaPullModel(c *gin.Context) {
  1533. var req struct {
  1534. ChannelID int `json:"channel_id"`
  1535. ModelName string `json:"model_name"`
  1536. }
  1537. if err := c.ShouldBindJSON(&req); err != nil {
  1538. c.JSON(http.StatusBadRequest, gin.H{
  1539. "success": false,
  1540. "message": "Invalid request parameters",
  1541. })
  1542. return
  1543. }
  1544. if req.ChannelID == 0 || req.ModelName == "" {
  1545. c.JSON(http.StatusBadRequest, gin.H{
  1546. "success": false,
  1547. "message": "Channel ID and model name are required",
  1548. })
  1549. return
  1550. }
  1551. // 获取渠道信息
  1552. channel, err := model.GetChannelById(req.ChannelID, true)
  1553. if err != nil {
  1554. c.JSON(http.StatusNotFound, gin.H{
  1555. "success": false,
  1556. "message": "Channel not found",
  1557. })
  1558. return
  1559. }
  1560. // 检查是否是 Ollama 渠道
  1561. if channel.Type != constant.ChannelTypeOllama {
  1562. c.JSON(http.StatusBadRequest, gin.H{
  1563. "success": false,
  1564. "message": "This operation is only supported for Ollama channels",
  1565. })
  1566. return
  1567. }
  1568. baseURL := constant.ChannelBaseURLs[channel.Type]
  1569. if channel.GetBaseURL() != "" {
  1570. baseURL = channel.GetBaseURL()
  1571. }
  1572. key := strings.Split(channel.Key, "\n")[0]
  1573. err = ollama.PullOllamaModel(baseURL, key, req.ModelName)
  1574. if err != nil {
  1575. c.JSON(http.StatusInternalServerError, gin.H{
  1576. "success": false,
  1577. "message": fmt.Sprintf("Failed to pull model: %s", err.Error()),
  1578. })
  1579. return
  1580. }
  1581. c.JSON(http.StatusOK, gin.H{
  1582. "success": true,
  1583. "message": fmt.Sprintf("Model %s pulled successfully", req.ModelName),
  1584. })
  1585. }
  1586. // OllamaPullModelStream 流式拉取 Ollama 模型
  1587. func OllamaPullModelStream(c *gin.Context) {
  1588. var req struct {
  1589. ChannelID int `json:"channel_id"`
  1590. ModelName string `json:"model_name"`
  1591. }
  1592. if err := c.ShouldBindJSON(&req); err != nil {
  1593. c.JSON(http.StatusBadRequest, gin.H{
  1594. "success": false,
  1595. "message": "Invalid request parameters",
  1596. })
  1597. return
  1598. }
  1599. if req.ChannelID == 0 || req.ModelName == "" {
  1600. c.JSON(http.StatusBadRequest, gin.H{
  1601. "success": false,
  1602. "message": "Channel ID and model name are required",
  1603. })
  1604. return
  1605. }
  1606. // 获取渠道信息
  1607. channel, err := model.GetChannelById(req.ChannelID, true)
  1608. if err != nil {
  1609. c.JSON(http.StatusNotFound, gin.H{
  1610. "success": false,
  1611. "message": "Channel not found",
  1612. })
  1613. return
  1614. }
  1615. // 检查是否是 Ollama 渠道
  1616. if channel.Type != constant.ChannelTypeOllama {
  1617. c.JSON(http.StatusBadRequest, gin.H{
  1618. "success": false,
  1619. "message": "This operation is only supported for Ollama channels",
  1620. })
  1621. return
  1622. }
  1623. baseURL := constant.ChannelBaseURLs[channel.Type]
  1624. if channel.GetBaseURL() != "" {
  1625. baseURL = channel.GetBaseURL()
  1626. }
  1627. // 设置 SSE 头部
  1628. c.Header("Content-Type", "text/event-stream")
  1629. c.Header("Cache-Control", "no-cache")
  1630. c.Header("Connection", "keep-alive")
  1631. c.Header("Access-Control-Allow-Origin", "*")
  1632. key := strings.Split(channel.Key, "\n")[0]
  1633. // 创建进度回调函数
  1634. progressCallback := func(progress ollama.OllamaPullResponse) {
  1635. data, _ := json.Marshal(progress)
  1636. fmt.Fprintf(c.Writer, "data: %s\n\n", string(data))
  1637. c.Writer.Flush()
  1638. }
  1639. // 执行拉取
  1640. err = ollama.PullOllamaModelStream(baseURL, key, req.ModelName, progressCallback)
  1641. if err != nil {
  1642. errorData, _ := json.Marshal(gin.H{
  1643. "error": err.Error(),
  1644. })
  1645. fmt.Fprintf(c.Writer, "data: %s\n\n", string(errorData))
  1646. } else {
  1647. successData, _ := json.Marshal(gin.H{
  1648. "message": fmt.Sprintf("Model %s pulled successfully", req.ModelName),
  1649. })
  1650. fmt.Fprintf(c.Writer, "data: %s\n\n", string(successData))
  1651. }
  1652. // 发送结束标志
  1653. fmt.Fprintf(c.Writer, "data: [DONE]\n\n")
  1654. c.Writer.Flush()
  1655. }
  1656. // OllamaDeleteModel 删除 Ollama 模型
  1657. func OllamaDeleteModel(c *gin.Context) {
  1658. var req struct {
  1659. ChannelID int `json:"channel_id"`
  1660. ModelName string `json:"model_name"`
  1661. }
  1662. if err := c.ShouldBindJSON(&req); err != nil {
  1663. c.JSON(http.StatusBadRequest, gin.H{
  1664. "success": false,
  1665. "message": "Invalid request parameters",
  1666. })
  1667. return
  1668. }
  1669. if req.ChannelID == 0 || req.ModelName == "" {
  1670. c.JSON(http.StatusBadRequest, gin.H{
  1671. "success": false,
  1672. "message": "Channel ID and model name are required",
  1673. })
  1674. return
  1675. }
  1676. // 获取渠道信息
  1677. channel, err := model.GetChannelById(req.ChannelID, true)
  1678. if err != nil {
  1679. c.JSON(http.StatusNotFound, gin.H{
  1680. "success": false,
  1681. "message": "Channel not found",
  1682. })
  1683. return
  1684. }
  1685. // 检查是否是 Ollama 渠道
  1686. if channel.Type != constant.ChannelTypeOllama {
  1687. c.JSON(http.StatusBadRequest, gin.H{
  1688. "success": false,
  1689. "message": "This operation is only supported for Ollama channels",
  1690. })
  1691. return
  1692. }
  1693. baseURL := constant.ChannelBaseURLs[channel.Type]
  1694. if channel.GetBaseURL() != "" {
  1695. baseURL = channel.GetBaseURL()
  1696. }
  1697. key := strings.Split(channel.Key, "\n")[0]
  1698. err = ollama.DeleteOllamaModel(baseURL, key, req.ModelName)
  1699. if err != nil {
  1700. c.JSON(http.StatusInternalServerError, gin.H{
  1701. "success": false,
  1702. "message": fmt.Sprintf("Failed to delete model: %s", err.Error()),
  1703. })
  1704. return
  1705. }
  1706. c.JSON(http.StatusOK, gin.H{
  1707. "success": true,
  1708. "message": fmt.Sprintf("Model %s deleted successfully", req.ModelName),
  1709. })
  1710. }
  1711. // OllamaVersion 获取 Ollama 服务版本信息
  1712. func OllamaVersion(c *gin.Context) {
  1713. id, err := strconv.Atoi(c.Param("id"))
  1714. if err != nil {
  1715. c.JSON(http.StatusBadRequest, gin.H{
  1716. "success": false,
  1717. "message": "Invalid channel id",
  1718. })
  1719. return
  1720. }
  1721. channel, err := model.GetChannelById(id, true)
  1722. if err != nil {
  1723. c.JSON(http.StatusNotFound, gin.H{
  1724. "success": false,
  1725. "message": "Channel not found",
  1726. })
  1727. return
  1728. }
  1729. if channel.Type != constant.ChannelTypeOllama {
  1730. c.JSON(http.StatusBadRequest, gin.H{
  1731. "success": false,
  1732. "message": "This operation is only supported for Ollama channels",
  1733. })
  1734. return
  1735. }
  1736. baseURL := constant.ChannelBaseURLs[channel.Type]
  1737. if channel.GetBaseURL() != "" {
  1738. baseURL = channel.GetBaseURL()
  1739. }
  1740. key := strings.Split(channel.Key, "\n")[0]
  1741. version, err := ollama.FetchOllamaVersion(baseURL, key)
  1742. if err != nil {
  1743. c.JSON(http.StatusOK, gin.H{
  1744. "success": false,
  1745. "message": fmt.Sprintf("获取Ollama版本失败: %s", err.Error()),
  1746. })
  1747. return
  1748. }
  1749. c.JSON(http.StatusOK, gin.H{
  1750. "success": true,
  1751. "data": gin.H{
  1752. "version": version,
  1753. },
  1754. })
  1755. }