refactor: refactor relay part (#935)

This commit is contained in:
JustSong
2024-01-14 19:21:03 +08:00
committed by GitHub
parent f2c51a494c
commit eed9f5fdf0
37 changed files with 1527 additions and 1437 deletions

View File

@@ -9,6 +9,8 @@ import (
"net/http"
"one-api/common"
"one-api/model"
"one-api/relay/channel/openai"
"one-api/relay/util"
"strconv"
"sync"
"time"
@@ -16,7 +18,7 @@ import (
"github.com/gin-gonic/gin"
)
func testChannel(channel *model.Channel, request ChatRequest) (err error, openaiErr *OpenAIError) {
func testChannel(channel *model.Channel, request openai.ChatRequest) (err error, openaiErr *openai.Error) {
switch channel.Type {
case common.ChannelTypePaLM:
fallthrough
@@ -46,13 +48,13 @@ func testChannel(channel *model.Channel, request ChatRequest) (err error, openai
}
requestURL := common.ChannelBaseURLs[channel.Type]
if channel.Type == common.ChannelTypeAzure {
requestURL = getFullRequestURL(channel.GetBaseURL(), fmt.Sprintf("/openai/deployments/%s/chat/completions?api-version=2023-03-15-preview", request.Model), channel.Type)
requestURL = util.GetFullRequestURL(channel.GetBaseURL(), fmt.Sprintf("/openai/deployments/%s/chat/completions?api-version=2023-03-15-preview", request.Model), channel.Type)
} else {
if baseURL := channel.GetBaseURL(); len(baseURL) > 0 {
requestURL = baseURL
}
requestURL = getFullRequestURL(requestURL, "/v1/chat/completions", channel.Type)
requestURL = util.GetFullRequestURL(requestURL, "/v1/chat/completions", channel.Type)
}
jsonData, err := json.Marshal(request)
if err != nil {
@@ -68,12 +70,12 @@ func testChannel(channel *model.Channel, request ChatRequest) (err error, openai
req.Header.Set("Authorization", "Bearer "+channel.Key)
}
req.Header.Set("Content-Type", "application/json")
resp, err := httpClient.Do(req)
resp, err := util.HTTPClient.Do(req)
if err != nil {
return err, nil
}
defer resp.Body.Close()
var response TextResponse
var response openai.SlimTextResponse
body, err := io.ReadAll(resp.Body)
if err != nil {
return err, nil
@@ -91,12 +93,12 @@ func testChannel(channel *model.Channel, request ChatRequest) (err error, openai
return nil, nil
}
func buildTestRequest() *ChatRequest {
testRequest := &ChatRequest{
func buildTestRequest() *openai.ChatRequest {
testRequest := &openai.ChatRequest{
Model: "", // this will be set later
MaxTokens: 1,
}
testMessage := Message{
testMessage := openai.Message{
Role: "user",
Content: "hi",
}
@@ -204,10 +206,10 @@ func testAllChannels(notify bool) error {
err = errors.New(fmt.Sprintf("响应时间 %.2fs 超过阈值 %.2fs", float64(milliseconds)/1000.0, float64(disableThreshold)/1000.0))
disableChannel(channel.Id, channel.Name, err.Error())
}
if isChannelEnabled && shouldDisableChannel(openaiErr, -1) {
if isChannelEnabled && util.ShouldDisableChannel(openaiErr, -1) {
disableChannel(channel.Id, channel.Name, err.Error())
}
if !isChannelEnabled && shouldEnableChannel(err, openaiErr) {
if !isChannelEnabled && util.ShouldEnableChannel(err, openaiErr) {
enableChannel(channel.Id, channel.Name)
}
channel.UpdateResponseTime(milliseconds)