You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

2201 lines
57 KiB

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