channel.go 34 KB

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