channel.go 50 KB

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