channel-test.go 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940
  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 := validateTestResponseBody(respBody, isStream); 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 validateStreamTestResponseBody(respBody []byte) error {
  531. b := bytes.TrimSpace(respBody)
  532. if len(b) == 0 {
  533. return errors.New("stream response body is empty")
  534. }
  535. for _, line := range bytes.Split(b, []byte{'\n'}) {
  536. line = bytes.TrimSpace(line)
  537. if len(line) == 0 || !bytes.HasPrefix(line, []byte("data:")) {
  538. continue
  539. }
  540. payload := bytes.TrimSpace(bytes.TrimPrefix(line, []byte("data:")))
  541. if len(payload) == 0 || bytes.Equal(payload, []byte("[DONE]")) {
  542. continue
  543. }
  544. return nil
  545. }
  546. return errors.New("stream response body does not contain a valid stream event")
  547. }
  548. func validateTestResponseBody(respBody []byte, isStream bool) error {
  549. if bodyErr := detectErrorFromTestResponseBody(respBody); bodyErr != nil {
  550. return bodyErr
  551. }
  552. if isStream {
  553. return validateStreamTestResponseBody(respBody)
  554. }
  555. return nil
  556. }
  557. func shouldUseStreamForAutomaticChannelTest(channel *model.Channel) bool {
  558. return channel != nil && channel.Type == constant.ChannelTypeCodex
  559. }
  560. func detectErrorMessageFromJSONBytes(jsonBytes []byte) string {
  561. if len(jsonBytes) == 0 {
  562. return ""
  563. }
  564. if jsonBytes[0] != '{' && jsonBytes[0] != '[' {
  565. return ""
  566. }
  567. errVal := gjson.GetBytes(jsonBytes, "error")
  568. if !errVal.Exists() || errVal.Type == gjson.Null {
  569. return ""
  570. }
  571. message := gjson.GetBytes(jsonBytes, "error.message").String()
  572. if message == "" {
  573. message = gjson.GetBytes(jsonBytes, "error.error.message").String()
  574. }
  575. if message == "" && errVal.Type == gjson.String {
  576. message = errVal.String()
  577. }
  578. if message == "" {
  579. message = errVal.Raw
  580. }
  581. message = strings.TrimSpace(message)
  582. if message == "" {
  583. return "upstream returned error payload"
  584. }
  585. return message
  586. }
  587. func buildTestRequest(model string, endpointType string, channel *model.Channel, isStream bool) dto.Request {
  588. testResponsesInput := json.RawMessage(`[{"role":"user","content":"hi"}]`)
  589. // 根据端点类型构建不同的测试请求
  590. if endpointType != "" {
  591. switch constant.EndpointType(endpointType) {
  592. case constant.EndpointTypeEmbeddings:
  593. // 返回 EmbeddingRequest
  594. return &dto.EmbeddingRequest{
  595. Model: model,
  596. Input: []any{"hello world"},
  597. }
  598. case constant.EndpointTypeImageGeneration:
  599. // 返回 ImageRequest
  600. return &dto.ImageRequest{
  601. Model: model,
  602. Prompt: "a cute cat",
  603. N: lo.ToPtr(uint(1)),
  604. Size: "1024x1024",
  605. }
  606. case constant.EndpointTypeJinaRerank:
  607. // 返回 RerankRequest
  608. return &dto.RerankRequest{
  609. Model: model,
  610. Query: "What is Deep Learning?",
  611. Documents: []any{"Deep Learning is a subset of machine learning.", "Machine learning is a field of artificial intelligence."},
  612. TopN: lo.ToPtr(2),
  613. }
  614. case constant.EndpointTypeOpenAIResponse:
  615. // 返回 OpenAIResponsesRequest
  616. return &dto.OpenAIResponsesRequest{
  617. Model: model,
  618. Input: json.RawMessage(`[{"role":"user","content":"hi"}]`),
  619. Stream: lo.ToPtr(isStream),
  620. }
  621. case constant.EndpointTypeOpenAIResponseCompact:
  622. // 返回 OpenAIResponsesCompactionRequest
  623. return &dto.OpenAIResponsesCompactionRequest{
  624. Model: model,
  625. Input: testResponsesInput,
  626. }
  627. case constant.EndpointTypeAnthropic, constant.EndpointTypeGemini, constant.EndpointTypeOpenAI:
  628. // 返回 GeneralOpenAIRequest
  629. maxTokens := uint(16)
  630. if constant.EndpointType(endpointType) == constant.EndpointTypeGemini {
  631. maxTokens = 3000
  632. }
  633. req := &dto.GeneralOpenAIRequest{
  634. Model: model,
  635. Stream: lo.ToPtr(isStream),
  636. Messages: []dto.Message{
  637. {
  638. Role: "user",
  639. Content: "hi",
  640. },
  641. },
  642. MaxTokens: lo.ToPtr(maxTokens),
  643. }
  644. if isStream {
  645. req.StreamOptions = &dto.StreamOptions{IncludeUsage: true}
  646. }
  647. return req
  648. }
  649. }
  650. // 自动检测逻辑(保持原有行为)
  651. if strings.Contains(strings.ToLower(model), "rerank") {
  652. return &dto.RerankRequest{
  653. Model: model,
  654. Query: "What is Deep Learning?",
  655. Documents: []any{"Deep Learning is a subset of machine learning.", "Machine learning is a field of artificial intelligence."},
  656. TopN: lo.ToPtr(2),
  657. }
  658. }
  659. // 先判断是否为 Embedding 模型
  660. if strings.Contains(strings.ToLower(model), "embedding") ||
  661. strings.HasPrefix(model, "m3e") ||
  662. strings.Contains(model, "bge-") {
  663. // 返回 EmbeddingRequest
  664. return &dto.EmbeddingRequest{
  665. Model: model,
  666. Input: []any{"hello world"},
  667. }
  668. }
  669. // Responses compaction models (must use /v1/responses/compact)
  670. if strings.HasSuffix(model, ratio_setting.CompactModelSuffix) {
  671. return &dto.OpenAIResponsesCompactionRequest{
  672. Model: model,
  673. Input: testResponsesInput,
  674. }
  675. }
  676. // Responses-only models (e.g. codex series)
  677. if strings.Contains(strings.ToLower(model), "codex") {
  678. return &dto.OpenAIResponsesRequest{
  679. Model: model,
  680. Input: json.RawMessage(`[{"role":"user","content":"hi"}]`),
  681. Stream: lo.ToPtr(isStream),
  682. }
  683. }
  684. // Chat/Completion 请求 - 返回 GeneralOpenAIRequest
  685. testRequest := &dto.GeneralOpenAIRequest{
  686. Model: model,
  687. Stream: lo.ToPtr(isStream),
  688. Messages: []dto.Message{
  689. {
  690. Role: "user",
  691. Content: "hi",
  692. },
  693. },
  694. }
  695. if isStream {
  696. testRequest.StreamOptions = &dto.StreamOptions{IncludeUsage: true}
  697. }
  698. if strings.HasPrefix(model, "o") {
  699. testRequest.MaxCompletionTokens = lo.ToPtr(uint(16))
  700. } else if strings.Contains(model, "thinking") {
  701. if !strings.Contains(model, "claude") {
  702. testRequest.MaxTokens = lo.ToPtr(uint(50))
  703. }
  704. } else if strings.Contains(model, "gemini") {
  705. testRequest.MaxTokens = lo.ToPtr(uint(3000))
  706. } else {
  707. testRequest.MaxTokens = lo.ToPtr(uint(16))
  708. }
  709. return testRequest
  710. }
  711. func TestChannel(c *gin.Context) {
  712. channelId, err := strconv.Atoi(c.Param("id"))
  713. if err != nil {
  714. common.ApiError(c, err)
  715. return
  716. }
  717. channel, err := model.CacheGetChannel(channelId)
  718. if err != nil {
  719. channel, err = model.GetChannelById(channelId, true)
  720. if err != nil {
  721. common.ApiError(c, err)
  722. return
  723. }
  724. }
  725. //defer func() {
  726. // if channel.ChannelInfo.IsMultiKey {
  727. // go func() { _ = channel.SaveChannelInfo() }()
  728. // }
  729. //}()
  730. testModel := c.Query("model")
  731. endpointType := c.Query("endpoint_type")
  732. isStream, _ := strconv.ParseBool(c.Query("stream"))
  733. tik := time.Now()
  734. result := testChannel(channel, testModel, endpointType, isStream)
  735. if result.localErr != nil {
  736. resp := gin.H{
  737. "success": false,
  738. "message": result.localErr.Error(),
  739. "time": 0.0,
  740. }
  741. if result.newAPIError != nil {
  742. resp["error_code"] = result.newAPIError.GetErrorCode()
  743. }
  744. c.JSON(http.StatusOK, resp)
  745. return
  746. }
  747. tok := time.Now()
  748. milliseconds := tok.Sub(tik).Milliseconds()
  749. go channel.UpdateResponseTime(milliseconds)
  750. consumedTime := float64(milliseconds) / 1000.0
  751. if result.newAPIError != nil {
  752. c.JSON(http.StatusOK, gin.H{
  753. "success": false,
  754. "message": result.newAPIError.Error(),
  755. "time": consumedTime,
  756. "error_code": result.newAPIError.GetErrorCode(),
  757. })
  758. return
  759. }
  760. c.JSON(http.StatusOK, gin.H{
  761. "success": true,
  762. "message": "",
  763. "time": consumedTime,
  764. })
  765. }
  766. var testAllChannelsLock sync.Mutex
  767. var testAllChannelsRunning bool = false
  768. func testAllChannels(notify bool) error {
  769. testAllChannelsLock.Lock()
  770. if testAllChannelsRunning {
  771. testAllChannelsLock.Unlock()
  772. return errors.New("测试已在运行中")
  773. }
  774. testAllChannelsRunning = true
  775. testAllChannelsLock.Unlock()
  776. channels, getChannelErr := model.GetAllChannels(0, 0, true, false)
  777. if getChannelErr != nil {
  778. return getChannelErr
  779. }
  780. var disableThreshold = int64(common.ChannelDisableThreshold * 1000)
  781. if disableThreshold == 0 {
  782. disableThreshold = 10000000 // a impossible value
  783. }
  784. gopool.Go(func() {
  785. // 使用 defer 确保无论如何都会重置运行状态,防止死锁
  786. defer func() {
  787. testAllChannelsLock.Lock()
  788. testAllChannelsRunning = false
  789. testAllChannelsLock.Unlock()
  790. }()
  791. for _, channel := range channels {
  792. if channel.Status == common.ChannelStatusManuallyDisabled {
  793. continue
  794. }
  795. isChannelEnabled := channel.Status == common.ChannelStatusEnabled
  796. tik := time.Now()
  797. result := testChannel(channel, "", "", shouldUseStreamForAutomaticChannelTest(channel))
  798. tok := time.Now()
  799. milliseconds := tok.Sub(tik).Milliseconds()
  800. shouldBanChannel := false
  801. newAPIError := result.newAPIError
  802. // request error disables the channel
  803. if newAPIError != nil {
  804. shouldBanChannel = service.ShouldDisableChannel(result.newAPIError)
  805. }
  806. // 当错误检查通过,才检查响应时间
  807. if common.AutomaticDisableChannelEnabled && !shouldBanChannel {
  808. if milliseconds > disableThreshold {
  809. err := fmt.Errorf("响应时间 %.2fs 超过阈值 %.2fs", float64(milliseconds)/1000.0, float64(disableThreshold)/1000.0)
  810. newAPIError = types.NewOpenAIError(err, types.ErrorCodeChannelResponseTimeExceeded, http.StatusRequestTimeout)
  811. shouldBanChannel = true
  812. }
  813. }
  814. // disable channel
  815. if isChannelEnabled && shouldBanChannel && channel.GetAutoBan() {
  816. processChannelError(result.context, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(result.context, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError)
  817. }
  818. // enable channel
  819. if !isChannelEnabled && service.ShouldEnableChannel(newAPIError, channel.Status) {
  820. service.EnableChannel(channel.Id, common.GetContextKeyString(result.context, constant.ContextKeyChannelKey), channel.Name)
  821. }
  822. channel.UpdateResponseTime(milliseconds)
  823. time.Sleep(common.RequestInterval)
  824. }
  825. if notify {
  826. service.NotifyRootUser(dto.NotifyTypeChannelTest, "通道测试完成", "所有通道测试已完成")
  827. }
  828. })
  829. return nil
  830. }
  831. func TestAllChannels(c *gin.Context) {
  832. err := testAllChannels(true)
  833. if err != nil {
  834. common.ApiError(c, err)
  835. return
  836. }
  837. c.JSON(http.StatusOK, gin.H{
  838. "success": true,
  839. "message": "",
  840. })
  841. }
  842. var autoTestChannelsOnce sync.Once
  843. func AutomaticallyTestChannels() {
  844. // 只在Master节点定时测试渠道
  845. if !common.IsMasterNode {
  846. return
  847. }
  848. autoTestChannelsOnce.Do(func() {
  849. for {
  850. if !operation_setting.GetMonitorSetting().AutoTestChannelEnabled {
  851. time.Sleep(1 * time.Minute)
  852. continue
  853. }
  854. for {
  855. frequency := operation_setting.GetMonitorSetting().AutoTestChannelMinutes
  856. time.Sleep(time.Duration(int(math.Round(frequency))) * time.Minute)
  857. common.SysLog(fmt.Sprintf("automatically test channels with interval %f minutes", frequency))
  858. common.SysLog("automatically testing all channels")
  859. _ = testAllChannels(false)
  860. common.SysLog("automatically channel test finished")
  861. if !operation_setting.GetMonitorSetting().AutoTestChannelEnabled {
  862. break
  863. }
  864. }
  865. }
  866. })
  867. }