mirror of
https://github.com/linux-do/new-api.git
synced 2025-09-18 00:16:37 +08:00
287 lines
8.5 KiB
Go
287 lines
8.5 KiB
Go
package controller
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
"github.com/gin-gonic/gin"
|
|
"io"
|
|
"log"
|
|
"net/http"
|
|
"one-api/common"
|
|
"one-api/dto"
|
|
"one-api/middleware"
|
|
"one-api/model"
|
|
"one-api/relay"
|
|
"one-api/relay/constant"
|
|
relayconstant "one-api/relay/constant"
|
|
"one-api/service"
|
|
"strings"
|
|
)
|
|
|
|
func relayHandler(c *gin.Context, relayMode int) *dto.OpenAIErrorWithStatusCode {
|
|
var err *dto.OpenAIErrorWithStatusCode
|
|
switch relayMode {
|
|
case relayconstant.RelayModeImagesGenerations:
|
|
err = relay.ImageHelper(c, relayMode)
|
|
case relayconstant.RelayModeAudioSpeech:
|
|
fallthrough
|
|
case relayconstant.RelayModeAudioTranslation:
|
|
fallthrough
|
|
case relayconstant.RelayModeAudioTranscription:
|
|
err = relay.AudioHelper(c)
|
|
case relayconstant.RelayModeRerank:
|
|
err = relay.RerankHelper(c, relayMode)
|
|
default:
|
|
err = relay.TextHelper(c)
|
|
}
|
|
return err
|
|
}
|
|
|
|
func Relay(c *gin.Context) {
|
|
relayMode := constant.Path2RelayMode(c.Request.URL.Path)
|
|
retryTimes := common.RetryTimes
|
|
requestId := c.GetString(common.RequestIdKey)
|
|
channelId := c.GetInt("channel_id")
|
|
channelType := c.GetInt("channel_type")
|
|
group := c.GetString("group")
|
|
originalModel := c.GetString("original_model")
|
|
openaiErr := relayHandler(c, relayMode)
|
|
c.Set("use_channel", []string{fmt.Sprintf("%d", channelId)})
|
|
if openaiErr != nil {
|
|
go processChannelError(c, channelId, channelType, openaiErr)
|
|
} else {
|
|
retryTimes = 0
|
|
}
|
|
for i := 0; shouldRetry(c, channelId, openaiErr, retryTimes) && i < retryTimes; i++ {
|
|
channel, err := model.CacheGetRandomSatisfiedChannel(group, originalModel, i)
|
|
if err != nil {
|
|
common.LogError(c.Request.Context(), fmt.Sprintf("CacheGetRandomSatisfiedChannel failed: %s", err.Error()))
|
|
break
|
|
}
|
|
channelId = channel.Id
|
|
useChannel := c.GetStringSlice("use_channel")
|
|
useChannel = append(useChannel, fmt.Sprintf("%d", channelId))
|
|
c.Set("use_channel", useChannel)
|
|
common.LogInfo(c.Request.Context(), fmt.Sprintf("using channel #%d to retry (remain times %d)", channel.Id, i))
|
|
middleware.SetupContextForSelectedChannel(c, channel, originalModel)
|
|
|
|
requestBody, err := common.GetRequestBody(c)
|
|
c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody))
|
|
openaiErr = relayHandler(c, relayMode)
|
|
if openaiErr != nil {
|
|
go processChannelError(c, channelId, channel.Type, openaiErr)
|
|
}
|
|
}
|
|
useChannel := c.GetStringSlice("use_channel")
|
|
if len(useChannel) > 1 {
|
|
retryLogStr := fmt.Sprintf("重试：%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
|
|
common.LogInfo(c.Request.Context(), retryLogStr)
|
|
}
|
|
|
|
if openaiErr != nil {
|
|
if openaiErr.StatusCode == http.StatusTooManyRequests {
|
|
openaiErr.Error.Message = "当前分组上游负载已饱和，请稍后再试"
|
|
}
|
|
openaiErr.Error.Message = common.MessageWithRequestId(openaiErr.Error.Message, requestId)
|
|
c.JSON(openaiErr.StatusCode, gin.H{
|
|
"error": openaiErr.Error,
|
|
})
|
|
}
|
|
}
|
|
|
|
func shouldRetry(c *gin.Context, channelId int, openaiErr *dto.OpenAIErrorWithStatusCode, retryTimes int) bool {
|
|
if openaiErr == nil {
|
|
return false
|
|
}
|
|
if retryTimes <= 0 {
|
|
return false
|
|
}
|
|
if _, ok := c.Get("specific_channel_id"); ok {
|
|
return false
|
|
}
|
|
if openaiErr.StatusCode == http.StatusTooManyRequests {
|
|
return true
|
|
}
|
|
if openaiErr.StatusCode == 307 {
|
|
return true
|
|
}
|
|
if openaiErr.StatusCode/100 == 5 {
|
|
// 超时不重试
|
|
if openaiErr.StatusCode == 504 || openaiErr.StatusCode == 524 {
|
|
return false
|
|
}
|
|
return true
|
|
}
|
|
if openaiErr.StatusCode == http.StatusBadRequest {
|
|
return false
|
|
}
|
|
if openaiErr.StatusCode == 408 {
|
|
// azure处理超时不重试
|
|
return false
|
|
}
|
|
if openaiErr.LocalError {
|
|
return false
|
|
}
|
|
if openaiErr.StatusCode/100 == 2 {
|
|
return false
|
|
}
|
|
return true
|
|
}
|
|
|
|
func processChannelError(c *gin.Context, channelId int, channelType int, err *dto.OpenAIErrorWithStatusCode) {
|
|
autoBan := c.GetBool("auto_ban")
|
|
common.LogError(c.Request.Context(), fmt.Sprintf("relay error (channel #%d, status code: %d): %s", channelId, err.StatusCode, err.Error.Message))
|
|
if service.ShouldDisableChannel(channelType, err) && autoBan {
|
|
channelName := c.GetString("channel_name")
|
|
service.DisableChannel(channelId, channelName, err.Error.Message)
|
|
}
|
|
}
|
|
|
|
func RelayMidjourney(c *gin.Context) {
|
|
relayMode := c.GetInt("relay_mode")
|
|
var err *dto.MidjourneyResponse
|
|
switch relayMode {
|
|
case relayconstant.RelayModeMidjourneyNotify:
|
|
err = relay.RelayMidjourneyNotify(c)
|
|
case relayconstant.RelayModeMidjourneyTaskFetch, relayconstant.RelayModeMidjourneyTaskFetchByCondition:
|
|
err = relay.RelayMidjourneyTask(c, relayMode)
|
|
case relayconstant.RelayModeMidjourneyTaskImageSeed:
|
|
err = relay.RelayMidjourneyTaskImageSeed(c)
|
|
case relayconstant.RelayModeSwapFace:
|
|
err = relay.RelaySwapFace(c)
|
|
default:
|
|
err = relay.RelayMidjourneySubmit(c, relayMode)
|
|
}
|
|
//err = relayMidjourneySubmit(c, relayMode)
|
|
log.Println(err)
|
|
if err != nil {
|
|
statusCode := http.StatusBadRequest
|
|
if err.Code == 30 {
|
|
err.Result = "当前分组负载已饱和，请稍后再试，或升级账户以提升服务质量。"
|
|
statusCode = http.StatusTooManyRequests
|
|
}
|
|
c.JSON(statusCode, gin.H{
|
|
"description": fmt.Sprintf("%s %s", err.Description, err.Result),
|
|
"type": "upstream_error",
|
|
"code": err.Code,
|
|
})
|
|
channelId := c.GetInt("channel_id")
|
|
common.LogError(c, fmt.Sprintf("relay error (channel #%d, status code %d): %s", channelId, statusCode, fmt.Sprintf("%s %s", err.Description, err.Result)))
|
|
}
|
|
}
|
|
|
|
func RelayNotImplemented(c *gin.Context) {
|
|
err := dto.OpenAIError{
|
|
Message: "API not implemented",
|
|
Type: "new_api_error",
|
|
Param: "",
|
|
Code: "api_not_implemented",
|
|
}
|
|
c.JSON(http.StatusNotImplemented, gin.H{
|
|
"error": err,
|
|
})
|
|
}
|
|
|
|
func RelayNotFound(c *gin.Context) {
|
|
err := dto.OpenAIError{
|
|
Message: fmt.Sprintf("Invalid URL (%s %s)", c.Request.Method, c.Request.URL.Path),
|
|
Type: "invalid_request_error",
|
|
Param: "",
|
|
Code: "",
|
|
}
|
|
c.JSON(http.StatusNotFound, gin.H{
|
|
"error": err,
|
|
})
|
|
}
|
|
|
|
func RelayTask(c *gin.Context) {
|
|
retryTimes := common.RetryTimes
|
|
channelId := c.GetInt("channel_id")
|
|
relayMode := c.GetInt("relay_mode")
|
|
group := c.GetString("group")
|
|
originalModel := c.GetString("original_model")
|
|
c.Set("use_channel", []string{fmt.Sprintf("%d", channelId)})
|
|
taskErr := taskRelayHandler(c, relayMode)
|
|
if taskErr == nil {
|
|
retryTimes = 0
|
|
}
|
|
for i := 0; shouldRetryTaskRelay(c, channelId, taskErr, retryTimes) && i < retryTimes; i++ {
|
|
channel, err := model.CacheGetRandomSatisfiedChannel(group, originalModel, i)
|
|
if err != nil {
|
|
common.LogError(c.Request.Context(), fmt.Sprintf("CacheGetRandomSatisfiedChannel failed: %s", err.Error()))
|
|
break
|
|
}
|
|
channelId = channel.Id
|
|
useChannel := c.GetStringSlice("use_channel")
|
|
useChannel = append(useChannel, fmt.Sprintf("%d", channelId))
|
|
c.Set("use_channel", useChannel)
|
|
common.LogInfo(c.Request.Context(), fmt.Sprintf("using channel #%d to retry (remain times %d)", channel.Id, i))
|
|
middleware.SetupContextForSelectedChannel(c, channel, originalModel)
|
|
|
|
requestBody, err := common.GetRequestBody(c)
|
|
c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody))
|
|
taskErr = taskRelayHandler(c, relayMode)
|
|
}
|
|
useChannel := c.GetStringSlice("use_channel")
|
|
if len(useChannel) > 1 {
|
|
retryLogStr := fmt.Sprintf("重试：%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
|
|
common.LogInfo(c.Request.Context(), retryLogStr)
|
|
}
|
|
if taskErr != nil {
|
|
if taskErr.StatusCode == http.StatusTooManyRequests {
|
|
taskErr.Message = "当前分组上游负载已饱和，请稍后再试"
|
|
}
|
|
c.JSON(taskErr.StatusCode, taskErr)
|
|
}
|
|
}
|
|
|
|
func taskRelayHandler(c *gin.Context, relayMode int) *dto.TaskError {
|
|
var err *dto.TaskError
|
|
switch relayMode {
|
|
case relayconstant.RelayModeSunoFetch, relayconstant.RelayModeSunoFetchByID:
|
|
err = relay.RelayTaskFetch(c, relayMode)
|
|
default:
|
|
err = relay.RelayTaskSubmit(c, relayMode)
|
|
}
|
|
return err
|
|
}
|
|
|
|
func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *dto.TaskError, retryTimes int) bool {
|
|
if taskErr == nil {
|
|
return false
|
|
}
|
|
if retryTimes <= 0 {
|
|
return false
|
|
}
|
|
if _, ok := c.Get("specific_channel_id"); ok {
|
|
return false
|
|
}
|
|
if taskErr.StatusCode == http.StatusTooManyRequests {
|
|
return true
|
|
}
|
|
if taskErr.StatusCode == 307 {
|
|
return true
|
|
}
|
|
if taskErr.StatusCode/100 == 5 {
|
|
// 超时不重试
|
|
if taskErr.StatusCode == 504 || taskErr.StatusCode == 524 {
|
|
return false
|
|
}
|
|
return true
|
|
}
|
|
if taskErr.StatusCode == http.StatusBadRequest {
|
|
return false
|
|
}
|
|
if taskErr.StatusCode == 408 {
|
|
// azure处理超时不重试
|
|
return false
|
|
}
|
|
if taskErr.LocalError {
|
|
return false
|
|
}
|
|
if taskErr.StatusCode/100 == 2 {
|
|
return false
|
|
}
|
|
return true
|
|
}
|