ratio_sync.go 8.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293
  1. package controller
  2. import (
  3. "encoding/json"
  4. "net/http"
  5. "one-api/model"
  6. "one-api/setting/ratio_setting"
  7. "one-api/dto"
  8. "sync"
  9. "time"
  10. "github.com/gin-gonic/gin"
  11. )
  12. type upstreamResult struct {
  13. Name string `json:"name"`
  14. Data map[string]any `json:"data,omitempty"`
  15. Err string `json:"err,omitempty"`
  16. }
  17. type TestResult struct {
  18. Name string `json:"name"`
  19. Status string `json:"status"`
  20. Error string `json:"error,omitempty"`
  21. }
  22. type DifferenceItem struct {
  23. Current interface{} `json:"current"` // 当前本地值,可能为null
  24. Upstreams map[string]interface{} `json:"upstreams"` // 上游值:具体值/"same"/null
  25. }
  26. type SyncableChannel struct {
  27. ID int `json:"id"`
  28. Name string `json:"name"`
  29. BaseURL string `json:"base_url"`
  30. Status int `json:"status"`
  31. }
  32. func FetchUpstreamRatios(c *gin.Context) {
  33. var req dto.UpstreamRequest
  34. if err := c.ShouldBindJSON(&req); err != nil {
  35. c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()})
  36. return
  37. }
  38. if req.Timeout <= 0 {
  39. req.Timeout = 10
  40. }
  41. var upstreams []dto.UpstreamDTO
  42. if len(req.ChannelIDs) > 0 {
  43. intIds := make([]int, 0, len(req.ChannelIDs))
  44. for _, id64 := range req.ChannelIDs {
  45. intIds = append(intIds, int(id64))
  46. }
  47. dbChannels, _ := model.GetChannelsByIds(intIds)
  48. for _, ch := range dbChannels {
  49. upstreams = append(upstreams, dto.UpstreamDTO{
  50. Name: ch.Name,
  51. BaseURL: ch.GetBaseURL(),
  52. Endpoint: "",
  53. })
  54. }
  55. }
  56. var wg sync.WaitGroup
  57. ch := make(chan upstreamResult, len(upstreams))
  58. for _, chn := range upstreams {
  59. wg.Add(1)
  60. go func(chItem dto.UpstreamDTO) {
  61. defer wg.Done()
  62. endpoint := chItem.Endpoint
  63. if endpoint == "" {
  64. endpoint = "/api/ratio_config"
  65. }
  66. url := chItem.BaseURL + endpoint
  67. client := http.Client{Timeout: time.Duration(req.Timeout) * time.Second}
  68. resp, err := client.Get(url)
  69. if err != nil {
  70. ch <- upstreamResult{Name: chItem.Name, Err: err.Error()}
  71. return
  72. }
  73. defer resp.Body.Close()
  74. if resp.StatusCode != http.StatusOK {
  75. ch <- upstreamResult{Name: chItem.Name, Err: resp.Status}
  76. return
  77. }
  78. var body struct {
  79. Success bool `json:"success"`
  80. Data map[string]any `json:"data"`
  81. Message string `json:"message"`
  82. }
  83. if err := json.NewDecoder(resp.Body).Decode(&body); err != nil {
  84. ch <- upstreamResult{Name: chItem.Name, Err: err.Error()}
  85. return
  86. }
  87. if !body.Success {
  88. ch <- upstreamResult{Name: chItem.Name, Err: body.Message}
  89. return
  90. }
  91. ch <- upstreamResult{Name: chItem.Name, Data: body.Data}
  92. }(chn)
  93. }
  94. wg.Wait()
  95. close(ch)
  96. localData := ratio_setting.GetExposedData()
  97. var testResults []dto.TestResult
  98. var successfulChannels []struct {
  99. name string
  100. data map[string]any
  101. }
  102. for r := range ch {
  103. if r.Err != "" {
  104. testResults = append(testResults, dto.TestResult{
  105. Name: r.Name,
  106. Status: "error",
  107. Error: r.Err,
  108. })
  109. } else {
  110. testResults = append(testResults, dto.TestResult{
  111. Name: r.Name,
  112. Status: "success",
  113. })
  114. successfulChannels = append(successfulChannels, struct {
  115. name string
  116. data map[string]any
  117. }{name: r.Name, data: r.Data})
  118. }
  119. }
  120. differences := buildDifferences(localData, successfulChannels)
  121. c.JSON(http.StatusOK, gin.H{
  122. "success": true,
  123. "data": gin.H{
  124. "differences": differences,
  125. "test_results": testResults,
  126. },
  127. })
  128. }
  129. func buildDifferences(localData map[string]any, successfulChannels []struct {
  130. name string
  131. data map[string]any
  132. }) map[string]map[string]dto.DifferenceItem {
  133. differences := make(map[string]map[string]dto.DifferenceItem)
  134. ratioTypes := []string{"model_ratio", "completion_ratio", "cache_ratio", "model_price"}
  135. allModels := make(map[string]struct{})
  136. for _, ratioType := range ratioTypes {
  137. if localRatioAny, ok := localData[ratioType]; ok {
  138. if localRatio, ok := localRatioAny.(map[string]float64); ok {
  139. for modelName := range localRatio {
  140. allModels[modelName] = struct{}{}
  141. }
  142. }
  143. }
  144. }
  145. for _, channel := range successfulChannels {
  146. for _, ratioType := range ratioTypes {
  147. if upstreamRatio, ok := channel.data[ratioType].(map[string]any); ok {
  148. for modelName := range upstreamRatio {
  149. allModels[modelName] = struct{}{}
  150. }
  151. }
  152. }
  153. }
  154. for modelName := range allModels {
  155. for _, ratioType := range ratioTypes {
  156. var localValue interface{} = nil
  157. if localRatioAny, ok := localData[ratioType]; ok {
  158. if localRatio, ok := localRatioAny.(map[string]float64); ok {
  159. if val, exists := localRatio[modelName]; exists {
  160. localValue = val
  161. }
  162. }
  163. }
  164. upstreamValues := make(map[string]interface{})
  165. hasUpstreamValue := false
  166. hasDifference := false
  167. for _, channel := range successfulChannels {
  168. var upstreamValue interface{} = nil
  169. if upstreamRatio, ok := channel.data[ratioType].(map[string]any); ok {
  170. if val, exists := upstreamRatio[modelName]; exists {
  171. upstreamValue = val
  172. hasUpstreamValue = true
  173. if localValue != nil && localValue != val {
  174. hasDifference = true
  175. } else if localValue == val {
  176. upstreamValue = "same"
  177. }
  178. }
  179. }
  180. if upstreamValue == nil && localValue == nil {
  181. upstreamValue = "same"
  182. }
  183. if localValue == nil && upstreamValue != nil && upstreamValue != "same" {
  184. hasDifference = true
  185. }
  186. upstreamValues[channel.name] = upstreamValue
  187. }
  188. shouldInclude := false
  189. if localValue != nil {
  190. if hasDifference {
  191. shouldInclude = true
  192. }
  193. } else {
  194. if hasUpstreamValue {
  195. shouldInclude = true
  196. }
  197. }
  198. if shouldInclude {
  199. if differences[modelName] == nil {
  200. differences[modelName] = make(map[string]dto.DifferenceItem)
  201. }
  202. differences[modelName][ratioType] = dto.DifferenceItem{
  203. Current: localValue,
  204. Upstreams: upstreamValues,
  205. }
  206. }
  207. }
  208. }
  209. channelHasDiff := make(map[string]bool)
  210. for _, ratioMap := range differences {
  211. for _, item := range ratioMap {
  212. for chName, val := range item.Upstreams {
  213. if val != nil && val != "same" {
  214. channelHasDiff[chName] = true
  215. }
  216. }
  217. }
  218. }
  219. for modelName, ratioMap := range differences {
  220. for ratioType, item := range ratioMap {
  221. for chName := range item.Upstreams {
  222. if !channelHasDiff[chName] {
  223. delete(item.Upstreams, chName)
  224. }
  225. }
  226. differences[modelName][ratioType] = item
  227. }
  228. }
  229. return differences
  230. }
  231. func GetSyncableChannels(c *gin.Context) {
  232. channels, err := model.GetAllChannels(0, 0, true, false)
  233. if err != nil {
  234. c.JSON(http.StatusOK, gin.H{
  235. "success": false,
  236. "message": err.Error(),
  237. })
  238. return
  239. }
  240. var syncableChannels []dto.SyncableChannel
  241. for _, channel := range channels {
  242. if channel.GetBaseURL() != "" {
  243. syncableChannels = append(syncableChannels, dto.SyncableChannel{
  244. ID: channel.Id,
  245. Name: channel.Name,
  246. BaseURL: channel.GetBaseURL(),
  247. Status: channel.Status,
  248. })
  249. }
  250. }
  251. c.JSON(http.StatusOK, gin.H{
  252. "success": true,
  253. "message": "",
  254. "data": syncableChannels,
  255. })
  256. }