channel-test.go 27 KB

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