25개 이상의 토픽을 선택하실 수 없습니다. Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

895 lines
27 KiB

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