xinyin025/new-api
1
0
Fork 0
mirror of https://github.com/linux-do/new-api.git synced 2025-09-19 00:46:37 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: support o1 channel test

Browse Source
This commit is contained in:
CalciumIon 2024-09-13 03:17:04 +08:00
parent 804aad3f37
commit cb73889353
1 changed files with 17 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
controller/channel-test.go
View File

@ -20,6 +20,7 @@ import (
"one-api/relay/constant" "one-api/relay/constant"
"one-api/service" "one-api/service"
"strconv" "strconv"
"strings"
"sync" "sync"
"time" "time"
@ -81,8 +82,7 @@ func testChannel(channel *model.Channel, testModel string) (err error, openAIErr
return fmt.Errorf("invalid api type: %d, adaptor is nil", apiType), nil return fmt.Errorf("invalid api type: %d, adaptor is nil", apiType), nil
} }
request := buildTestRequest() request := buildTestRequest(testModel)
request.Model = testModel
meta.UpstreamModelName = testModel meta.UpstreamModelName = testModel
common.SysLog(fmt.Sprintf("testing channel %d with model %s", channel.Id, testModel)) common.SysLog(fmt.Sprintf("testing channel %d with model %s", channel.Id, testModel))
@ -141,17 +141,22 @@ func testChannel(channel *model.Channel, testModel string) (err error, openAIErr
return nil, nil return nil, nil
} }
func buildTestRequest() *dto.GeneralOpenAIRequest { func buildTestRequest(model string) *dto.GeneralOpenAIRequest {
testRequest := &dto.GeneralOpenAIRequest{ testRequest := &dto.GeneralOpenAIRequest{
Model: "", // this will be set later Model: "", // this will be set later
MaxTokens: 1,
Stream: false, Stream: false,
} }
if strings.HasPrefix(model, "o1-") {
testRequest.MaxCompletionTokens = 1
} else {
testRequest.MaxTokens = 1
}
content, _ := json.Marshal("hi") content, _ := json.Marshal("hi")
testMessage := dto.Message{ testMessage := dto.Message{
Role: "user", Role: "user",
Content: content, Content: content,
} }
testRequest.Model = model
testRequest.Messages = append(testRequest.Messages, testMessage) testRequest.Messages = append(testRequest.Messages, testMessage)
return testRequest return testRequest
} }
@ -226,26 +231,22 @@ func testAllChannels(notify bool) error {
tok := time.Now() tok := time.Now()
milliseconds := tok.Sub(tik).Milliseconds() milliseconds := tok.Sub(tik).Milliseconds()
ban := false shouldBanChannel := false
if milliseconds > disableThreshold {
err = errors.New(fmt.Sprintf("响应时间 %.2fs 超过阈值 %.2fs", float64(milliseconds)/1000.0, float64(disableThreshold)/1000.0))
ban = true
}
// request error disables the channel // request error disables the channel
if openaiWithStatusErr != nil { if openaiWithStatusErr != nil {
oaiErr := openaiWithStatusErr.Error oaiErr := openaiWithStatusErr.Error
err = errors.New(fmt.Sprintf("type %s, httpCode %d, code %v, message %s", oaiErr.Type, openaiWithStatusErr.StatusCode, oaiErr.Code, oaiErr.Message)) err = errors.New(fmt.Sprintf("type %s, httpCode %d, code %v, message %s", oaiErr.Type, openaiWithStatusErr.StatusCode, oaiErr.Code, oaiErr.Message))
ban = service.ShouldDisableChannel(channel.Type, openaiWithStatusErr) shouldBanChannel = service.ShouldDisableChannel(channel.Type, openaiWithStatusErr)
} }
// parse *int to bool if milliseconds > disableThreshold {
if !channel.GetAutoBan() { err = errors.New(fmt.Sprintf("响应时间 %.2fs 超过阈值 %.2fs", float64(milliseconds)/1000.0, float64(disableThreshold)/1000.0))
ban = false shouldBanChannel = true
} }
// disable channel // disable channel
if ban && isChannelEnabled { if isChannelEnabled && shouldBanChannel && channel.GetAutoBan() {
service.DisableChannel(channel.Id, channel.Name, err.Error()) service.DisableChannel(channel.Id, channel.Name, err.Error())
} }