channel.go 54 KB

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