channel.go 32 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290
  1. package controller
  2. import (
  3. "encoding/json"
  4. "fmt"
  5. "net/http"
  6. "one-api/common"
  7. "one-api/constant"
  8. "one-api/model"
  9. "strconv"
  10. "strings"
  11. "github.com/gin-gonic/gin"
  12. )
  13. type OpenAIModel struct {
  14. ID string `json:"id"`
  15. Object string `json:"object"`
  16. Created int64 `json:"created"`
  17. OwnedBy string `json:"owned_by"`
  18. Permission []struct {
  19. ID string `json:"id"`
  20. Object string `json:"object"`
  21. Created int64 `json:"created"`
  22. AllowCreateEngine bool `json:"allow_create_engine"`
  23. AllowSampling bool `json:"allow_sampling"`
  24. AllowLogprobs bool `json:"allow_logprobs"`
  25. AllowSearchIndices bool `json:"allow_search_indices"`
  26. AllowView bool `json:"allow_view"`
  27. AllowFineTuning bool `json:"allow_fine_tuning"`
  28. Organization string `json:"organization"`
  29. Group string `json:"group"`
  30. IsBlocking bool `json:"is_blocking"`
  31. } `json:"permission"`
  32. Root string `json:"root"`
  33. Parent string `json:"parent"`
  34. }
  35. type GoogleOpenAICompatibleModels []struct {
  36. Name string `json:"name"`
  37. Version string `json:"version"`
  38. DisplayName string `json:"displayName"`
  39. Description string `json:"description,omitempty"`
  40. InputTokenLimit int `json:"inputTokenLimit"`
  41. OutputTokenLimit int `json:"outputTokenLimit"`
  42. SupportedGenerationMethods []string `json:"supportedGenerationMethods"`
  43. Temperature float64 `json:"temperature,omitempty"`
  44. TopP float64 `json:"topP,omitempty"`
  45. TopK int `json:"topK,omitempty"`
  46. MaxTemperature int `json:"maxTemperature,omitempty"`
  47. }
  48. type OpenAIModelsResponse struct {
  49. Data []OpenAIModel `json:"data"`
  50. Success bool `json:"success"`
  51. }
  52. type GoogleOpenAICompatibleResponse struct {
  53. Models []GoogleOpenAICompatibleModels `json:"models"`
  54. NextPageToken string `json:"nextPageToken"`
  55. }
  56. func parseStatusFilter(statusParam string) int {
  57. switch strings.ToLower(statusParam) {
  58. case "enabled", "1":
  59. return common.ChannelStatusEnabled
  60. case "disabled", "0":
  61. return 0
  62. default:
  63. return -1
  64. }
  65. }
  66. func GetAllChannels(c *gin.Context) {
  67. pageInfo := common.GetPageQuery(c)
  68. channelData := make([]*model.Channel, 0)
  69. idSort, _ := strconv.ParseBool(c.Query("id_sort"))
  70. enableTagMode, _ := strconv.ParseBool(c.Query("tag_mode"))
  71. statusParam := c.Query("status")
  72. // statusFilter: -1 all, 1 enabled, 0 disabled (include auto & manual)
  73. statusFilter := parseStatusFilter(statusParam)
  74. // type filter
  75. typeStr := c.Query("type")
  76. typeFilter := -1
  77. if typeStr != "" {
  78. if t, err := strconv.Atoi(typeStr); err == nil {
  79. typeFilter = t
  80. }
  81. }
  82. var total int64
  83. if enableTagMode {
  84. tags, err := model.GetPaginatedTags(pageInfo.GetStartIdx(), pageInfo.GetPageSize())
  85. if err != nil {
  86. c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()})
  87. return
  88. }
  89. for _, tag := range tags {
  90. if tag == nil || *tag == "" {
  91. continue
  92. }
  93. tagChannels, err := model.GetChannelsByTag(*tag, idSort)
  94. if err != nil {
  95. continue
  96. }
  97. filtered := make([]*model.Channel, 0)
  98. for _, ch := range tagChannels {
  99. if statusFilter == common.ChannelStatusEnabled && ch.Status != common.ChannelStatusEnabled {
  100. continue
  101. }
  102. if statusFilter == 0 && ch.Status == common.ChannelStatusEnabled {
  103. continue
  104. }
  105. if typeFilter >= 0 && ch.Type != typeFilter {
  106. continue
  107. }
  108. filtered = append(filtered, ch)
  109. }
  110. channelData = append(channelData, filtered...)
  111. }
  112. total, _ = model.CountAllTags()
  113. } else {
  114. baseQuery := model.DB.Model(&model.Channel{})
  115. if typeFilter >= 0 {
  116. baseQuery = baseQuery.Where("type = ?", typeFilter)
  117. }
  118. if statusFilter == common.ChannelStatusEnabled {
  119. baseQuery = baseQuery.Where("status = ?", common.ChannelStatusEnabled)
  120. } else if statusFilter == 0 {
  121. baseQuery = baseQuery.Where("status != ?", common.ChannelStatusEnabled)
  122. }
  123. baseQuery.Count(&total)
  124. order := "priority desc"
  125. if idSort {
  126. order = "id desc"
  127. }
  128. err := baseQuery.Order(order).Limit(pageInfo.GetPageSize()).Offset(pageInfo.GetStartIdx()).Omit("key").Find(&channelData).Error
  129. if err != nil {
  130. c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()})
  131. return
  132. }
  133. }
  134. countQuery := model.DB.Model(&model.Channel{})
  135. if statusFilter == common.ChannelStatusEnabled {
  136. countQuery = countQuery.Where("status = ?", common.ChannelStatusEnabled)
  137. } else if statusFilter == 0 {
  138. countQuery = countQuery.Where("status != ?", common.ChannelStatusEnabled)
  139. }
  140. var results []struct {
  141. Type int64
  142. Count int64
  143. }
  144. _ = countQuery.Select("type, count(*) as count").Group("type").Find(&results).Error
  145. typeCounts := make(map[int64]int64)
  146. for _, r := range results {
  147. typeCounts[r.Type] = r.Count
  148. }
  149. common.ApiSuccess(c, gin.H{
  150. "items": channelData,
  151. "total": total,
  152. "page": pageInfo.GetPage(),
  153. "page_size": pageInfo.GetPageSize(),
  154. "type_counts": typeCounts,
  155. })
  156. return
  157. }
  158. func FetchUpstreamModels(c *gin.Context) {
  159. id, err := strconv.Atoi(c.Param("id"))
  160. if err != nil {
  161. common.ApiError(c, err)
  162. return
  163. }
  164. channel, err := model.GetChannelById(id, true)
  165. if err != nil {
  166. common.ApiError(c, err)
  167. return
  168. }
  169. baseURL := constant.ChannelBaseURLs[channel.Type]
  170. if channel.GetBaseURL() != "" {
  171. baseURL = channel.GetBaseURL()
  172. }
  173. var url string
  174. switch channel.Type {
  175. case constant.ChannelTypeGemini:
  176. // curl https://example.com/v1beta/models?key=$GEMINI_API_KEY
  177. url = fmt.Sprintf("%s/v1beta/openai/models?key=%s", baseURL, channel.Key)
  178. case constant.ChannelTypeAli:
  179. url = fmt.Sprintf("%s/compatible-mode/v1/models", baseURL)
  180. default:
  181. url = fmt.Sprintf("%s/v1/models", baseURL)
  182. }
  183. // 获取响应体 - 根据渠道类型决定是否添加 AuthHeader
  184. var body []byte
  185. if channel.Type == constant.ChannelTypeGemini {
  186. body, err = GetResponseBody("GET", url, channel, nil) // I don't know why, but Gemini requires no AuthHeader
  187. } else {
  188. body, err = GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
  189. }
  190. if err != nil {
  191. common.ApiError(c, err)
  192. return
  193. }
  194. var result OpenAIModelsResponse
  195. var parseSuccess bool
  196. // 适配特殊格式
  197. switch channel.Type {
  198. case constant.ChannelTypeGemini:
  199. var googleResult GoogleOpenAICompatibleResponse
  200. if err = json.Unmarshal(body, &googleResult); err == nil {
  201. // 转换Google格式到OpenAI格式
  202. for _, model := range googleResult.Models {
  203. for _, gModel := range model {
  204. result.Data = append(result.Data, OpenAIModel{
  205. ID: gModel.Name,
  206. })
  207. }
  208. }
  209. parseSuccess = true
  210. }
  211. }
  212. // 如果解析失败,尝试OpenAI格式
  213. if !parseSuccess {
  214. if err = json.Unmarshal(body, &result); err != nil {
  215. c.JSON(http.StatusOK, gin.H{
  216. "success": false,
  217. "message": fmt.Sprintf("解析响应失败: %s", err.Error()),
  218. })
  219. return
  220. }
  221. }
  222. var ids []string
  223. for _, model := range result.Data {
  224. id := model.ID
  225. if channel.Type == constant.ChannelTypeGemini {
  226. id = strings.TrimPrefix(id, "models/")
  227. }
  228. ids = append(ids, id)
  229. }
  230. c.JSON(http.StatusOK, gin.H{
  231. "success": true,
  232. "message": "",
  233. "data": ids,
  234. })
  235. }
  236. func FixChannelsAbilities(c *gin.Context) {
  237. success, fails, err := model.FixAbility()
  238. if err != nil {
  239. common.ApiError(c, err)
  240. return
  241. }
  242. c.JSON(http.StatusOK, gin.H{
  243. "success": true,
  244. "message": "",
  245. "data": gin.H{
  246. "success": success,
  247. "fails": fails,
  248. },
  249. })
  250. }
  251. func SearchChannels(c *gin.Context) {
  252. keyword := c.Query("keyword")
  253. group := c.Query("group")
  254. modelKeyword := c.Query("model")
  255. statusParam := c.Query("status")
  256. statusFilter := parseStatusFilter(statusParam)
  257. idSort, _ := strconv.ParseBool(c.Query("id_sort"))
  258. enableTagMode, _ := strconv.ParseBool(c.Query("tag_mode"))
  259. channelData := make([]*model.Channel, 0)
  260. if enableTagMode {
  261. tags, err := model.SearchTags(keyword, group, modelKeyword, idSort)
  262. if err != nil {
  263. c.JSON(http.StatusOK, gin.H{
  264. "success": false,
  265. "message": err.Error(),
  266. })
  267. return
  268. }
  269. for _, tag := range tags {
  270. if tag != nil && *tag != "" {
  271. tagChannel, err := model.GetChannelsByTag(*tag, idSort)
  272. if err == nil {
  273. channelData = append(channelData, tagChannel...)
  274. }
  275. }
  276. }
  277. } else {
  278. channels, err := model.SearchChannels(keyword, group, modelKeyword, idSort)
  279. if err != nil {
  280. c.JSON(http.StatusOK, gin.H{
  281. "success": false,
  282. "message": err.Error(),
  283. })
  284. return
  285. }
  286. channelData = channels
  287. }
  288. if statusFilter == common.ChannelStatusEnabled || statusFilter == 0 {
  289. filtered := make([]*model.Channel, 0, len(channelData))
  290. for _, ch := range channelData {
  291. if statusFilter == common.ChannelStatusEnabled && ch.Status != common.ChannelStatusEnabled {
  292. continue
  293. }
  294. if statusFilter == 0 && ch.Status == common.ChannelStatusEnabled {
  295. continue
  296. }
  297. filtered = append(filtered, ch)
  298. }
  299. channelData = filtered
  300. }
  301. // calculate type counts for search results
  302. typeCounts := make(map[int64]int64)
  303. for _, channel := range channelData {
  304. typeCounts[int64(channel.Type)]++
  305. }
  306. typeParam := c.Query("type")
  307. typeFilter := -1
  308. if typeParam != "" {
  309. if tp, err := strconv.Atoi(typeParam); err == nil {
  310. typeFilter = tp
  311. }
  312. }
  313. if typeFilter >= 0 {
  314. filtered := make([]*model.Channel, 0, len(channelData))
  315. for _, ch := range channelData {
  316. if ch.Type == typeFilter {
  317. filtered = append(filtered, ch)
  318. }
  319. }
  320. channelData = filtered
  321. }
  322. page, _ := strconv.Atoi(c.DefaultQuery("p", "1"))
  323. pageSize, _ := strconv.Atoi(c.DefaultQuery("page_size", "20"))
  324. if page < 1 {
  325. page = 1
  326. }
  327. if pageSize <= 0 {
  328. pageSize = 20
  329. }
  330. total := len(channelData)
  331. startIdx := (page - 1) * pageSize
  332. if startIdx > total {
  333. startIdx = total
  334. }
  335. endIdx := startIdx + pageSize
  336. if endIdx > total {
  337. endIdx = total
  338. }
  339. pagedData := channelData[startIdx:endIdx]
  340. c.JSON(http.StatusOK, gin.H{
  341. "success": true,
  342. "message": "",
  343. "data": gin.H{
  344. "items": pagedData,
  345. "total": total,
  346. "type_counts": typeCounts,
  347. },
  348. })
  349. return
  350. }
  351. func GetChannel(c *gin.Context) {
  352. id, err := strconv.Atoi(c.Param("id"))
  353. if err != nil {
  354. common.ApiError(c, err)
  355. return
  356. }
  357. channel, err := model.GetChannelById(id, false)
  358. if err != nil {
  359. common.ApiError(c, err)
  360. return
  361. }
  362. c.JSON(http.StatusOK, gin.H{
  363. "success": true,
  364. "message": "",
  365. "data": channel,
  366. })
  367. return
  368. }
  369. // validateChannel 通用的渠道校验函数
  370. func validateChannel(channel *model.Channel, isAdd bool) error {
  371. // 校验 channel settings
  372. if err := channel.ValidateSettings(); err != nil {
  373. return fmt.Errorf("渠道额外设置[channel setting] 格式错误:%s", err.Error())
  374. }
  375. // 如果是添加操作,检查 channel 和 key 是否为空
  376. if isAdd {
  377. if channel == nil || channel.Key == "" {
  378. return fmt.Errorf("channel cannot be empty")
  379. }
  380. // 检查模型名称长度是否超过 255
  381. for _, m := range channel.GetModels() {
  382. if len(m) > 255 {
  383. return fmt.Errorf("模型名称过长: %s", m)
  384. }
  385. }
  386. }
  387. // VertexAI 特殊校验
  388. if channel.Type == constant.ChannelTypeVertexAi {
  389. if channel.Other == "" {
  390. return fmt.Errorf("部署地区不能为空")
  391. }
  392. regionMap, err := common.StrToMap(channel.Other)
  393. if err != nil {
  394. return fmt.Errorf("部署地区必须是标准的Json格式,例如{\"default\": \"us-central1\", \"region2\": \"us-east1\"}")
  395. }
  396. if regionMap["default"] == nil {
  397. return fmt.Errorf("部署地区必须包含default字段")
  398. }
  399. }
  400. return nil
  401. }
  402. type AddChannelRequest struct {
  403. Mode string `json:"mode"`
  404. MultiKeyMode constant.MultiKeyMode `json:"multi_key_mode"`
  405. Channel *model.Channel `json:"channel"`
  406. }
  407. func getVertexArrayKeys(keys string) ([]string, error) {
  408. if keys == "" {
  409. return nil, nil
  410. }
  411. var keyArray []interface{}
  412. err := common.Unmarshal([]byte(keys), &keyArray)
  413. if err != nil {
  414. return nil, fmt.Errorf("批量添加 Vertex AI 必须使用标准的JsonArray格式,例如[{key1}, {key2}...],请检查输入: %w", err)
  415. }
  416. cleanKeys := make([]string, 0, len(keyArray))
  417. for _, key := range keyArray {
  418. var keyStr string
  419. switch v := key.(type) {
  420. case string:
  421. keyStr = strings.TrimSpace(v)
  422. default:
  423. bytes, err := json.Marshal(v)
  424. if err != nil {
  425. return nil, fmt.Errorf("Vertex AI key JSON 编码失败: %w", err)
  426. }
  427. keyStr = string(bytes)
  428. }
  429. if keyStr != "" {
  430. cleanKeys = append(cleanKeys, keyStr)
  431. }
  432. }
  433. if len(cleanKeys) == 0 {
  434. return nil, fmt.Errorf("批量添加 Vertex AI 的 keys 不能为空")
  435. }
  436. return cleanKeys, nil
  437. }
  438. func AddChannel(c *gin.Context) {
  439. addChannelRequest := AddChannelRequest{}
  440. err := c.ShouldBindJSON(&addChannelRequest)
  441. if err != nil {
  442. common.ApiError(c, err)
  443. return
  444. }
  445. // 使用统一的校验函数
  446. if err := validateChannel(addChannelRequest.Channel, true); err != nil {
  447. c.JSON(http.StatusOK, gin.H{
  448. "success": false,
  449. "message": err.Error(),
  450. })
  451. return
  452. }
  453. addChannelRequest.Channel.CreatedTime = common.GetTimestamp()
  454. keys := make([]string, 0)
  455. switch addChannelRequest.Mode {
  456. case "multi_to_single":
  457. addChannelRequest.Channel.ChannelInfo.IsMultiKey = true
  458. addChannelRequest.Channel.ChannelInfo.MultiKeyMode = addChannelRequest.MultiKeyMode
  459. if addChannelRequest.Channel.Type == constant.ChannelTypeVertexAi {
  460. array, err := getVertexArrayKeys(addChannelRequest.Channel.Key)
  461. if err != nil {
  462. c.JSON(http.StatusOK, gin.H{
  463. "success": false,
  464. "message": err.Error(),
  465. })
  466. return
  467. }
  468. addChannelRequest.Channel.ChannelInfo.MultiKeySize = len(array)
  469. addChannelRequest.Channel.Key = strings.Join(array, "\n")
  470. } else {
  471. cleanKeys := make([]string, 0)
  472. for _, key := range strings.Split(addChannelRequest.Channel.Key, "\n") {
  473. if key == "" {
  474. continue
  475. }
  476. key = strings.TrimSpace(key)
  477. cleanKeys = append(cleanKeys, key)
  478. }
  479. addChannelRequest.Channel.ChannelInfo.MultiKeySize = len(cleanKeys)
  480. addChannelRequest.Channel.Key = strings.Join(cleanKeys, "\n")
  481. }
  482. keys = []string{addChannelRequest.Channel.Key}
  483. case "batch":
  484. if addChannelRequest.Channel.Type == constant.ChannelTypeVertexAi {
  485. // multi json
  486. keys, err = getVertexArrayKeys(addChannelRequest.Channel.Key)
  487. if err != nil {
  488. c.JSON(http.StatusOK, gin.H{
  489. "success": false,
  490. "message": err.Error(),
  491. })
  492. return
  493. }
  494. } else {
  495. keys = strings.Split(addChannelRequest.Channel.Key, "\n")
  496. }
  497. case "single":
  498. keys = []string{addChannelRequest.Channel.Key}
  499. default:
  500. c.JSON(http.StatusOK, gin.H{
  501. "success": false,
  502. "message": "不支持的添加模式",
  503. })
  504. return
  505. }
  506. channels := make([]model.Channel, 0, len(keys))
  507. for _, key := range keys {
  508. if key == "" {
  509. continue
  510. }
  511. localChannel := addChannelRequest.Channel
  512. localChannel.Key = key
  513. channels = append(channels, *localChannel)
  514. }
  515. err = model.BatchInsertChannels(channels)
  516. if err != nil {
  517. common.ApiError(c, err)
  518. return
  519. }
  520. c.JSON(http.StatusOK, gin.H{
  521. "success": true,
  522. "message": "",
  523. })
  524. return
  525. }
  526. func DeleteChannel(c *gin.Context) {
  527. id, _ := strconv.Atoi(c.Param("id"))
  528. channel := model.Channel{Id: id}
  529. err := channel.Delete()
  530. if err != nil {
  531. common.ApiError(c, err)
  532. return
  533. }
  534. model.InitChannelCache()
  535. c.JSON(http.StatusOK, gin.H{
  536. "success": true,
  537. "message": "",
  538. })
  539. return
  540. }
  541. func DeleteDisabledChannel(c *gin.Context) {
  542. rows, err := model.DeleteDisabledChannel()
  543. if err != nil {
  544. common.ApiError(c, err)
  545. return
  546. }
  547. model.InitChannelCache()
  548. c.JSON(http.StatusOK, gin.H{
  549. "success": true,
  550. "message": "",
  551. "data": rows,
  552. })
  553. return
  554. }
  555. type ChannelTag struct {
  556. Tag string `json:"tag"`
  557. NewTag *string `json:"new_tag"`
  558. Priority *int64 `json:"priority"`
  559. Weight *uint `json:"weight"`
  560. ModelMapping *string `json:"model_mapping"`
  561. Models *string `json:"models"`
  562. Groups *string `json:"groups"`
  563. }
  564. func DisableTagChannels(c *gin.Context) {
  565. channelTag := ChannelTag{}
  566. err := c.ShouldBindJSON(&channelTag)
  567. if err != nil || channelTag.Tag == "" {
  568. c.JSON(http.StatusOK, gin.H{
  569. "success": false,
  570. "message": "参数错误",
  571. })
  572. return
  573. }
  574. err = model.DisableChannelByTag(channelTag.Tag)
  575. if err != nil {
  576. common.ApiError(c, err)
  577. return
  578. }
  579. model.InitChannelCache()
  580. c.JSON(http.StatusOK, gin.H{
  581. "success": true,
  582. "message": "",
  583. })
  584. return
  585. }
  586. func EnableTagChannels(c *gin.Context) {
  587. channelTag := ChannelTag{}
  588. err := c.ShouldBindJSON(&channelTag)
  589. if err != nil || channelTag.Tag == "" {
  590. c.JSON(http.StatusOK, gin.H{
  591. "success": false,
  592. "message": "参数错误",
  593. })
  594. return
  595. }
  596. err = model.EnableChannelByTag(channelTag.Tag)
  597. if err != nil {
  598. common.ApiError(c, err)
  599. return
  600. }
  601. model.InitChannelCache()
  602. c.JSON(http.StatusOK, gin.H{
  603. "success": true,
  604. "message": "",
  605. })
  606. return
  607. }
  608. func EditTagChannels(c *gin.Context) {
  609. channelTag := ChannelTag{}
  610. err := c.ShouldBindJSON(&channelTag)
  611. if err != nil {
  612. c.JSON(http.StatusOK, gin.H{
  613. "success": false,
  614. "message": "参数错误",
  615. })
  616. return
  617. }
  618. if channelTag.Tag == "" {
  619. c.JSON(http.StatusOK, gin.H{
  620. "success": false,
  621. "message": "tag不能为空",
  622. })
  623. return
  624. }
  625. err = model.EditChannelByTag(channelTag.Tag, channelTag.NewTag, channelTag.ModelMapping, channelTag.Models, channelTag.Groups, channelTag.Priority, channelTag.Weight)
  626. if err != nil {
  627. common.ApiError(c, err)
  628. return
  629. }
  630. model.InitChannelCache()
  631. c.JSON(http.StatusOK, gin.H{
  632. "success": true,
  633. "message": "",
  634. })
  635. return
  636. }
  637. type ChannelBatch struct {
  638. Ids []int `json:"ids"`
  639. Tag *string `json:"tag"`
  640. }
  641. func DeleteChannelBatch(c *gin.Context) {
  642. channelBatch := ChannelBatch{}
  643. err := c.ShouldBindJSON(&channelBatch)
  644. if err != nil || len(channelBatch.Ids) == 0 {
  645. c.JSON(http.StatusOK, gin.H{
  646. "success": false,
  647. "message": "参数错误",
  648. })
  649. return
  650. }
  651. err = model.BatchDeleteChannels(channelBatch.Ids)
  652. if err != nil {
  653. common.ApiError(c, err)
  654. return
  655. }
  656. model.InitChannelCache()
  657. c.JSON(http.StatusOK, gin.H{
  658. "success": true,
  659. "message": "",
  660. "data": len(channelBatch.Ids),
  661. })
  662. return
  663. }
  664. type PatchChannel struct {
  665. model.Channel
  666. MultiKeyMode *string `json:"multi_key_mode"`
  667. KeyMode *string `json:"key_mode"` // 多key模式下密钥覆盖或者追加
  668. }
  669. func UpdateChannel(c *gin.Context) {
  670. channel := PatchChannel{}
  671. err := c.ShouldBindJSON(&channel)
  672. if err != nil {
  673. common.ApiError(c, err)
  674. return
  675. }
  676. // 使用统一的校验函数
  677. if err := validateChannel(&channel.Channel, false); err != nil {
  678. c.JSON(http.StatusOK, gin.H{
  679. "success": false,
  680. "message": err.Error(),
  681. })
  682. return
  683. }
  684. // Preserve existing ChannelInfo to ensure multi-key channels keep correct state even if the client does not send ChannelInfo in the request.
  685. originChannel, err := model.GetChannelById(channel.Id, true)
  686. if err != nil {
  687. c.JSON(http.StatusOK, gin.H{
  688. "success": false,
  689. "message": err.Error(),
  690. })
  691. return
  692. }
  693. // Always copy the original ChannelInfo so that fields like IsMultiKey and MultiKeySize are retained.
  694. channel.ChannelInfo = originChannel.ChannelInfo
  695. // If the request explicitly specifies a new MultiKeyMode, apply it on top of the original info.
  696. if channel.MultiKeyMode != nil && *channel.MultiKeyMode != "" {
  697. channel.ChannelInfo.MultiKeyMode = constant.MultiKeyMode(*channel.MultiKeyMode)
  698. }
  699. // 处理多key模式下的密钥追加/覆盖逻辑
  700. if channel.KeyMode != nil && channel.ChannelInfo.IsMultiKey {
  701. switch *channel.KeyMode {
  702. case "append":
  703. // 追加模式:将新密钥添加到现有密钥列表
  704. if originChannel.Key != "" {
  705. var newKeys []string
  706. var existingKeys []string
  707. // 解析现有密钥
  708. if strings.HasPrefix(strings.TrimSpace(originChannel.Key), "[") {
  709. // JSON数组格式
  710. var arr []json.RawMessage
  711. if err := json.Unmarshal([]byte(strings.TrimSpace(originChannel.Key)), &arr); err == nil {
  712. existingKeys = make([]string, len(arr))
  713. for i, v := range arr {
  714. existingKeys[i] = string(v)
  715. }
  716. }
  717. } else {
  718. // 换行分隔格式
  719. existingKeys = strings.Split(strings.Trim(originChannel.Key, "\n"), "\n")
  720. }
  721. // 处理 Vertex AI 的特殊情况
  722. if channel.Type == constant.ChannelTypeVertexAi {
  723. // 尝试解析新密钥为JSON数组
  724. if strings.HasPrefix(strings.TrimSpace(channel.Key), "[") {
  725. array, err := getVertexArrayKeys(channel.Key)
  726. if err != nil {
  727. c.JSON(http.StatusOK, gin.H{
  728. "success": false,
  729. "message": "追加密钥解析失败: " + err.Error(),
  730. })
  731. return
  732. }
  733. newKeys = array
  734. } else {
  735. // 单个JSON密钥
  736. newKeys = []string{channel.Key}
  737. }
  738. // 合并密钥
  739. allKeys := append(existingKeys, newKeys...)
  740. channel.Key = strings.Join(allKeys, "\n")
  741. } else {
  742. // 普通渠道的处理
  743. inputKeys := strings.Split(channel.Key, "\n")
  744. for _, key := range inputKeys {
  745. key = strings.TrimSpace(key)
  746. if key != "" {
  747. newKeys = append(newKeys, key)
  748. }
  749. }
  750. // 合并密钥
  751. allKeys := append(existingKeys, newKeys...)
  752. channel.Key = strings.Join(allKeys, "\n")
  753. }
  754. }
  755. case "replace":
  756. // 覆盖模式:直接使用新密钥(默认行为,不需要特殊处理)
  757. }
  758. }
  759. err = channel.Update()
  760. if err != nil {
  761. common.ApiError(c, err)
  762. return
  763. }
  764. model.InitChannelCache()
  765. channel.Key = ""
  766. c.JSON(http.StatusOK, gin.H{
  767. "success": true,
  768. "message": "",
  769. "data": channel,
  770. })
  771. return
  772. }
  773. func FetchModels(c *gin.Context) {
  774. var req struct {
  775. BaseURL string `json:"base_url"`
  776. Type int `json:"type"`
  777. Key string `json:"key"`
  778. }
  779. if err := c.ShouldBindJSON(&req); err != nil {
  780. c.JSON(http.StatusBadRequest, gin.H{
  781. "success": false,
  782. "message": "Invalid request",
  783. })
  784. return
  785. }
  786. baseURL := req.BaseURL
  787. if baseURL == "" {
  788. baseURL = constant.ChannelBaseURLs[req.Type]
  789. }
  790. client := &http.Client{}
  791. url := fmt.Sprintf("%s/v1/models", baseURL)
  792. request, err := http.NewRequest("GET", url, nil)
  793. if err != nil {
  794. c.JSON(http.StatusInternalServerError, gin.H{
  795. "success": false,
  796. "message": err.Error(),
  797. })
  798. return
  799. }
  800. // remove line breaks and extra spaces.
  801. key := strings.TrimSpace(req.Key)
  802. // If the key contains a line break, only take the first part.
  803. key = strings.Split(key, "\n")[0]
  804. request.Header.Set("Authorization", "Bearer "+key)
  805. response, err := client.Do(request)
  806. if err != nil {
  807. c.JSON(http.StatusInternalServerError, gin.H{
  808. "success": false,
  809. "message": err.Error(),
  810. })
  811. return
  812. }
  813. //check status code
  814. if response.StatusCode != http.StatusOK {
  815. c.JSON(http.StatusInternalServerError, gin.H{
  816. "success": false,
  817. "message": "Failed to fetch models",
  818. })
  819. return
  820. }
  821. defer response.Body.Close()
  822. var result struct {
  823. Data []struct {
  824. ID string `json:"id"`
  825. } `json:"data"`
  826. }
  827. if err := json.NewDecoder(response.Body).Decode(&result); err != nil {
  828. c.JSON(http.StatusInternalServerError, gin.H{
  829. "success": false,
  830. "message": err.Error(),
  831. })
  832. return
  833. }
  834. var models []string
  835. for _, model := range result.Data {
  836. models = append(models, model.ID)
  837. }
  838. c.JSON(http.StatusOK, gin.H{
  839. "success": true,
  840. "data": models,
  841. })
  842. }
  843. func BatchSetChannelTag(c *gin.Context) {
  844. channelBatch := ChannelBatch{}
  845. err := c.ShouldBindJSON(&channelBatch)
  846. if err != nil || len(channelBatch.Ids) == 0 {
  847. c.JSON(http.StatusOK, gin.H{
  848. "success": false,
  849. "message": "参数错误",
  850. })
  851. return
  852. }
  853. err = model.BatchSetChannelTag(channelBatch.Ids, channelBatch.Tag)
  854. if err != nil {
  855. common.ApiError(c, err)
  856. return
  857. }
  858. model.InitChannelCache()
  859. c.JSON(http.StatusOK, gin.H{
  860. "success": true,
  861. "message": "",
  862. "data": len(channelBatch.Ids),
  863. })
  864. return
  865. }
  866. func GetTagModels(c *gin.Context) {
  867. tag := c.Query("tag")
  868. if tag == "" {
  869. c.JSON(http.StatusBadRequest, gin.H{
  870. "success": false,
  871. "message": "tag不能为空",
  872. })
  873. return
  874. }
  875. channels, err := model.GetChannelsByTag(tag, false) // Assuming false for idSort is fine here
  876. if err != nil {
  877. c.JSON(http.StatusInternalServerError, gin.H{
  878. "success": false,
  879. "message": err.Error(),
  880. })
  881. return
  882. }
  883. var longestModels string
  884. maxLength := 0
  885. // Find the longest models string among all channels with the given tag
  886. for _, channel := range channels {
  887. if channel.Models != "" {
  888. currentModels := strings.Split(channel.Models, ",")
  889. if len(currentModels) > maxLength {
  890. maxLength = len(currentModels)
  891. longestModels = channel.Models
  892. }
  893. }
  894. }
  895. c.JSON(http.StatusOK, gin.H{
  896. "success": true,
  897. "message": "",
  898. "data": longestModels,
  899. })
  900. return
  901. }
  902. // CopyChannel handles cloning an existing channel with its key.
  903. // POST /api/channel/copy/:id
  904. // Optional query params:
  905. //
  906. // suffix - string appended to the original name (default "_复制")
  907. // reset_balance - bool, when true will reset balance & used_quota to 0 (default true)
  908. func CopyChannel(c *gin.Context) {
  909. id, err := strconv.Atoi(c.Param("id"))
  910. if err != nil {
  911. c.JSON(http.StatusOK, gin.H{"success": false, "message": "invalid id"})
  912. return
  913. }
  914. suffix := c.DefaultQuery("suffix", "_复制")
  915. resetBalance := true
  916. if rbStr := c.DefaultQuery("reset_balance", "true"); rbStr != "" {
  917. if v, err := strconv.ParseBool(rbStr); err == nil {
  918. resetBalance = v
  919. }
  920. }
  921. // fetch original channel with key
  922. origin, err := model.GetChannelById(id, true)
  923. if err != nil {
  924. c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()})
  925. return
  926. }
  927. // clone channel
  928. clone := *origin // shallow copy is sufficient as we will overwrite primitives
  929. clone.Id = 0 // let DB auto-generate
  930. clone.CreatedTime = common.GetTimestamp()
  931. clone.Name = origin.Name + suffix
  932. clone.TestTime = 0
  933. clone.ResponseTime = 0
  934. if resetBalance {
  935. clone.Balance = 0
  936. clone.UsedQuota = 0
  937. }
  938. // insert
  939. if err := model.BatchInsertChannels([]model.Channel{clone}); err != nil {
  940. c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()})
  941. return
  942. }
  943. model.InitChannelCache()
  944. // success
  945. c.JSON(http.StatusOK, gin.H{"success": true, "message": "", "data": gin.H{"id": clone.Id}})
  946. }
  947. // MultiKeyManageRequest represents the request for multi-key management operations
  948. type MultiKeyManageRequest struct {
  949. ChannelId int `json:"channel_id"`
  950. Action string `json:"action"` // "disable_key", "enable_key", "delete_disabled_keys", "get_key_status"
  951. KeyIndex *int `json:"key_index,omitempty"` // for disable_key and enable_key actions
  952. }
  953. // MultiKeyStatusResponse represents the response for key status query
  954. type MultiKeyStatusResponse struct {
  955. Keys []KeyStatus `json:"keys"`
  956. }
  957. type KeyStatus struct {
  958. Index int `json:"index"`
  959. Status int `json:"status"` // 1: enabled, 2: disabled
  960. DisabledTime int64 `json:"disabled_time,omitempty"`
  961. Reason string `json:"reason,omitempty"`
  962. KeyPreview string `json:"key_preview"` // first 10 chars of key for identification
  963. }
  964. // ManageMultiKeys handles multi-key management operations
  965. func ManageMultiKeys(c *gin.Context) {
  966. request := MultiKeyManageRequest{}
  967. err := c.ShouldBindJSON(&request)
  968. if err != nil {
  969. common.ApiError(c, err)
  970. return
  971. }
  972. channel, err := model.GetChannelById(request.ChannelId, true)
  973. if err != nil {
  974. c.JSON(http.StatusOK, gin.H{
  975. "success": false,
  976. "message": "渠道不存在",
  977. })
  978. return
  979. }
  980. if !channel.ChannelInfo.IsMultiKey {
  981. c.JSON(http.StatusOK, gin.H{
  982. "success": false,
  983. "message": "该渠道不是多密钥模式",
  984. })
  985. return
  986. }
  987. switch request.Action {
  988. case "get_key_status":
  989. keys := channel.GetKeys()
  990. var keyStatusList []KeyStatus
  991. for i, key := range keys {
  992. status := 1 // default enabled
  993. var disabledTime int64
  994. var reason string
  995. if channel.ChannelInfo.MultiKeyStatusList != nil {
  996. if s, exists := channel.ChannelInfo.MultiKeyStatusList[i]; exists {
  997. status = s
  998. }
  999. }
  1000. if status != 1 {
  1001. if channel.ChannelInfo.MultiKeyDisabledTime != nil {
  1002. disabledTime = channel.ChannelInfo.MultiKeyDisabledTime[i]
  1003. }
  1004. if channel.ChannelInfo.MultiKeyDisabledReason != nil {
  1005. reason = channel.ChannelInfo.MultiKeyDisabledReason[i]
  1006. }
  1007. }
  1008. // Create key preview (first 10 chars)
  1009. keyPreview := key
  1010. if len(key) > 10 {
  1011. keyPreview = key[:10] + "..."
  1012. }
  1013. keyStatusList = append(keyStatusList, KeyStatus{
  1014. Index: i,
  1015. Status: status,
  1016. DisabledTime: disabledTime,
  1017. Reason: reason,
  1018. KeyPreview: keyPreview,
  1019. })
  1020. }
  1021. c.JSON(http.StatusOK, gin.H{
  1022. "success": true,
  1023. "message": "",
  1024. "data": MultiKeyStatusResponse{Keys: keyStatusList},
  1025. })
  1026. return
  1027. case "disable_key":
  1028. if request.KeyIndex == nil {
  1029. c.JSON(http.StatusOK, gin.H{
  1030. "success": false,
  1031. "message": "未指定要禁用的密钥索引",
  1032. })
  1033. return
  1034. }
  1035. keyIndex := *request.KeyIndex
  1036. if keyIndex < 0 || keyIndex >= channel.ChannelInfo.MultiKeySize {
  1037. c.JSON(http.StatusOK, gin.H{
  1038. "success": false,
  1039. "message": "密钥索引超出范围",
  1040. })
  1041. return
  1042. }
  1043. if channel.ChannelInfo.MultiKeyStatusList == nil {
  1044. channel.ChannelInfo.MultiKeyStatusList = make(map[int]int)
  1045. }
  1046. if channel.ChannelInfo.MultiKeyDisabledTime == nil {
  1047. channel.ChannelInfo.MultiKeyDisabledTime = make(map[int]int64)
  1048. }
  1049. if channel.ChannelInfo.MultiKeyDisabledReason == nil {
  1050. channel.ChannelInfo.MultiKeyDisabledReason = make(map[int]string)
  1051. }
  1052. channel.ChannelInfo.MultiKeyStatusList[keyIndex] = 2 // disabled
  1053. channel.ChannelInfo.MultiKeyDisabledTime[keyIndex] = common.GetTimestamp()
  1054. channel.ChannelInfo.MultiKeyDisabledReason[keyIndex] = "手动禁用"
  1055. err = channel.Update()
  1056. if err != nil {
  1057. common.ApiError(c, err)
  1058. return
  1059. }
  1060. model.InitChannelCache()
  1061. c.JSON(http.StatusOK, gin.H{
  1062. "success": true,
  1063. "message": "密钥已禁用",
  1064. })
  1065. return
  1066. case "enable_key":
  1067. if request.KeyIndex == nil {
  1068. c.JSON(http.StatusOK, gin.H{
  1069. "success": false,
  1070. "message": "未指定要启用的密钥索引",
  1071. })
  1072. return
  1073. }
  1074. keyIndex := *request.KeyIndex
  1075. if keyIndex < 0 || keyIndex >= channel.ChannelInfo.MultiKeySize {
  1076. c.JSON(http.StatusOK, gin.H{
  1077. "success": false,
  1078. "message": "密钥索引超出范围",
  1079. })
  1080. return
  1081. }
  1082. // 从状态列表中删除该密钥的记录,使其回到默认启用状态
  1083. if channel.ChannelInfo.MultiKeyStatusList != nil {
  1084. delete(channel.ChannelInfo.MultiKeyStatusList, keyIndex)
  1085. }
  1086. if channel.ChannelInfo.MultiKeyDisabledTime != nil {
  1087. delete(channel.ChannelInfo.MultiKeyDisabledTime, keyIndex)
  1088. }
  1089. if channel.ChannelInfo.MultiKeyDisabledReason != nil {
  1090. delete(channel.ChannelInfo.MultiKeyDisabledReason, keyIndex)
  1091. }
  1092. err = channel.Update()
  1093. if err != nil {
  1094. common.ApiError(c, err)
  1095. return
  1096. }
  1097. model.InitChannelCache()
  1098. c.JSON(http.StatusOK, gin.H{
  1099. "success": true,
  1100. "message": "密钥已启用",
  1101. })
  1102. return
  1103. case "delete_disabled_keys":
  1104. keys := channel.GetKeys()
  1105. var remainingKeys []string
  1106. var deletedCount int
  1107. var newStatusList = make(map[int]int)
  1108. var newDisabledTime = make(map[int]int64)
  1109. var newDisabledReason = make(map[int]string)
  1110. newIndex := 0
  1111. for i, key := range keys {
  1112. status := 1 // default enabled
  1113. if channel.ChannelInfo.MultiKeyStatusList != nil {
  1114. if s, exists := channel.ChannelInfo.MultiKeyStatusList[i]; exists {
  1115. status = s
  1116. }
  1117. }
  1118. // 只删除自动禁用(status == 3)的密钥,保留启用(status == 1)和手动禁用(status == 2)的密钥
  1119. if status == 3 {
  1120. deletedCount++
  1121. } else {
  1122. remainingKeys = append(remainingKeys, key)
  1123. // 保留非自动禁用密钥的状态信息,重新索引
  1124. if status != 1 {
  1125. newStatusList[newIndex] = status
  1126. if channel.ChannelInfo.MultiKeyDisabledTime != nil {
  1127. if t, exists := channel.ChannelInfo.MultiKeyDisabledTime[i]; exists {
  1128. newDisabledTime[newIndex] = t
  1129. }
  1130. }
  1131. if channel.ChannelInfo.MultiKeyDisabledReason != nil {
  1132. if r, exists := channel.ChannelInfo.MultiKeyDisabledReason[i]; exists {
  1133. newDisabledReason[newIndex] = r
  1134. }
  1135. }
  1136. }
  1137. newIndex++
  1138. }
  1139. }
  1140. if deletedCount == 0 {
  1141. c.JSON(http.StatusOK, gin.H{
  1142. "success": false,
  1143. "message": "没有需要删除的自动禁用密钥",
  1144. })
  1145. return
  1146. }
  1147. // Update channel with remaining keys
  1148. channel.Key = strings.Join(remainingKeys, "\n")
  1149. channel.ChannelInfo.MultiKeySize = len(remainingKeys)
  1150. channel.ChannelInfo.MultiKeyStatusList = newStatusList
  1151. channel.ChannelInfo.MultiKeyDisabledTime = newDisabledTime
  1152. channel.ChannelInfo.MultiKeyDisabledReason = newDisabledReason
  1153. err = channel.Update()
  1154. if err != nil {
  1155. common.ApiError(c, err)
  1156. return
  1157. }
  1158. model.InitChannelCache()
  1159. c.JSON(http.StatusOK, gin.H{
  1160. "success": true,
  1161. "message": fmt.Sprintf("已删除 %d 个自动禁用的密钥", deletedCount),
  1162. "data": deletedCount,
  1163. })
  1164. return
  1165. default:
  1166. c.JSON(http.StatusOK, gin.H{
  1167. "success": false,
  1168. "message": "不支持的操作",
  1169. })
  1170. return
  1171. }
  1172. }