channel-test.go 26 KB

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