Nevar pievienot vairāk kā 25 tēmas Tēmai ir jāsākas ar burtu vai ciparu, tā var saturēt domu zīmes ('-') un var būt līdz 35 simboliem gara.
 
 
 

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