mirror of
https://github.com/linux-do/new-api.git
synced 2025-09-17 16:06:38 +08:00
278 lines
7.4 KiB
Go
278 lines
7.4 KiB
Go
package controller
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"net/url"
|
|
"one-api/common"
|
|
"one-api/dto"
|
|
"one-api/model"
|
|
"one-api/relay"
|
|
relaycommon "one-api/relay/common"
|
|
"one-api/relay/constant"
|
|
"one-api/service"
|
|
"strconv"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func testChannel(channel *model.Channel, testModel string) (err error, openaiErr *dto.OpenAIError) {
|
|
if channel.Type == common.ChannelTypeMidjourney {
|
|
return errors.New("midjourney channel test is not supported"), nil
|
|
}
|
|
if channel.Type == common.ChannelTypeSunoAPI {
|
|
return errors.New("suno channel test is not supported"), nil
|
|
}
|
|
w := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(w)
|
|
c.Request = &http.Request{
|
|
Method: "POST",
|
|
URL: &url.URL{Path: "/v1/chat/completions"},
|
|
Body: nil,
|
|
Header: make(http.Header),
|
|
}
|
|
c.Request.Header.Set("Authorization", "Bearer "+channel.Key)
|
|
c.Request.Header.Set("Content-Type", "application/json")
|
|
c.Set("channel", channel.Type)
|
|
c.Set("base_url", channel.GetBaseURL())
|
|
switch channel.Type {
|
|
case common.ChannelTypeAzure:
|
|
c.Set("api_version", channel.Other)
|
|
case common.ChannelTypeXunfei:
|
|
c.Set("api_version", channel.Other)
|
|
//case common.ChannelTypeAIProxyLibrary:
|
|
// c.Set("library_id", channel.Other)
|
|
case common.ChannelTypeGemini:
|
|
c.Set("api_version", channel.Other)
|
|
case common.ChannelTypeAli:
|
|
c.Set("plugin", channel.Other)
|
|
}
|
|
|
|
meta := relaycommon.GenRelayInfo(c)
|
|
apiType, _ := constant.ChannelType2APIType(channel.Type)
|
|
adaptor := relay.GetAdaptor(apiType)
|
|
if adaptor == nil {
|
|
return fmt.Errorf("invalid api type: %d, adaptor is nil", apiType), nil
|
|
}
|
|
if testModel == "" {
|
|
if channel.TestModel != nil && *channel.TestModel != "" {
|
|
testModel = *channel.TestModel
|
|
} else {
|
|
testModel = adaptor.GetModelList()[0]
|
|
}
|
|
} else {
|
|
modelMapping := *channel.ModelMapping
|
|
if modelMapping != "" && modelMapping != "{}" {
|
|
modelMap := make(map[string]string)
|
|
err := json.Unmarshal([]byte(modelMapping), &modelMap)
|
|
if err != nil {
|
|
openaiErr := service.OpenAIErrorWrapperLocal(err, "unmarshal_model_mapping_failed", http.StatusInternalServerError).Error
|
|
return err, &openaiErr
|
|
}
|
|
if modelMap[testModel] != "" {
|
|
testModel = modelMap[testModel]
|
|
}
|
|
}
|
|
}
|
|
|
|
request := buildTestRequest()
|
|
request.Model = testModel
|
|
meta.UpstreamModelName = testModel
|
|
common.SysLog(fmt.Sprintf("testing channel %d with model %s", channel.Id, testModel))
|
|
|
|
adaptor.Init(meta, *request)
|
|
|
|
convertedRequest, err := adaptor.ConvertRequest(c, constant.RelayModeChatCompletions, request)
|
|
if err != nil {
|
|
return err, nil
|
|
}
|
|
jsonData, err := json.Marshal(convertedRequest)
|
|
if err != nil {
|
|
return err, nil
|
|
}
|
|
requestBody := bytes.NewBuffer(jsonData)
|
|
c.Request.Body = io.NopCloser(requestBody)
|
|
resp, err := adaptor.DoRequest(c, meta, requestBody)
|
|
if err != nil {
|
|
return err, nil
|
|
}
|
|
if resp != nil && resp.StatusCode != http.StatusOK {
|
|
err := relaycommon.RelayErrorHandler(resp)
|
|
return fmt.Errorf("status code %d: %s", resp.StatusCode, err.Error.Message), &err.Error
|
|
}
|
|
usage, respErr := adaptor.DoResponse(c, resp, meta)
|
|
if respErr != nil {
|
|
return fmt.Errorf("%s", respErr.Error.Message), &respErr.Error
|
|
}
|
|
if usage == nil {
|
|
return errors.New("usage is nil"), nil
|
|
}
|
|
result := w.Result()
|
|
// print result.Body
|
|
respBody, err := io.ReadAll(result.Body)
|
|
if err != nil {
|
|
return err, nil
|
|
}
|
|
common.SysLog(fmt.Sprintf("testing channel #%d, response: \n%s", channel.Id, string(respBody)))
|
|
return nil, nil
|
|
}
|
|
|
|
func buildTestRequest() *dto.GeneralOpenAIRequest {
|
|
testRequest := &dto.GeneralOpenAIRequest{
|
|
Model: "", // this will be set later
|
|
MaxTokens: 1,
|
|
Stream: false,
|
|
}
|
|
content, _ := json.Marshal("hi")
|
|
testMessage := dto.Message{
|
|
Role: "user",
|
|
Content: content,
|
|
}
|
|
testRequest.Messages = append(testRequest.Messages, testMessage)
|
|
return testRequest
|
|
}
|
|
|
|
func TestChannel(c *gin.Context) {
|
|
id, err := strconv.Atoi(c.Param("id"))
|
|
if err != nil {
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"success": false,
|
|
"message": err.Error(),
|
|
})
|
|
return
|
|
}
|
|
channel, err := model.GetChannelById(id, true)
|
|
if err != nil {
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"success": false,
|
|
"message": err.Error(),
|
|
})
|
|
return
|
|
}
|
|
testModel := c.Query("model")
|
|
tik := time.Now()
|
|
err, _ = testChannel(channel, testModel)
|
|
tok := time.Now()
|
|
milliseconds := tok.Sub(tik).Milliseconds()
|
|
go channel.UpdateResponseTime(milliseconds)
|
|
consumedTime := float64(milliseconds) / 1000.0
|
|
if err != nil {
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"success": false,
|
|
"message": err.Error(),
|
|
"time": consumedTime,
|
|
})
|
|
return
|
|
}
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"success": true,
|
|
"message": "",
|
|
"time": consumedTime,
|
|
})
|
|
return
|
|
}
|
|
|
|
var testAllChannelsLock sync.Mutex
|
|
var testAllChannelsRunning bool = false
|
|
|
|
func testAllChannels(notify bool) error {
|
|
if common.RootUserEmail == "" {
|
|
common.RootUserEmail = model.GetRootUserEmail()
|
|
}
|
|
testAllChannelsLock.Lock()
|
|
if testAllChannelsRunning {
|
|
testAllChannelsLock.Unlock()
|
|
return errors.New("测试已在运行中")
|
|
}
|
|
testAllChannelsRunning = true
|
|
testAllChannelsLock.Unlock()
|
|
channels, err := model.GetAllChannels(0, 0, true, false)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
var disableThreshold = int64(common.ChannelDisableThreshold * 1000)
|
|
if disableThreshold == 0 {
|
|
disableThreshold = 10000000 // a impossible value
|
|
}
|
|
go func() {
|
|
for _, channel := range channels {
|
|
isChannelEnabled := channel.Status == common.ChannelStatusEnabled
|
|
tik := time.Now()
|
|
err, openaiErr := testChannel(channel, "")
|
|
tok := time.Now()
|
|
milliseconds := tok.Sub(tik).Milliseconds()
|
|
|
|
ban := false
|
|
if milliseconds > disableThreshold {
|
|
err = errors.New(fmt.Sprintf("响应时间 %.2fs 超过阈值 %.2fs", float64(milliseconds)/1000.0, float64(disableThreshold)/1000.0))
|
|
ban = true
|
|
}
|
|
if openaiErr != nil {
|
|
err = errors.New(fmt.Sprintf("type %s, code %v, message %s", openaiErr.Type, openaiErr.Code, openaiErr.Message))
|
|
ban = true
|
|
}
|
|
// parse *int to bool
|
|
if channel.AutoBan != nil && *channel.AutoBan == 0 {
|
|
ban = false
|
|
}
|
|
if openaiErr != nil {
|
|
openAiErrWithStatus := dto.OpenAIErrorWithStatusCode{
|
|
StatusCode: -1,
|
|
Error: *openaiErr,
|
|
LocalError: false,
|
|
}
|
|
if isChannelEnabled && service.ShouldDisableChannel(channel.Type, &openAiErrWithStatus) && ban {
|
|
service.DisableChannel(channel.Id, channel.Name, err.Error())
|
|
}
|
|
if !isChannelEnabled && service.ShouldEnableChannel(err, openaiErr, channel.Status) {
|
|
service.EnableChannel(channel.Id, channel.Name)
|
|
}
|
|
}
|
|
channel.UpdateResponseTime(milliseconds)
|
|
time.Sleep(common.RequestInterval)
|
|
}
|
|
testAllChannelsLock.Lock()
|
|
testAllChannelsRunning = false
|
|
testAllChannelsLock.Unlock()
|
|
if notify {
|
|
err := common.SendEmail("通道测试完成", common.RootUserEmail, "通道测试完成，如果没有收到禁用通知，说明所有通道都正常")
|
|
if err != nil {
|
|
common.SysError(fmt.Sprintf("failed to send email: %s", err.Error()))
|
|
}
|
|
}
|
|
}()
|
|
return nil
|
|
}
|
|
|
|
func TestAllChannels(c *gin.Context) {
|
|
err := testAllChannels(true)
|
|
if err != nil {
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"success": false,
|
|
"message": err.Error(),
|
|
})
|
|
return
|
|
}
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"success": true,
|
|
"message": "",
|
|
})
|
|
return
|
|
}
|
|
|
|
func AutomaticallyTestChannels(frequency int) {
|
|
for {
|
|
time.Sleep(time.Duration(frequency) * time.Minute)
|
|
common.SysLog("testing all channels")
|
|
_ = testAllChannels(false)
|
|
common.SysLog("channel test finished")
|
|
}
|
|
}
|