channel-test.go 30 KB

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