mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-03-30 00:46:42 +00:00
feat(task): add model redirection, per-call billing, and multipart retry fix for async tasks
1. Async task model redirection (aligned with sync tasks):
- Integrate ModelMappedHelper in RelayTaskSubmit after model name
determination, populating OriginModelName / UpstreamModelName on RelayInfo.
- All task adaptors now send UpstreamModelName to upstream providers:
- Gemini & Vertex: BuildRequestURL uses UpstreamModelName.
- Doubao & Ali: BuildRequestBody conditionally overwrites body.Model.
- Vidu, Kling, Hailuo, Jimeng: convertToRequestPayload accepts RelayInfo
and unconditionally uses info.UpstreamModelName.
- Sora: BuildRequestBody parses JSON and multipart bodies to replace
the "model" field with UpstreamModelName.
- Frontend log visibility: LogTaskConsumption and taskBillingOther now
emit is_model_mapped / upstream_model_name in the "other" JSON field.
- Billing safety: RecalculateTaskQuotaByTokens reads model name from
BillingContext.OriginModelName (via taskModelName) instead of
task.Data["model"], preventing billing leaks from upstream model names.
2. Per-call billing (TaskPricePatches lifecycle):
- Rename TaskBillingContext.ModelName → OriginModelName; add PerCallBilling
bool field, populated from TaskPricePatches at submission time.
- settleTaskBillingOnComplete short-circuits when PerCallBilling is true,
skipping both adaptor adjustments and token-based recalculation.
- Remove ModelName from TaskSubmitResult; use relayInfo.OriginModelName
consistently in controller/relay.go for billing context and logging.
3. Multipart retry boundary mismatch fix:
- Root cause: after Sora (or OpenAI audio) rebuilds a multipart body with a
new boundary and overwrites c.Request.Header["Content-Type"], subsequent
calls to ParseMultipartFormReusable on retry would parse the cached
original body with the wrong boundary, causing "NextPart: EOF".
- Fix: ParseMultipartFormReusable now caches the original Content-Type in
gin context key "_original_multipart_ct" on first call and reuses it for
all subsequent parses, making multipart parsing retry-safe globally.
- Sora adaptor reverted to the standard pattern (direct header set/get),
which is now safe thanks to the root fix.
4. Tests:
- task_billing_test.go: update makeTask to use OriginModelName; add
PerCallBilling settlement tests (skip adaptor adjust, skip token recalc);
add non-per-call adaptor adjustment test with refund verification.
This commit is contained in:
@@ -253,8 +253,12 @@ func ProcessAliOtherRatios(aliReq *AliVideoRequest) (map[string]float64, error)
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) convertToAliRequest(info *relaycommon.RelayInfo, req relaycommon.TaskSubmitReq) (*AliVideoRequest, error) {
|
||||
upstreamModel := req.Model
|
||||
if info.IsModelMapped {
|
||||
upstreamModel = info.UpstreamModelName
|
||||
}
|
||||
aliReq := &AliVideoRequest{
|
||||
Model: req.Model,
|
||||
Model: upstreamModel,
|
||||
Input: AliVideoInput{
|
||||
Prompt: req.Prompt,
|
||||
ImgURL: req.InputReference,
|
||||
@@ -332,7 +336,7 @@ func (a *TaskAdaptor) convertToAliRequest(info *relaycommon.RelayInfo, req relay
|
||||
}
|
||||
}
|
||||
|
||||
if aliReq.Model != req.Model {
|
||||
if aliReq.Model != upstreamModel {
|
||||
return nil, errors.New("can't change model with metadata")
|
||||
}
|
||||
|
||||
|
||||
@@ -131,7 +131,11 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "convert request payload failed")
|
||||
}
|
||||
info.UpstreamModelName = body.Model
|
||||
if info.IsModelMapped {
|
||||
body.Model = info.UpstreamModelName
|
||||
} else {
|
||||
info.UpstreamModelName = body.Model
|
||||
}
|
||||
data, err := common.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
||||
@@ -105,7 +105,7 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
|
||||
|
||||
// BuildRequestURL constructs the upstream URL.
|
||||
func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) {
|
||||
modelName := info.OriginModelName
|
||||
modelName := info.UpstreamModelName
|
||||
version := model_setting.GetGeminiVersionSetting(modelName)
|
||||
|
||||
return fmt.Sprintf(
|
||||
|
||||
@@ -61,7 +61,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
|
||||
return nil, fmt.Errorf("invalid request type in context")
|
||||
}
|
||||
|
||||
body, err := a.convertToRequestPayload(&req)
|
||||
body, err := a.convertToRequestPayload(&req, info)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "convert request payload failed")
|
||||
}
|
||||
@@ -142,8 +142,8 @@ func (a *TaskAdaptor) GetChannelName() string {
|
||||
return ChannelName
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*VideoRequest, error) {
|
||||
modelConfig := GetModelConfig(req.Model)
|
||||
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq, info *relaycommon.RelayInfo) (*VideoRequest, error) {
|
||||
modelConfig := GetModelConfig(info.UpstreamModelName)
|
||||
duration := DefaultDuration
|
||||
if req.Duration > 0 {
|
||||
duration = req.Duration
|
||||
@@ -154,7 +154,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
|
||||
}
|
||||
|
||||
videoRequest := &VideoRequest{
|
||||
Model: req.Model,
|
||||
Model: info.UpstreamModelName,
|
||||
Prompt: req.Prompt,
|
||||
Duration: &duration,
|
||||
Resolution: resolution,
|
||||
|
||||
@@ -165,7 +165,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
|
||||
}
|
||||
}
|
||||
|
||||
body, err := a.convertToRequestPayload(&req)
|
||||
body, err := a.convertToRequestPayload(&req, info)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "convert request payload failed")
|
||||
}
|
||||
@@ -378,9 +378,9 @@ func hmacSHA256(key []byte, data []byte) []byte {
|
||||
return h.Sum(nil)
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*requestPayload, error) {
|
||||
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq, info *relaycommon.RelayInfo) (*requestPayload, error) {
|
||||
r := requestPayload{
|
||||
ReqKey: req.Model,
|
||||
ReqKey: info.UpstreamModelName,
|
||||
Prompt: req.Prompt,
|
||||
}
|
||||
|
||||
|
||||
@@ -150,7 +150,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
|
||||
}
|
||||
req := v.(relaycommon.TaskSubmitReq)
|
||||
|
||||
body, err := a.convertToRequestPayload(&req)
|
||||
body, err := a.convertToRequestPayload(&req, info)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -248,15 +248,15 @@ func (a *TaskAdaptor) GetChannelName() string {
|
||||
// helpers
|
||||
// ============================
|
||||
|
||||
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*requestPayload, error) {
|
||||
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq, info *relaycommon.RelayInfo) (*requestPayload, error) {
|
||||
r := requestPayload{
|
||||
Prompt: req.Prompt,
|
||||
Image: req.Image,
|
||||
Mode: taskcommon.DefaultString(req.Mode, "std"),
|
||||
Duration: fmt.Sprintf("%d", taskcommon.DefaultInt(req.Duration, 5)),
|
||||
AspectRatio: a.getAspectRatio(req.Size),
|
||||
ModelName: req.Model,
|
||||
Model: req.Model, // Keep consistent with model_name, double writing improves compatibility
|
||||
ModelName: info.UpstreamModelName,
|
||||
Model: info.UpstreamModelName,
|
||||
CfgScale: 0.5,
|
||||
StaticMask: "",
|
||||
DynamicMasks: []DynamicMask{},
|
||||
@@ -266,6 +266,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
|
||||
}
|
||||
if r.ModelName == "" {
|
||||
r.ModelName = "kling-v1"
|
||||
r.Model = "kling-v1"
|
||||
}
|
||||
if err := taskcommon.UnmarshalMetadata(req.Metadata, &r); err != nil {
|
||||
return nil, errors.Wrap(err, "unmarshal metadata failed")
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
package sora
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
@@ -145,6 +147,59 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "get_request_body_failed")
|
||||
}
|
||||
cachedBody, err := storage.Bytes()
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "read_body_bytes_failed")
|
||||
}
|
||||
contentType := c.GetHeader("Content-Type")
|
||||
|
||||
if strings.HasPrefix(contentType, "application/json") {
|
||||
var bodyMap map[string]interface{}
|
||||
if err := common.Unmarshal(cachedBody, &bodyMap); err == nil {
|
||||
bodyMap["model"] = info.UpstreamModelName
|
||||
if newBody, err := common.Marshal(bodyMap); err == nil {
|
||||
return bytes.NewReader(newBody), nil
|
||||
}
|
||||
}
|
||||
return bytes.NewReader(cachedBody), nil
|
||||
}
|
||||
|
||||
if strings.Contains(contentType, "multipart/form-data") {
|
||||
formData, err := common.ParseMultipartFormReusable(c)
|
||||
if err != nil {
|
||||
return bytes.NewReader(cachedBody), nil
|
||||
}
|
||||
var buf bytes.Buffer
|
||||
writer := multipart.NewWriter(&buf)
|
||||
writer.WriteField("model", info.UpstreamModelName)
|
||||
for key, values := range formData.Value {
|
||||
if key == "model" {
|
||||
continue
|
||||
}
|
||||
for _, v := range values {
|
||||
writer.WriteField(key, v)
|
||||
}
|
||||
}
|
||||
for fieldName, fileHeaders := range formData.File {
|
||||
for _, fh := range fileHeaders {
|
||||
f, err := fh.Open()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
part, err := writer.CreateFormFile(fieldName, fh.Filename)
|
||||
if err != nil {
|
||||
f.Close()
|
||||
continue
|
||||
}
|
||||
io.Copy(part, f)
|
||||
f.Close()
|
||||
}
|
||||
}
|
||||
writer.Close()
|
||||
c.Request.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
return &buf, nil
|
||||
}
|
||||
|
||||
return common.ReaderOnly(storage), nil
|
||||
}
|
||||
|
||||
|
||||
@@ -86,7 +86,7 @@ func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, erro
|
||||
if err := common.Unmarshal([]byte(a.apiKey), adc); err != nil {
|
||||
return "", fmt.Errorf("failed to decode credentials: %w", err)
|
||||
}
|
||||
modelName := info.OriginModelName
|
||||
modelName := info.UpstreamModelName
|
||||
if modelName == "" {
|
||||
modelName = "veo-3.0-generate-001"
|
||||
}
|
||||
|
||||
@@ -116,7 +116,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
|
||||
}
|
||||
req := v.(relaycommon.TaskSubmitReq)
|
||||
|
||||
body, err := a.convertToRequestPayload(&req)
|
||||
body, err := a.convertToRequestPayload(&req, info)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -224,9 +224,9 @@ func (a *TaskAdaptor) GetChannelName() string {
|
||||
// helpers
|
||||
// ============================
|
||||
|
||||
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*requestPayload, error) {
|
||||
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq, info *relaycommon.RelayInfo) (*requestPayload, error) {
|
||||
r := requestPayload{
|
||||
Model: taskcommon.DefaultString(req.Model, "viduq1"),
|
||||
Model: taskcommon.DefaultString(info.UpstreamModelName, "viduq1"),
|
||||
Images: req.Images,
|
||||
Prompt: req.Prompt,
|
||||
Duration: taskcommon.DefaultInt(req.Duration, 5),
|
||||
|
||||
@@ -26,7 +26,6 @@ type TaskSubmitResult struct {
|
||||
UpstreamTaskID string
|
||||
TaskData []byte
|
||||
Platform constant.TaskPlatform
|
||||
ModelName string
|
||||
Quota int
|
||||
//PerCallPrice types.PriceData
|
||||
}
|
||||
@@ -163,6 +162,13 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe
|
||||
modelName = service.CoverTaskActionToModelName(platform, info.Action)
|
||||
}
|
||||
|
||||
// 2.5 应用渠道的模型映射(与同步任务对齐)
|
||||
info.OriginModelName = modelName
|
||||
info.UpstreamModelName = modelName
|
||||
if err := helper.ModelMappedHelper(c, info, nil); err != nil {
|
||||
return nil, service.TaskErrorWrapperLocal(err, "model_mapping_failed", http.StatusBadRequest)
|
||||
}
|
||||
|
||||
// 3. 预生成公开 task ID(仅首次)
|
||||
if info.PublicTaskID == "" {
|
||||
info.PublicTaskID = model.GenerateTaskID()
|
||||
@@ -241,7 +247,6 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe
|
||||
UpstreamTaskID: upstreamTaskID,
|
||||
TaskData: taskData,
|
||||
Platform: platform,
|
||||
ModelName: modelName,
|
||||
Quota: finalQuota,
|
||||
}, nil
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user