channel-test.go 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760
  1. package controller
  2. import (
  3. "bytes"
  4. "encoding/json"
  5. "errors"
  6. "fmt"
  7. "io"
  8. "math"
  9. "net/http"
  10. "net/http/httptest"
  11. "net/url"
  12. "strconv"
  13. "strings"
  14. "sync"
  15. "time"
  16. "github.com/QuantumNous/new-api/common"
  17. "github.com/QuantumNous/new-api/constant"
  18. "github.com/QuantumNous/new-api/dto"
  19. "github.com/QuantumNous/new-api/middleware"
  20. "github.com/QuantumNous/new-api/model"
  21. "github.com/QuantumNous/new-api/relay"
  22. relaycommon "github.com/QuantumNous/new-api/relay/common"
  23. relayconstant "github.com/QuantumNous/new-api/relay/constant"
  24. "github.com/QuantumNous/new-api/relay/helper"
  25. "github.com/QuantumNous/new-api/service"
  26. "github.com/QuantumNous/new-api/setting/operation_setting"
  27. "github.com/QuantumNous/new-api/setting/ratio_setting"
  28. "github.com/QuantumNous/new-api/types"
  29. "github.com/bytedance/gopkg/util/gopool"
  30. "github.com/samber/lo"
  31. "github.com/gin-gonic/gin"
  32. )
  33. type testResult struct {
  34. context *gin.Context
  35. localErr error
  36. newAPIError *types.NewAPIError
  37. }
  38. func testChannel(channel *model.Channel, testModel string, endpointType string) testResult {
  39. tik := time.Now()
  40. var unsupportedTestChannelTypes = []int{
  41. constant.ChannelTypeMidjourney,
  42. constant.ChannelTypeMidjourneyPlus,
  43. constant.ChannelTypeSunoAPI,
  44. constant.ChannelTypeKling,
  45. constant.ChannelTypeJimeng,
  46. constant.ChannelTypeDoubaoVideo,
  47. constant.ChannelTypeVidu,
  48. }
  49. if lo.Contains(unsupportedTestChannelTypes, channel.Type) {
  50. channelTypeName := constant.GetChannelTypeName(channel.Type)
  51. return testResult{
  52. localErr: fmt.Errorf("%s channel test is not supported", channelTypeName),
  53. }
  54. }
  55. w := httptest.NewRecorder()
  56. c, _ := gin.CreateTestContext(w)
  57. testModel = strings.TrimSpace(testModel)
  58. if testModel == "" {
  59. if channel.TestModel != nil && *channel.TestModel != "" {
  60. testModel = strings.TrimSpace(*channel.TestModel)
  61. } else {
  62. models := channel.GetModels()
  63. if len(models) > 0 {
  64. testModel = strings.TrimSpace(models[0])
  65. }
  66. if testModel == "" {
  67. testModel = "gpt-4o-mini"
  68. }
  69. }
  70. }
  71. requestPath := "/v1/chat/completions"
  72. // 如果指定了端点类型,使用指定的端点类型
  73. if endpointType != "" {
  74. if endpointInfo, ok := common.GetDefaultEndpointInfo(constant.EndpointType(endpointType)); ok {
  75. requestPath = endpointInfo.Path
  76. }
  77. } else {
  78. // 如果没有指定端点类型,使用原有的自动检测逻辑
  79. if strings.Contains(strings.ToLower(testModel), "rerank") {
  80. requestPath = "/v1/rerank"
  81. }
  82. // 先判断是否为 Embedding 模型
  83. if strings.Contains(strings.ToLower(testModel), "embedding") ||
  84. strings.HasPrefix(testModel, "m3e") || // m3e 系列模型
  85. strings.Contains(testModel, "bge-") || // bge 系列模型
  86. strings.Contains(testModel, "embed") ||
  87. channel.Type == constant.ChannelTypeMokaAI { // 其他 embedding 模型
  88. requestPath = "/v1/embeddings" // 修改请求路径
  89. }
  90. // VolcEngine 图像生成模型
  91. if channel.Type == constant.ChannelTypeVolcEngine && strings.Contains(testModel, "seedream") {
  92. requestPath = "/v1/images/generations"
  93. }
  94. // responses-only models
  95. if strings.Contains(strings.ToLower(testModel), "codex") {
  96. requestPath = "/v1/responses"
  97. }
  98. // responses compaction models (must use /v1/responses/compact)
  99. if strings.HasSuffix(testModel, ratio_setting.CompactModelSuffix) {
  100. requestPath = "/v1/responses/compact"
  101. }
  102. }
  103. if strings.HasPrefix(requestPath, "/v1/responses/compact") {
  104. testModel = ratio_setting.WithCompactModelSuffix(testModel)
  105. }
  106. c.Request = &http.Request{
  107. Method: "POST",
  108. URL: &url.URL{Path: requestPath}, // 使用动态路径
  109. Body: nil,
  110. Header: make(http.Header),
  111. }
  112. cache, err := model.GetUserCache(1)
  113. if err != nil {
  114. return testResult{
  115. localErr: err,
  116. newAPIError: nil,
  117. }
  118. }
  119. cache.WriteContext(c)
  120. //c.Request.Header.Set("Authorization", "Bearer "+channel.Key)
  121. c.Request.Header.Set("Content-Type", "application/json")
  122. c.Set("channel", channel.Type)
  123. c.Set("base_url", channel.GetBaseURL())
  124. group, _ := model.GetUserGroup(1, false)
  125. c.Set("group", group)
  126. newAPIError := middleware.SetupContextForSelectedChannel(c, channel, testModel)
  127. if newAPIError != nil {
  128. return testResult{
  129. context: c,
  130. localErr: newAPIError,
  131. newAPIError: newAPIError,
  132. }
  133. }
  134. // Determine relay format based on endpoint type or request path
  135. var relayFormat types.RelayFormat
  136. if endpointType != "" {
  137. // 根据指定的端点类型设置 relayFormat
  138. switch constant.EndpointType(endpointType) {
  139. case constant.EndpointTypeOpenAI:
  140. relayFormat = types.RelayFormatOpenAI
  141. case constant.EndpointTypeOpenAIResponse:
  142. relayFormat = types.RelayFormatOpenAIResponses
  143. case constant.EndpointTypeOpenAIResponseCompact:
  144. relayFormat = types.RelayFormatOpenAIResponsesCompaction
  145. case constant.EndpointTypeAnthropic:
  146. relayFormat = types.RelayFormatClaude
  147. case constant.EndpointTypeGemini:
  148. relayFormat = types.RelayFormatGemini
  149. case constant.EndpointTypeJinaRerank:
  150. relayFormat = types.RelayFormatRerank
  151. case constant.EndpointTypeImageGeneration:
  152. relayFormat = types.RelayFormatOpenAIImage
  153. case constant.EndpointTypeEmbeddings:
  154. relayFormat = types.RelayFormatEmbedding
  155. default:
  156. relayFormat = types.RelayFormatOpenAI
  157. }
  158. } else {
  159. // 根据请求路径自动检测
  160. relayFormat = types.RelayFormatOpenAI
  161. if c.Request.URL.Path == "/v1/embeddings" {
  162. relayFormat = types.RelayFormatEmbedding
  163. }
  164. if c.Request.URL.Path == "/v1/images/generations" {
  165. relayFormat = types.RelayFormatOpenAIImage
  166. }
  167. if c.Request.URL.Path == "/v1/messages" {
  168. relayFormat = types.RelayFormatClaude
  169. }
  170. if strings.Contains(c.Request.URL.Path, "/v1beta/models") {
  171. relayFormat = types.RelayFormatGemini
  172. }
  173. if c.Request.URL.Path == "/v1/rerank" || c.Request.URL.Path == "/rerank" {
  174. relayFormat = types.RelayFormatRerank
  175. }
  176. if c.Request.URL.Path == "/v1/responses" {
  177. relayFormat = types.RelayFormatOpenAIResponses
  178. }
  179. if strings.HasPrefix(c.Request.URL.Path, "/v1/responses/compact") {
  180. relayFormat = types.RelayFormatOpenAIResponsesCompaction
  181. }
  182. }
  183. request := buildTestRequest(testModel, endpointType, channel)
  184. info, err := relaycommon.GenRelayInfo(c, relayFormat, request, nil)
  185. if err != nil {
  186. return testResult{
  187. context: c,
  188. localErr: err,
  189. newAPIError: types.NewError(err, types.ErrorCodeGenRelayInfoFailed),
  190. }
  191. }
  192. info.IsChannelTest = true
  193. info.InitChannelMeta(c)
  194. err = helper.ModelMappedHelper(c, info, request)
  195. if err != nil {
  196. return testResult{
  197. context: c,
  198. localErr: err,
  199. newAPIError: types.NewError(err, types.ErrorCodeChannelModelMappedError),
  200. }
  201. }
  202. testModel = info.UpstreamModelName
  203. // 更新请求中的模型名称
  204. request.SetModelName(testModel)
  205. apiType, _ := common.ChannelType2APIType(channel.Type)
  206. if info.RelayMode == relayconstant.RelayModeResponsesCompact &&
  207. apiType != constant.APITypeOpenAI &&
  208. apiType != constant.APITypeCodex {
  209. return testResult{
  210. context: c,
  211. localErr: fmt.Errorf("responses compaction test only supports openai/codex channels, got api type %d", apiType),
  212. newAPIError: types.NewError(fmt.Errorf("unsupported api type: %d", apiType), types.ErrorCodeInvalidApiType),
  213. }
  214. }
  215. adaptor := relay.GetAdaptor(apiType)
  216. if adaptor == nil {
  217. return testResult{
  218. context: c,
  219. localErr: fmt.Errorf("invalid api type: %d, adaptor is nil", apiType),
  220. newAPIError: types.NewError(fmt.Errorf("invalid api type: %d, adaptor is nil", apiType), types.ErrorCodeInvalidApiType),
  221. }
  222. }
  223. //// 创建一个用于日志的 info 副本,移除 ApiKey
  224. //logInfo := info
  225. //logInfo.ApiKey = ""
  226. common.SysLog(fmt.Sprintf("testing channel %d with model %s , info %+v ", channel.Id, testModel, info.ToString()))
  227. priceData, err := helper.ModelPriceHelper(c, info, 0, request.GetTokenCountMeta())
  228. if err != nil {
  229. return testResult{
  230. context: c,
  231. localErr: err,
  232. newAPIError: types.NewError(err, types.ErrorCodeModelPriceError),
  233. }
  234. }
  235. adaptor.Init(info)
  236. var convertedRequest any
  237. // 根据 RelayMode 选择正确的转换函数
  238. switch info.RelayMode {
  239. case relayconstant.RelayModeEmbeddings:
  240. // Embedding 请求 - request 已经是正确的类型
  241. if embeddingReq, ok := request.(*dto.EmbeddingRequest); ok {
  242. convertedRequest, err = adaptor.ConvertEmbeddingRequest(c, info, *embeddingReq)
  243. } else {
  244. return testResult{
  245. context: c,
  246. localErr: errors.New("invalid embedding request type"),
  247. newAPIError: types.NewError(errors.New("invalid embedding request type"), types.ErrorCodeConvertRequestFailed),
  248. }
  249. }
  250. case relayconstant.RelayModeImagesGenerations:
  251. // 图像生成请求 - request 已经是正确的类型
  252. if imageReq, ok := request.(*dto.ImageRequest); ok {
  253. convertedRequest, err = adaptor.ConvertImageRequest(c, info, *imageReq)
  254. } else {
  255. return testResult{
  256. context: c,
  257. localErr: errors.New("invalid image request type"),
  258. newAPIError: types.NewError(errors.New("invalid image request type"), types.ErrorCodeConvertRequestFailed),
  259. }
  260. }
  261. case relayconstant.RelayModeRerank:
  262. // Rerank 请求 - request 已经是正确的类型
  263. if rerankReq, ok := request.(*dto.RerankRequest); ok {
  264. convertedRequest, err = adaptor.ConvertRerankRequest(c, info.RelayMode, *rerankReq)
  265. } else {
  266. return testResult{
  267. context: c,
  268. localErr: errors.New("invalid rerank request type"),
  269. newAPIError: types.NewError(errors.New("invalid rerank request type"), types.ErrorCodeConvertRequestFailed),
  270. }
  271. }
  272. case relayconstant.RelayModeResponses:
  273. // Response 请求 - request 已经是正确的类型
  274. if responseReq, ok := request.(*dto.OpenAIResponsesRequest); ok {
  275. convertedRequest, err = adaptor.ConvertOpenAIResponsesRequest(c, info, *responseReq)
  276. } else {
  277. return testResult{
  278. context: c,
  279. localErr: errors.New("invalid response request type"),
  280. newAPIError: types.NewError(errors.New("invalid response request type"), types.ErrorCodeConvertRequestFailed),
  281. }
  282. }
  283. case relayconstant.RelayModeResponsesCompact:
  284. // Response compaction request - convert to OpenAIResponsesRequest before adapting
  285. switch req := request.(type) {
  286. case *dto.OpenAIResponsesCompactionRequest:
  287. convertedRequest, err = adaptor.ConvertOpenAIResponsesRequest(c, info, dto.OpenAIResponsesRequest{
  288. Model: req.Model,
  289. Input: req.Input,
  290. Instructions: req.Instructions,
  291. PreviousResponseID: req.PreviousResponseID,
  292. })
  293. case *dto.OpenAIResponsesRequest:
  294. convertedRequest, err = adaptor.ConvertOpenAIResponsesRequest(c, info, *req)
  295. default:
  296. return testResult{
  297. context: c,
  298. localErr: errors.New("invalid response compaction request type"),
  299. newAPIError: types.NewError(errors.New("invalid response compaction request type"), types.ErrorCodeConvertRequestFailed),
  300. }
  301. }
  302. default:
  303. // Chat/Completion 等其他请求类型
  304. if generalReq, ok := request.(*dto.GeneralOpenAIRequest); ok {
  305. convertedRequest, err = adaptor.ConvertOpenAIRequest(c, info, generalReq)
  306. } else {
  307. return testResult{
  308. context: c,
  309. localErr: errors.New("invalid general request type"),
  310. newAPIError: types.NewError(errors.New("invalid general request type"), types.ErrorCodeConvertRequestFailed),
  311. }
  312. }
  313. }
  314. if err != nil {
  315. return testResult{
  316. context: c,
  317. localErr: err,
  318. newAPIError: types.NewError(err, types.ErrorCodeConvertRequestFailed),
  319. }
  320. }
  321. jsonData, err := json.Marshal(convertedRequest)
  322. if err != nil {
  323. return testResult{
  324. context: c,
  325. localErr: err,
  326. newAPIError: types.NewError(err, types.ErrorCodeJsonMarshalFailed),
  327. }
  328. }
  329. //jsonData, err = relaycommon.RemoveDisabledFields(jsonData, info.ChannelOtherSettings)
  330. //if err != nil {
  331. // return testResult{
  332. // context: c,
  333. // localErr: err,
  334. // newAPIError: types.NewError(err, types.ErrorCodeConvertRequestFailed),
  335. // }
  336. //}
  337. if len(info.ParamOverride) > 0 {
  338. jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info))
  339. if err != nil {
  340. return testResult{
  341. context: c,
  342. localErr: err,
  343. newAPIError: types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid),
  344. }
  345. }
  346. }
  347. requestBody := bytes.NewBuffer(jsonData)
  348. c.Request.Body = io.NopCloser(bytes.NewBuffer(jsonData))
  349. resp, err := adaptor.DoRequest(c, info, requestBody)
  350. if err != nil {
  351. return testResult{
  352. context: c,
  353. localErr: err,
  354. newAPIError: types.NewOpenAIError(err, types.ErrorCodeDoRequestFailed, http.StatusInternalServerError),
  355. }
  356. }
  357. var httpResp *http.Response
  358. if resp != nil {
  359. httpResp = resp.(*http.Response)
  360. if httpResp.StatusCode != http.StatusOK {
  361. err := service.RelayErrorHandler(c.Request.Context(), httpResp, true)
  362. common.SysError(fmt.Sprintf(
  363. "channel test bad response: channel_id=%d name=%s type=%d model=%s endpoint_type=%s status=%d err=%v",
  364. channel.Id,
  365. channel.Name,
  366. channel.Type,
  367. testModel,
  368. endpointType,
  369. httpResp.StatusCode,
  370. err,
  371. ))
  372. return testResult{
  373. context: c,
  374. localErr: err,
  375. newAPIError: types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError),
  376. }
  377. }
  378. }
  379. usageA, respErr := adaptor.DoResponse(c, httpResp, info)
  380. if respErr != nil {
  381. return testResult{
  382. context: c,
  383. localErr: respErr,
  384. newAPIError: respErr,
  385. }
  386. }
  387. if usageA == nil {
  388. return testResult{
  389. context: c,
  390. localErr: errors.New("usage is nil"),
  391. newAPIError: types.NewOpenAIError(errors.New("usage is nil"), types.ErrorCodeBadResponseBody, http.StatusInternalServerError),
  392. }
  393. }
  394. usage := usageA.(*dto.Usage)
  395. result := w.Result()
  396. respBody, err := io.ReadAll(result.Body)
  397. if err != nil {
  398. return testResult{
  399. context: c,
  400. localErr: err,
  401. newAPIError: types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError),
  402. }
  403. }
  404. info.SetEstimatePromptTokens(usage.PromptTokens)
  405. quota := 0
  406. if !priceData.UsePrice {
  407. quota = usage.PromptTokens + int(math.Round(float64(usage.CompletionTokens)*priceData.CompletionRatio))
  408. quota = int(math.Round(float64(quota) * priceData.ModelRatio))
  409. if priceData.ModelRatio != 0 && quota <= 0 {
  410. quota = 1
  411. }
  412. } else {
  413. quota = int(priceData.ModelPrice * common.QuotaPerUnit)
  414. }
  415. tok := time.Now()
  416. milliseconds := tok.Sub(tik).Milliseconds()
  417. consumedTime := float64(milliseconds) / 1000.0
  418. other := service.GenerateTextOtherInfo(c, info, priceData.ModelRatio, priceData.GroupRatioInfo.GroupRatio, priceData.CompletionRatio,
  419. usage.PromptTokensDetails.CachedTokens, priceData.CacheRatio, priceData.ModelPrice, priceData.GroupRatioInfo.GroupSpecialRatio)
  420. model.RecordConsumeLog(c, 1, model.RecordConsumeLogParams{
  421. ChannelId: channel.Id,
  422. PromptTokens: usage.PromptTokens,
  423. CompletionTokens: usage.CompletionTokens,
  424. ModelName: info.OriginModelName,
  425. TokenName: "模型测试",
  426. Quota: quota,
  427. Content: "模型测试",
  428. UseTimeSeconds: int(consumedTime),
  429. IsStream: info.IsStream,
  430. Group: info.UsingGroup,
  431. Other: other,
  432. })
  433. common.SysLog(fmt.Sprintf("testing channel #%d, response: \n%s", channel.Id, string(respBody)))
  434. return testResult{
  435. context: c,
  436. localErr: nil,
  437. newAPIError: nil,
  438. }
  439. }
  440. func buildTestRequest(model string, endpointType string, channel *model.Channel) dto.Request {
  441. testResponsesInput := json.RawMessage(`[{"role":"user","content":"hi"}]`)
  442. // 根据端点类型构建不同的测试请求
  443. if endpointType != "" {
  444. switch constant.EndpointType(endpointType) {
  445. case constant.EndpointTypeEmbeddings:
  446. // 返回 EmbeddingRequest
  447. return &dto.EmbeddingRequest{
  448. Model: model,
  449. Input: []any{"hello world"},
  450. }
  451. case constant.EndpointTypeImageGeneration:
  452. // 返回 ImageRequest
  453. return &dto.ImageRequest{
  454. Model: model,
  455. Prompt: "a cute cat",
  456. N: 1,
  457. Size: "1024x1024",
  458. }
  459. case constant.EndpointTypeJinaRerank:
  460. // 返回 RerankRequest
  461. return &dto.RerankRequest{
  462. Model: model,
  463. Query: "What is Deep Learning?",
  464. Documents: []any{"Deep Learning is a subset of machine learning.", "Machine learning is a field of artificial intelligence."},
  465. TopN: 2,
  466. }
  467. case constant.EndpointTypeOpenAIResponse:
  468. // 返回 OpenAIResponsesRequest
  469. return &dto.OpenAIResponsesRequest{
  470. Model: model,
  471. Input: json.RawMessage(`[{"role":"user","content":"hi"}]`),
  472. }
  473. case constant.EndpointTypeOpenAIResponseCompact:
  474. // 返回 OpenAIResponsesCompactionRequest
  475. return &dto.OpenAIResponsesCompactionRequest{
  476. Model: model,
  477. Input: testResponsesInput,
  478. }
  479. case constant.EndpointTypeAnthropic, constant.EndpointTypeGemini, constant.EndpointTypeOpenAI:
  480. // 返回 GeneralOpenAIRequest
  481. maxTokens := uint(16)
  482. if constant.EndpointType(endpointType) == constant.EndpointTypeGemini {
  483. maxTokens = 3000
  484. }
  485. return &dto.GeneralOpenAIRequest{
  486. Model: model,
  487. Stream: false,
  488. Messages: []dto.Message{
  489. {
  490. Role: "user",
  491. Content: "hi",
  492. },
  493. },
  494. MaxTokens: maxTokens,
  495. }
  496. }
  497. }
  498. // 自动检测逻辑(保持原有行为)
  499. if strings.Contains(strings.ToLower(model), "rerank") {
  500. return &dto.RerankRequest{
  501. Model: model,
  502. Query: "What is Deep Learning?",
  503. Documents: []any{"Deep Learning is a subset of machine learning.", "Machine learning is a field of artificial intelligence."},
  504. TopN: 2,
  505. }
  506. }
  507. // 先判断是否为 Embedding 模型
  508. if strings.Contains(strings.ToLower(model), "embedding") ||
  509. strings.HasPrefix(model, "m3e") ||
  510. strings.Contains(model, "bge-") {
  511. // 返回 EmbeddingRequest
  512. return &dto.EmbeddingRequest{
  513. Model: model,
  514. Input: []any{"hello world"},
  515. }
  516. }
  517. // Responses compaction models (must use /v1/responses/compact)
  518. if strings.HasSuffix(model, ratio_setting.CompactModelSuffix) {
  519. return &dto.OpenAIResponsesCompactionRequest{
  520. Model: model,
  521. Input: testResponsesInput,
  522. }
  523. }
  524. // Responses-only models (e.g. codex series)
  525. if strings.Contains(strings.ToLower(model), "codex") {
  526. return &dto.OpenAIResponsesRequest{
  527. Model: model,
  528. Input: json.RawMessage(`[{"role":"user","content":"hi"}]`),
  529. }
  530. }
  531. // Chat/Completion 请求 - 返回 GeneralOpenAIRequest
  532. testRequest := &dto.GeneralOpenAIRequest{
  533. Model: model,
  534. Stream: false,
  535. Messages: []dto.Message{
  536. {
  537. Role: "user",
  538. Content: "hi",
  539. },
  540. },
  541. }
  542. if strings.HasPrefix(model, "o") {
  543. testRequest.MaxCompletionTokens = 16
  544. } else if strings.Contains(model, "thinking") {
  545. if !strings.Contains(model, "claude") {
  546. testRequest.MaxTokens = 50
  547. }
  548. } else if strings.Contains(model, "gemini") {
  549. testRequest.MaxTokens = 3000
  550. } else {
  551. testRequest.MaxTokens = 16
  552. }
  553. return testRequest
  554. }
  555. func TestChannel(c *gin.Context) {
  556. channelId, err := strconv.Atoi(c.Param("id"))
  557. if err != nil {
  558. common.ApiError(c, err)
  559. return
  560. }
  561. channel, err := model.CacheGetChannel(channelId)
  562. if err != nil {
  563. channel, err = model.GetChannelById(channelId, true)
  564. if err != nil {
  565. common.ApiError(c, err)
  566. return
  567. }
  568. }
  569. //defer func() {
  570. // if channel.ChannelInfo.IsMultiKey {
  571. // go func() { _ = channel.SaveChannelInfo() }()
  572. // }
  573. //}()
  574. testModel := c.Query("model")
  575. endpointType := c.Query("endpoint_type")
  576. tik := time.Now()
  577. result := testChannel(channel, testModel, endpointType)
  578. if result.localErr != nil {
  579. c.JSON(http.StatusOK, gin.H{
  580. "success": false,
  581. "message": result.localErr.Error(),
  582. "time": 0.0,
  583. })
  584. return
  585. }
  586. tok := time.Now()
  587. milliseconds := tok.Sub(tik).Milliseconds()
  588. go channel.UpdateResponseTime(milliseconds)
  589. consumedTime := float64(milliseconds) / 1000.0
  590. if result.newAPIError != nil {
  591. c.JSON(http.StatusOK, gin.H{
  592. "success": false,
  593. "message": result.newAPIError.Error(),
  594. "time": consumedTime,
  595. })
  596. return
  597. }
  598. c.JSON(http.StatusOK, gin.H{
  599. "success": true,
  600. "message": "",
  601. "time": consumedTime,
  602. })
  603. }
  604. var testAllChannelsLock sync.Mutex
  605. var testAllChannelsRunning bool = false
  606. func testAllChannels(notify bool) error {
  607. testAllChannelsLock.Lock()
  608. if testAllChannelsRunning {
  609. testAllChannelsLock.Unlock()
  610. return errors.New("测试已在运行中")
  611. }
  612. testAllChannelsRunning = true
  613. testAllChannelsLock.Unlock()
  614. channels, getChannelErr := model.GetAllChannels(0, 0, true, false)
  615. if getChannelErr != nil {
  616. return getChannelErr
  617. }
  618. var disableThreshold = int64(common.ChannelDisableThreshold * 1000)
  619. if disableThreshold == 0 {
  620. disableThreshold = 10000000 // a impossible value
  621. }
  622. gopool.Go(func() {
  623. // 使用 defer 确保无论如何都会重置运行状态,防止死锁
  624. defer func() {
  625. testAllChannelsLock.Lock()
  626. testAllChannelsRunning = false
  627. testAllChannelsLock.Unlock()
  628. }()
  629. for _, channel := range channels {
  630. isChannelEnabled := channel.Status == common.ChannelStatusEnabled
  631. tik := time.Now()
  632. result := testChannel(channel, "", "")
  633. tok := time.Now()
  634. milliseconds := tok.Sub(tik).Milliseconds()
  635. shouldBanChannel := false
  636. newAPIError := result.newAPIError
  637. // request error disables the channel
  638. if newAPIError != nil {
  639. shouldBanChannel = service.ShouldDisableChannel(channel.Type, result.newAPIError)
  640. }
  641. // 当错误检查通过,才检查响应时间
  642. if common.AutomaticDisableChannelEnabled && !shouldBanChannel {
  643. if milliseconds > disableThreshold {
  644. err := fmt.Errorf("响应时间 %.2fs 超过阈值 %.2fs", float64(milliseconds)/1000.0, float64(disableThreshold)/1000.0)
  645. newAPIError = types.NewOpenAIError(err, types.ErrorCodeChannelResponseTimeExceeded, http.StatusRequestTimeout)
  646. shouldBanChannel = true
  647. }
  648. }
  649. // disable channel
  650. if isChannelEnabled && shouldBanChannel && channel.GetAutoBan() {
  651. processChannelError(result.context, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(result.context, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError)
  652. }
  653. // enable channel
  654. if !isChannelEnabled && service.ShouldEnableChannel(newAPIError, channel.Status) {
  655. service.EnableChannel(channel.Id, common.GetContextKeyString(result.context, constant.ContextKeyChannelKey), channel.Name)
  656. }
  657. channel.UpdateResponseTime(milliseconds)
  658. time.Sleep(common.RequestInterval)
  659. }
  660. if notify {
  661. service.NotifyRootUser(dto.NotifyTypeChannelTest, "通道测试完成", "所有通道测试已完成")
  662. }
  663. })
  664. return nil
  665. }
  666. func TestAllChannels(c *gin.Context) {
  667. err := testAllChannels(true)
  668. if err != nil {
  669. common.ApiError(c, err)
  670. return
  671. }
  672. c.JSON(http.StatusOK, gin.H{
  673. "success": true,
  674. "message": "",
  675. })
  676. }
  677. var autoTestChannelsOnce sync.Once
  678. func AutomaticallyTestChannels() {
  679. // 只在Master节点定时测试渠道
  680. if !common.IsMasterNode {
  681. return
  682. }
  683. autoTestChannelsOnce.Do(func() {
  684. for {
  685. if !operation_setting.GetMonitorSetting().AutoTestChannelEnabled {
  686. time.Sleep(1 * time.Minute)
  687. continue
  688. }
  689. for {
  690. frequency := operation_setting.GetMonitorSetting().AutoTestChannelMinutes
  691. time.Sleep(time.Duration(int(math.Round(frequency))) * time.Minute)
  692. common.SysLog(fmt.Sprintf("automatically test channels with interval %f minutes", frequency))
  693. common.SysLog("automatically testing all channels")
  694. _ = testAllChannels(false)
  695. common.SysLog("automatically channel test finished")
  696. if !operation_setting.GetMonitorSetting().AutoTestChannelEnabled {
  697. break
  698. }
  699. }
  700. }
  701. })
  702. }