channel-test.go 28 KB

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