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

feat: 初步兼容midjourney-proxy-plus

Browse Source
This commit is contained in:
CaIon
2024-03-13 15:37:01 +08:00
parent 95d8059c90
commit 37c0c8ebdd
14 changed files with 246 additions and 153 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
controller/channel-billing.go
View File

@@ -214,8 +214,8 @@ func updateChannelBalance(channel *model.Channel) (float64, error) {
return 0, errors.New("尚未实现")
case common.ChannelTypeCustom:
baseURL = channel.GetBaseURL()
case common.ChannelTypeOpenAISB:
return updateChannelOpenAISBBalance(channel)
//case common.ChannelTypeOpenAISB:
// return updateChannelOpenAISBBalance(channel)
case common.ChannelTypeAIProxy:
return updateChannelAIProxyBalance(channel)
case common.ChannelTypeAPI2GPT:

18
controller/midjourney.go
View File

@@ -10,8 +10,8 @@ import (
"log"
"net/http"
"one-api/common"
"one-api/dto"
"one-api/model"
relay2 "one-api/relay"
"one-api/service"
"strconv"
"strings"
@@ -75,11 +75,11 @@ import (
responseBody, err := io.ReadAll(resp.Body)
resp.Body.Close()
log.Printf("responseBody: %s", string(responseBody))
var responseItem Midjourney
var responseItem MidjourneyDto
// err = json.NewDecoder(resp.Body).Decode(&responseItem)
err = json.Unmarshal(responseBody, &responseItem)
if err != nil {
if strings.Contains(err.Error(), "cannot unmarshal number into Go struct field Midjourney.status of type string") {
if strings.Contains(err.Error(), "cannot unmarshal number into Go struct field MidjourneyDto.status of type string") {
var responseWithoutStatus MidjourneyWithoutStatus
var responseStatus MidjourneyStatus
err1 := json.Unmarshal(responseBody, &responseWithoutStatus)
@@ -228,12 +228,16 @@ func UpdateMidjourneyTaskBulk() {
common.LogError(ctx, fmt.Sprintf("Get Task Do req error: %v", err))
continue
}
if resp.StatusCode != http.StatusOK {
common.LogError(ctx, fmt.Sprintf("Get Task status code: %d", resp.StatusCode))
continue
}
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
common.LogError(ctx, fmt.Sprintf("Get Task parse body error: %v", err))
continue
}
var responseItems []relay2.Midjourney
var responseItems []dto.MidjourneyDto
err = json.Unmarshal(responseBody, &responseItems)
if err != nil {
common.LogError(ctx, fmt.Sprintf("Get Task parse body error2: %v, body: %s", err, string(responseBody)))
@@ -259,6 +263,10 @@ func UpdateMidjourneyTaskBulk() {
task.ImageUrl = responseItem.ImageUrl
task.Status = responseItem.Status
task.FailReason = responseItem.FailReason
if responseItem.Buttons != nil {
buttonStr, _ := json.Marshal(responseItem.Buttons)
task.Buttons = string(buttonStr)
}
if task.Progress != "100%" && responseItem.FailReason != "" {
common.LogInfo(ctx, task.MjId+" 构建失败，"+task.FailReason)
task.Progress = "100%"
@@ -286,7 +294,7 @@ func UpdateMidjourneyTaskBulk() {
}
}
func checkMjTaskNeedUpdate(oldTask *model.Midjourney, newTask relay2.Midjourney) bool {
func checkMjTaskNeedUpdate(oldTask *model.Midjourney, newTask dto.MidjourneyDto) bool {
if oldTask.Code != 1 {
return true
}

37
controller/relay.go
View File

@@ -62,7 +62,13 @@ func Relay(c *gin.Context) {
func RelayMidjourney(c *gin.Context) {
relayMode := relayconstant.RelayModeUnknown
if strings.HasPrefix(c.Request.URL.Path, "/mj/submit/imagine") {
if strings.HasPrefix(c.Request.URL.Path, "/mj/submit/action") {
// midjourney plus
relayMode = relayconstant.RelayModeMidjourneyAction
} else if strings.HasPrefix(c.Request.URL.Path, "/mj/submit/modal") {
// midjourney plus
relayMode = relayconstant.RelayModeMidjourneyModal
} else if strings.HasPrefix(c.Request.URL.Path, "/mj/submit/imagine") {
relayMode = relayconstant.RelayModeMidjourneyImagine
} else if strings.HasPrefix(c.Request.URL.Path, "/mj/submit/blend") {
relayMode = relayconstant.RelayModeMidjourneyBlend
@@ -86,35 +92,24 @@ func RelayMidjourney(c *gin.Context) {
err = relay.RelayMidjourneyNotify(c)
case relayconstant.RelayModeMidjourneyTaskFetch, relayconstant.RelayModeMidjourneyTaskFetchByCondition:
err = relay.RelayMidjourneyTask(c, relayMode)
//case relayconstant.RelayModeMidjourneyModal:
// err = relay.RelayMidjournneyModal(c)
default:
err = relay.RelayMidjourneySubmit(c, relayMode)
}
//err = relayMidjourneySubmit(c, relayMode)
log.Println(err)
if err != nil {
retryTimesStr := c.Query("retry")
retryTimes, _ := strconv.Atoi(retryTimesStr)
if retryTimesStr == "" {
retryTimes = common.RetryTimes
}
if retryTimes > 0 {
c.Redirect(http.StatusTemporaryRedirect, fmt.Sprintf("%s?retry=%d", c.Request.URL.Path, retryTimes-1))
} else {
if err.Code == 30 {
err.Result = "当前分组负载已饱和，请稍后再试，或升级账户以提升服务质量。"
}
c.JSON(429, gin.H{
"error": fmt.Sprintf("%s %s", err.Description, err.Result),
"type": "upstream_error",
})
if err.Code == 30 {
err.Result = "当前分组负载已饱和，请稍后再试，或升级账户以提升服务质量。"
}
c.JSON(429, gin.H{
"error": fmt.Sprintf("%s %s", err.Description, err.Result),
"type": "upstream_error",
"code": err.Code,
})
channelId := c.GetInt("channel_id")
common.SysError(fmt.Sprintf("relay error (channel #%d): %s", channelId, fmt.Sprintf("%s %s", err.Description, err.Result)))
//if shouldDisableChannel(&err.Error) {
// channelId := c.GetInt("channel_id")
// channelName := c.GetString("channel_name")
// disableChannel(channelId, channelName, err.Result)
//};''''''''''''''''''''''''''''''''
}
}