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

2211 lines
57 KiB

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