Non puoi selezionare più di 25 argomenti Gli argomenti devono iniziare con una lettera o un numero, possono includere trattini ('-') e possono essere lunghi fino a 35 caratteri.
 
 
 

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