您最多选择25个主题 主题必须以字母或数字开头,可以包含连字符 (-),并且长度不得超过35个字符
 
 
 

2226 行
58 KiB

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