fix(openai): keep image edit multipart body reusable

Author: Codex

relay/channel/openai/image_edit_test.go

@@ -0,0 +1,71 @@
+package openai
+
+import (
+	"bytes"
+	"io"
+	"mime/multipart"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/QuantumNous/new-api/dto"
+	relaycommon "github.com/QuantumNous/new-api/relay/common"
+	relayconstant "github.com/QuantumNous/new-api/relay/constant"
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/require"
+)
+
+func TestConvertImageEditRequestKeepsValidMultipartStreamFields(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	var body bytes.Buffer
+	writer := multipart.NewWriter(&body)
+	require.NoError(t, writer.WriteField("model", "gpt-image-1"))
+	require.NoError(t, writer.WriteField("prompt", "edit this image"))
+	require.NoError(t, writer.WriteField("stream", "true"))
+	require.NoError(t, writer.WriteField("partial_images", "3"))
+	part, err := writer.CreateFormFile("image", "input.png")
+	require.NoError(t, err)
+	_, err = part.Write([]byte("fake image"))
+	require.NoError(t, err)
+	require.NoError(t, writer.Close())
+
+	recorder := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(recorder)
+	c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body)
+	c.Request.Header.Set("Content-Type", writer.FormDataContentType())
+	require.NoError(t, c.Request.ParseMultipartForm(32<<20))
+
+	info := &relaycommon.RelayInfo{
+		RelayMode: relayconstant.RelayModeImagesEdits,
+	}
+	request := dto.ImageRequest{
+		Model:  "gpt-image-1",
+		Prompt: "edit this image",
+		Stream: true,
+	}
+
+	converted, err := (&Adaptor{}).ConvertImageRequest(c, info, request)
+	require.NoError(t, err)
+
+	convertedBody, ok := converted.(*bytes.Buffer)
+	require.True(t, ok)
+
+	contentType := c.Request.Header.Get("Content-Type")
+	replayedRequest := httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(convertedBody.Bytes()))
+	replayedRequest.Header.Set("Content-Type", contentType)
+	require.NoError(t, replayedRequest.ParseMultipartForm(32<<20))
+
+	require.Equal(t, "gpt-image-1", replayedRequest.PostForm.Get("model"))
+	require.Equal(t, "edit this image", replayedRequest.PostForm.Get("prompt"))
+	require.Equal(t, "true", replayedRequest.PostForm.Get("stream"))
+	require.Equal(t, "3", replayedRequest.PostForm.Get("partial_images"))
+	require.Len(t, replayedRequest.MultipartForm.File["image"], 1)
+
+	file, err := replayedRequest.MultipartForm.File["image"][0].Open()
+	require.NoError(t, err)
+	defer file.Close()
+	fileBytes, err := io.ReadAll(file)
+	require.NoError(t, err)
+	require.Equal(t, []byte("fake image"), fileBytes)
+}

relay/channel/openai/image_stream_test.go

@@ -24,7 +24,7 @@ func TestOpenaiImageStreamHandlerForwardsSSEAndUsage(t *testing.T) {
 		`event: image_generation.partial_image`,
 		`data: {"type":"image_generation.partial_image","b64_json":"partial"}`,
 		``,
-		`data: {"usage":{"prompt_tokens":3,"completion_tokens":4,"total_tokens":7}}`,
+		`data: {"usage":{"input_tokens":3,"output_tokens":4,"total_tokens":7,"input_tokens_details":{"image_tokens":2,"text_tokens":1}}}`,
 		``,
 		`data: [DONE]`,
 		``,
@@ -49,8 +49,10 @@ func TestOpenaiImageStreamHandlerForwardsSSEAndUsage(t *testing.T) {
 	require.Equal(t, 3, usage.PromptTokens)
 	require.Equal(t, 4, usage.CompletionTokens)
 	require.Equal(t, 7, usage.TotalTokens)
+	require.Equal(t, 2, usage.PromptTokensDetails.ImageTokens)
+	require.Equal(t, 1, usage.PromptTokensDetails.TextTokens)
 	require.Contains(t, recorder.Body.String(), `event: image_generation.partial_image`)
 	require.Contains(t, recorder.Body.String(), `data: {"type":"image_generation.partial_image","b64_json":"partial"}`)
-	require.Contains(t, recorder.Body.String(), `data: {"usage":{"prompt_tokens":3,"completion_tokens":4,"total_tokens":7}}`)
+	require.Contains(t, recorder.Body.String(), `data: {"usage":{"input_tokens":3,"output_tokens":4,"total_tokens":7,"input_tokens_details":{"image_tokens":2,"text_tokens":1}}}`)
 	require.Equal(t, "text/event-stream", recorder.Header().Get("Content-Type"))
 }

relay/channel/openai/relay-openai.go

@@ -575,22 +575,25 @@ func OpenaiHandlerWithUsage(c *gin.Context, info *relaycommon.RelayInfo, resp *h
 	// 写入新的 response body
 	service.IOCopyBytesGracefully(c, resp, responseBody)
 
-	// Once we've written to the client, we should not return errors anymore
-	// because the upstream has already consumed resources and returned content
-	// We should still perform billing even if parsing fails
-	// format
-	if usageResp.InputTokens > 0 {
-		usageResp.PromptTokens += usageResp.InputTokens
+	normalizeOpenAIUsage(&usageResp.Usage)
+	applyUsagePostProcessing(info, &usageResp.Usage, responseBody)
+	return &usageResp.Usage, nil
+}
+
+func normalizeOpenAIUsage(usage *dto.Usage) {
+	if usage == nil {
+		return
 	}
-	if usageResp.OutputTokens > 0 {
-		usageResp.CompletionTokens += usageResp.OutputTokens
+	if usage.InputTokens > 0 {
+		usage.PromptTokens += usage.InputTokens
 	}
-	if usageResp.InputTokensDetails != nil {
-		usageResp.PromptTokensDetails.ImageTokens += usageResp.InputTokensDetails.ImageTokens
-		usageResp.PromptTokensDetails.TextTokens += usageResp.InputTokensDetails.TextTokens
+	if usage.OutputTokens > 0 {
+		usage.CompletionTokens += usage.OutputTokens
+	}
+	if usage.InputTokensDetails != nil {
+		usage.PromptTokensDetails.ImageTokens += usage.InputTokensDetails.ImageTokens
+		usage.PromptTokensDetails.TextTokens += usage.InputTokensDetails.TextTokens
 	}
-	applyUsagePostProcessing(info, &usageResp.Usage, responseBody)
-	return &usageResp.Usage, nil
 }
 
 func OpenaiImageStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
@@ -609,7 +612,7 @@ func OpenaiImageStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp
 	}
 
 	scanner := bufio.NewScanner(resp.Body)
-	scanner.Buffer(make([]byte, helper.InitialScannerBufferSize), helper.DefaultMaxScannerBufferSize)
+	scanner.Buffer(make([]byte, helper.InitialScannerBufferSize), helper.GetScannerBufferSize())
 	for scanner.Scan() {
 		line := scanner.Text()
 		if strings.HasPrefix(line, "data:") {
@@ -621,8 +624,11 @@ func OpenaiImageStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp
 				info.ReceivedResponseCount++
 				lastStreamData = common.StringToByteSlice(data)
 				var usageResp dto.SimpleResponse
-				if err := common.Unmarshal(lastStreamData, &usageResp); err == nil && service.ValidUsage(&usageResp.Usage) {
-					usage = &usageResp.Usage
+				if err := common.Unmarshal(lastStreamData, &usageResp); err == nil {
+					normalizeOpenAIUsage(&usageResp.Usage)
+					if service.ValidUsage(&usageResp.Usage) {
+						usage = &usageResp.Usage
+					}
 				}
 			}
 		}

relay/helper/openai_image_request_test.go

@@ -2,11 +2,14 @@ package helper
 
 import (
 	"bytes"
+	"io"
 	"mime/multipart"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
 	"testing"
 
+	"github.com/QuantumNous/new-api/common"
 	relayconstant "github.com/QuantumNous/new-api/relay/constant"
 	"github.com/gin-gonic/gin"
 	"github.com/stretchr/testify/require"
@@ -26,6 +29,7 @@ func TestGetAndValidOpenAIImageRequestMultipartStream(t *testing.T) {
 	_, err = part.Write([]byte("fake image"))
 	require.NoError(t, err)
 	require.NoError(t, writer.Close())
+	originalBody := body.String()
 
 	recorder := httptest.NewRecorder()
 	c, _ := gin.CreateTestContext(recorder)
@@ -36,4 +40,13 @@ func TestGetAndValidOpenAIImageRequestMultipartStream(t *testing.T) {
 	require.NoError(t, err)
 	require.True(t, req.Stream)
 	require.True(t, req.IsStream(c))
+
+	bodyAfterValidation, err := io.ReadAll(c.Request.Body)
+	require.NoError(t, err)
+	require.Equal(t, originalBody, string(bodyAfterValidation))
+
+	form, err := common.ParseMultipartFormReusable(c)
+	require.NoError(t, err)
+	require.Equal(t, "true", url.Values(form.Value).Get("stream"))
+	require.Len(t, form.File["image"], 1)
 }

relay/helper/stream_scanner.go

@@ -27,7 +27,7 @@ const (
 	DefaultPingInterval         = 10 * time.Second
 )
 
-func getScannerBufferSize() int {
+func GetScannerBufferSize() int {
 	if constant.StreamScannerMaxBufferMB > 0 {
 		return constant.StreamScannerMaxBufferMB << 20
 	}
@@ -108,7 +108,7 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
 		close(stopChan)
 	}()
 
-	scanner.Buffer(make([]byte, InitialScannerBufferSize), getScannerBufferSize())
+	scanner.Buffer(make([]byte, InitialScannerBufferSize), GetScannerBufferSize())
 	scanner.Split(bufio.ScanLines)
 	SetEventStreamHeaders(c)
 

relay/helper/valid_request.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"math"
+	"net/url"
 	"strconv"
 	"strings"
 
@@ -146,11 +147,13 @@ func GetAndValidOpenAIImageRequest(c *gin.Context, relayMode int) (*dto.ImageReq
 	switch relayMode {
 	case relayconstant.RelayModeImagesEdits:
 		if strings.Contains(c.Request.Header.Get("Content-Type"), "multipart/form-data") {
-			_, err := c.MultipartForm()
+			form, err := common.ParseMultipartFormReusable(c)
 			if err != nil {
 				return nil, fmt.Errorf("failed to parse image edit form request: %w", err)
 			}
-			formData := c.Request.PostForm
+			formData := url.Values(form.Value)
+			c.Request.MultipartForm = form
+			c.Request.PostForm = formData
 			imageRequest.Prompt = formData.Get("prompt")
 			imageRequest.Model = formData.Get("model")
 			imageRequest.N = common.GetPointer(uint(common.String2Int(formData.Get("n"))))