cneicy
a5f7b8e586
Publish Docker image (Multi-arch) / Build & push (amd64) (push) Failing after 1m40s
Release (Linux, macOS, Windows) / Linux Release (push) Failing after 42m16s
Publish Docker image (Multi-arch) / Build & push (arm64) (push) Canceled after 0s
Release (Linux, macOS, Windows) / macOS Release (push) Canceled after 0s
Release (Linux, macOS, Windows) / Windows Release (push) Canceled after 0s
Publish Docker image (Multi-arch) / Create multi-arch manifests (push) Canceled after 0s
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
393 lines
14 KiB
Go
393 lines
14 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/dto"
|
|
"github.com/QuantumNous/new-api/logger"
|
|
"github.com/QuantumNous/new-api/model"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
|
"github.com/QuantumNous/new-api/setting"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func CovertMjpActionToModelName(mjAction string) string {
|
|
modelName := "mj_" + strings.ToLower(mjAction)
|
|
if mjAction == constant.MjActionSwapFace {
|
|
modelName = "swap_face"
|
|
}
|
|
return modelName
|
|
}
|
|
|
|
// PrepareMidjourneyTaskBilling sets the durable refund marker before the task is inserted.
|
|
func PrepareMidjourneyTaskBilling(relayInfo *relaycommon.RelayInfo, task *model.Midjourney, quota int, shouldBill bool) (bool, error) {
|
|
if task == nil {
|
|
return false, errors.New("Midjourney task is nil")
|
|
}
|
|
task.Quota = 0
|
|
task.TokenId = 0
|
|
task.BillingChannelId = 0
|
|
if !shouldBill {
|
|
return false, nil
|
|
}
|
|
if relayInfo == nil {
|
|
return false, errors.New("relay info is nil")
|
|
}
|
|
if quota < 0 {
|
|
return false, errors.New("quota cannot be negative")
|
|
}
|
|
if relayInfo.BillingSource == BillingSourceSubscription {
|
|
return false, errors.New("legacy Midjourney billing does not support subscriptions")
|
|
}
|
|
|
|
task.Quota = quota
|
|
task.BillingChannelId = task.ChannelId
|
|
if relayInfo.ChannelMeta != nil && relayInfo.ChannelId > 0 {
|
|
task.BillingChannelId = relayInfo.ChannelId
|
|
}
|
|
return true, nil
|
|
}
|
|
|
|
// SettleMidjourneyTaskBilling charges a persisted legacy task and records the applied stages.
|
|
func SettleMidjourneyTaskBilling(relayInfo *relaycommon.RelayInfo, task *model.Midjourney, prepared bool) (bool, error) {
|
|
if !prepared {
|
|
return false, nil
|
|
}
|
|
if relayInfo == nil {
|
|
return false, errors.New("relay info is nil")
|
|
}
|
|
if task == nil || task.Id == 0 {
|
|
return false, errors.New("Midjourney task must be persisted before billing")
|
|
}
|
|
|
|
result, billingErr := postConsumeQuotaWithResult(relayInfo, task.Quota, 0, true)
|
|
if !result.FundingApplied {
|
|
task.Quota = 0
|
|
task.TokenId = 0
|
|
task.BillingChannelId = 0
|
|
if updateErr := task.UpdateBillingState(); updateErr != nil {
|
|
return false, errors.Join(billingErr, fmt.Errorf("clear Midjourney billing state: %w", updateErr))
|
|
}
|
|
return false, billingErr
|
|
}
|
|
|
|
task.TokenId = 0
|
|
if result.TokenApplied {
|
|
task.TokenId = relayInfo.TokenId
|
|
}
|
|
if updateErr := task.UpdateBillingState(); updateErr != nil {
|
|
return true, errors.Join(billingErr, fmt.Errorf("update Midjourney billing state: %w", updateErr))
|
|
}
|
|
return true, billingErr
|
|
}
|
|
|
|
// RefundMidjourneyQuota reverses every accounting element recorded for a billed legacy task.
|
|
func RefundMidjourneyQuota(ctx context.Context, task *model.Midjourney, reason string) bool {
|
|
quota := task.Quota
|
|
if quota == 0 {
|
|
return true
|
|
}
|
|
|
|
if err := model.IncreaseUserQuota(task.UserId, quota, false); err != nil {
|
|
logger.LogWarn(ctx, fmt.Sprintf("退还 Midjourney 用户额度失败 task %s: %s", task.MjId, err.Error()))
|
|
return false
|
|
}
|
|
|
|
if task.TokenId > 0 {
|
|
tokenKey := resolveTokenKey(ctx, task.TokenId, task.MjId)
|
|
if tokenKey != "" {
|
|
if err := model.IncreaseTokenQuota(task.TokenId, tokenKey, quota); err != nil {
|
|
logger.LogWarn(ctx, fmt.Sprintf("退还 Midjourney 令牌额度失败 task %s: %s", task.MjId, err.Error()))
|
|
}
|
|
}
|
|
}
|
|
|
|
billingChannelId := task.GetBillingChannelId()
|
|
model.UpdateUserUsedQuota(task.UserId, -quota)
|
|
model.UpdateChannelUsedQuota(billingChannelId, -quota)
|
|
other := model.NewLogOther()
|
|
other.SetPublic("task_id", task.MjId)
|
|
other.SetPublic("reason", reason)
|
|
model.RecordTaskBillingLog(model.RecordTaskBillingLogParams{
|
|
UserId: task.UserId,
|
|
LogType: model.LogTypeRefund,
|
|
Content: "",
|
|
ChannelId: billingChannelId,
|
|
ModelName: CovertMjpActionToModelName(task.Action),
|
|
Quota: quota,
|
|
TokenId: task.TokenId,
|
|
Other: other,
|
|
})
|
|
|
|
task.Quota = 0
|
|
if err := task.UpdateBillingState(); err != nil {
|
|
logger.LogError(ctx, fmt.Sprintf("Midjourney 退款成功但清除 quota 失败 task %s: %s", task.MjId, err.Error()))
|
|
}
|
|
return true
|
|
}
|
|
|
|
func GetMjRequestModel(relayMode int, midjRequest *dto.MidjourneyRequest) (string, *dto.MidjourneyResponse, bool) {
|
|
action := ""
|
|
if relayMode == relayconstant.RelayModeMidjourneyAction {
|
|
// plus request
|
|
err := CoverPlusActionToNormalAction(midjRequest)
|
|
if err != nil {
|
|
return "", err, false
|
|
}
|
|
action = midjRequest.Action
|
|
} else {
|
|
switch relayMode {
|
|
case relayconstant.RelayModeMidjourneyImagine:
|
|
action = constant.MjActionImagine
|
|
case relayconstant.RelayModeMidjourneyVideo:
|
|
action = constant.MjActionVideo
|
|
case relayconstant.RelayModeMidjourneyEdits:
|
|
action = constant.MjActionEdits
|
|
case relayconstant.RelayModeMidjourneyDescribe:
|
|
action = constant.MjActionDescribe
|
|
case relayconstant.RelayModeMidjourneyBlend:
|
|
action = constant.MjActionBlend
|
|
case relayconstant.RelayModeMidjourneyShorten:
|
|
action = constant.MjActionShorten
|
|
case relayconstant.RelayModeMidjourneyChange:
|
|
action = midjRequest.Action
|
|
case relayconstant.RelayModeMidjourneyModal:
|
|
action = constant.MjActionModal
|
|
case relayconstant.RelayModeSwapFace:
|
|
action = constant.MjActionSwapFace
|
|
case relayconstant.RelayModeMidjourneyUpload:
|
|
action = constant.MjActionUpload
|
|
case relayconstant.RelayModeMidjourneySimpleChange:
|
|
params := ConvertSimpleChangeParams(midjRequest.Content)
|
|
if params == nil {
|
|
return "", MidjourneyErrorWrapper(constant.MjRequestError, "invalid_request"), false
|
|
}
|
|
action = params.Action
|
|
case relayconstant.RelayModeMidjourneyTaskFetch, relayconstant.RelayModeMidjourneyTaskFetchByCondition, relayconstant.RelayModeMidjourneyNotify:
|
|
return "", nil, true
|
|
default:
|
|
return "", MidjourneyErrorWrapper(constant.MjRequestError, "unknown_relay_action"), false
|
|
}
|
|
}
|
|
modelName := CovertMjpActionToModelName(action)
|
|
return modelName, nil, true
|
|
}
|
|
|
|
func CoverPlusActionToNormalAction(midjRequest *dto.MidjourneyRequest) *dto.MidjourneyResponse {
|
|
// "customId": "MJ::JOB::upsample::2::3dbbd469-36af-4a0f-8f02-df6c579e7011"
|
|
customId := midjRequest.CustomId
|
|
if customId == "" {
|
|
return MidjourneyErrorWrapper(constant.MjRequestError, "custom_id_is_required")
|
|
}
|
|
splits := strings.Split(customId, "::")
|
|
var action string
|
|
if splits[1] == "JOB" {
|
|
action = splits[2]
|
|
} else {
|
|
action = splits[1]
|
|
}
|
|
|
|
if action == "" {
|
|
return MidjourneyErrorWrapper(constant.MjRequestError, "unknown_action")
|
|
}
|
|
if strings.Contains(action, "upsample") {
|
|
index, err := strconv.Atoi(splits[3])
|
|
if err != nil {
|
|
return MidjourneyErrorWrapper(constant.MjRequestError, "index_parse_failed")
|
|
}
|
|
midjRequest.Index = index
|
|
midjRequest.Action = constant.MjActionUpscale
|
|
} else if strings.Contains(action, "variation") {
|
|
midjRequest.Index = 1
|
|
if action == "variation" {
|
|
index, err := strconv.Atoi(splits[3])
|
|
if err != nil {
|
|
return MidjourneyErrorWrapper(constant.MjRequestError, "index_parse_failed")
|
|
}
|
|
midjRequest.Index = index
|
|
midjRequest.Action = constant.MjActionVariation
|
|
} else if action == "low_variation" {
|
|
midjRequest.Action = constant.MjActionLowVariation
|
|
} else if action == "high_variation" {
|
|
midjRequest.Action = constant.MjActionHighVariation
|
|
}
|
|
} else if strings.Contains(action, "pan") {
|
|
midjRequest.Action = constant.MjActionPan
|
|
midjRequest.Index = 1
|
|
} else if strings.Contains(action, "reroll") {
|
|
midjRequest.Action = constant.MjActionReRoll
|
|
midjRequest.Index = 1
|
|
} else if action == "Outpaint" {
|
|
midjRequest.Action = constant.MjActionZoom
|
|
midjRequest.Index = 1
|
|
} else if action == "CustomZoom" {
|
|
midjRequest.Action = constant.MjActionCustomZoom
|
|
midjRequest.Index = 1
|
|
} else if action == "Inpaint" {
|
|
midjRequest.Action = constant.MjActionInPaint
|
|
midjRequest.Index = 1
|
|
} else {
|
|
return MidjourneyErrorWrapper(constant.MjRequestError, "unknown_action:"+customId)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func ConvertSimpleChangeParams(content string) *dto.MidjourneyRequest {
|
|
split := strings.Split(content, " ")
|
|
if len(split) != 2 {
|
|
return nil
|
|
}
|
|
|
|
action := strings.ToLower(split[1])
|
|
changeParams := &dto.MidjourneyRequest{}
|
|
changeParams.TaskId = split[0]
|
|
|
|
if action[0] == 'u' {
|
|
changeParams.Action = "UPSCALE"
|
|
} else if action[0] == 'v' {
|
|
changeParams.Action = "VARIATION"
|
|
} else if action == "r" {
|
|
changeParams.Action = "REROLL"
|
|
return changeParams
|
|
} else {
|
|
return nil
|
|
}
|
|
|
|
index, err := strconv.Atoi(action[1:2])
|
|
if err != nil || index < 1 || index > 4 {
|
|
return nil
|
|
}
|
|
changeParams.Index = index
|
|
return changeParams
|
|
}
|
|
|
|
// RecordMidjourneyPolicyResponse distinguishes accepted tasks from errors inside
|
|
// HTTP 200 responses. Submissions are single-attempt: an ambiguous transport
|
|
// failure must never create a duplicate task on another channel.
|
|
func RecordMidjourneyPolicyResponse(c *gin.Context, response *dto.MidjourneyResponseWithStatusCode, requestErr error) bool {
|
|
if response == nil {
|
|
response = MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "empty_response", http.StatusBadGateway)
|
|
}
|
|
accepted := requestErr == nil && response.StatusCode == http.StatusOK && (response.Response.Code == 1 || response.Response.Code == 21 || response.Response.Code == 22)
|
|
if accepted {
|
|
properties, _ := response.Response.Properties.(map[string]any)
|
|
if properties["status"] != "FAILURE" {
|
|
return true
|
|
}
|
|
}
|
|
state := RequestPolicy(c)
|
|
event := PolicyEvent{ChannelID: c.GetInt("channel_id"), Status: response.StatusCode, ErrorCode: strconv.Itoa(response.Response.Code), ErrorSource: "upstream", Decision: PolicyDecision{Action: "failure", Reason: "upstream_failure", Source: "upstream"}}
|
|
state.AddEvent(event)
|
|
event.Decision, event.Health = PolicyDecision{Action: "stop", Reason: "non_retryable_error", Source: "system"}, "unchanged"
|
|
if accepted {
|
|
event.Decision.Reason = "task_accepted"
|
|
}
|
|
state.AddEvent(event)
|
|
return false
|
|
}
|
|
|
|
func DoMidjourneyHttpRequest(c *gin.Context, timeout time.Duration, fullRequestURL string) (*dto.MidjourneyResponseWithStatusCode, []byte, error) {
|
|
var nullBytes []byte
|
|
//var requestBody io.Reader
|
|
//requestBody = c.Request.Body
|
|
// read request body to json, delete accountFilter and notifyHook
|
|
var mapResult map[string]any
|
|
// if get request, no need to read request body
|
|
if c.Request.Method != "GET" {
|
|
err := common.DecodeJson(c.Request.Body, &mapResult)
|
|
if err != nil {
|
|
return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "read_request_body_failed", http.StatusInternalServerError), nullBytes, err
|
|
}
|
|
if !setting.MjAccountFilterEnabled {
|
|
delete(mapResult, "accountFilter")
|
|
}
|
|
if !setting.MjNotifyEnabled {
|
|
delete(mapResult, "notifyHook")
|
|
}
|
|
//req, err := http.NewRequest(c.Request.Method, fullRequestURL, requestBody)
|
|
// make new request with mapResult
|
|
}
|
|
if setting.MjModeClearEnabled {
|
|
if prompt, ok := mapResult["prompt"].(string); ok {
|
|
prompt = strings.Replace(prompt, "--fast", "", -1)
|
|
prompt = strings.Replace(prompt, "--relax", "", -1)
|
|
prompt = strings.Replace(prompt, "--turbo", "", -1)
|
|
|
|
mapResult["prompt"] = prompt
|
|
}
|
|
}
|
|
reqBody, err := common.Marshal(mapResult)
|
|
if err != nil {
|
|
return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "marshal_request_body_failed", http.StatusInternalServerError), nullBytes, err
|
|
}
|
|
req, err := http.NewRequest(c.Request.Method, fullRequestURL, strings.NewReader(string(reqBody)))
|
|
if err != nil {
|
|
return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "create_request_failed", http.StatusInternalServerError), nullBytes, err
|
|
}
|
|
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
|
// 使用带有超时的 context 创建新的请求
|
|
req = req.WithContext(ctx)
|
|
req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
|
|
req.Header.Set("Accept", c.Request.Header.Get("Accept"))
|
|
auth := common.GetContextKeyString(c, constant.ContextKeyChannelKey)
|
|
if auth != "" {
|
|
auth = strings.TrimPrefix(auth, "Bearer ")
|
|
req.Header.Set("mj-api-secret", auth)
|
|
}
|
|
defer cancel()
|
|
resp, err := GetHttpClient().Do(req)
|
|
if err != nil {
|
|
common.SysLog("do request failed: " + err.Error())
|
|
return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "do_request_failed", http.StatusInternalServerError), nullBytes, err
|
|
}
|
|
statusCode := resp.StatusCode
|
|
//if statusCode != 200 {
|
|
// return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "bad_response_status_code", statusCode), nullBytes, nil
|
|
//}
|
|
err = req.Body.Close()
|
|
if err != nil {
|
|
return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "close_request_body_failed", statusCode), nullBytes, err
|
|
}
|
|
err = c.Request.Body.Close()
|
|
if err != nil {
|
|
return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "close_request_body_failed", statusCode), nullBytes, err
|
|
}
|
|
var midjResponse dto.MidjourneyResponse
|
|
var midjourneyUploadsResponse dto.MidjourneyUploadResponse
|
|
responseBody, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "read_response_body_failed", statusCode), nullBytes, err
|
|
}
|
|
CloseResponseBodyGracefully(resp)
|
|
logger.LogDebug(c, "midjourney response body: %s", responseBody)
|
|
if len(responseBody) == 0 {
|
|
return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "empty_response_body", statusCode), responseBody, nil
|
|
} else {
|
|
err = common.Unmarshal(responseBody, &midjResponse)
|
|
if err != nil {
|
|
err2 := common.Unmarshal(responseBody, &midjourneyUploadsResponse)
|
|
if err2 != nil {
|
|
return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "unmarshal_response_body_failed", statusCode), responseBody, err
|
|
}
|
|
}
|
|
}
|
|
//for k, v := range resp.Header {
|
|
// c.Writer.Header().Set(k, v[0])
|
|
//}
|
|
return &dto.MidjourneyResponseWithStatusCode{
|
|
StatusCode: statusCode,
|
|
Response: midjResponse,
|
|
}, responseBody, nil
|
|
}
|