fix(ali): honor image response format (#5513) (#7048)

relay/channel/ali/adaptor_test.go

@@ -1,16 +1,25 @@
 package ali
 
 import (
+	"encoding/base64"
 	"encoding/json"
+	"fmt"
+	"io"
 	"net/http"
 	"net/http/httptest"
+	"strings"
+	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/QuantumNous/new-api/common"
+	rootconstant "github.com/QuantumNous/new-api/constant"
 	relaycommon "github.com/QuantumNous/new-api/relay/common"
 	"github.com/QuantumNous/new-api/relay/constant"
 	relayhelper "github.com/QuantumNous/new-api/relay/helper"
 	"github.com/QuantumNous/new-api/relaykit/dto"
+	"github.com/QuantumNous/new-api/service"
+	"github.com/QuantumNous/new-api/setting/system_setting"
 	"github.com/gin-gonic/gin"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -159,3 +168,81 @@ func TestMappedAliImageModelUsesUpstreamProtocol(t *testing.T) {
 	assert.True(t, adaptor.IsSyncImageModel)
 	assert.IsType(t, &AliImageRequest{}, converted)
 }
+
+func TestAliImageHandlerHonorsRequestResponseFormat(t *testing.T) {
+	imageBytes := []byte("ali-image")
+	var downloads atomic.Int32
+	imageServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		downloads.Add(1)
+		w.Header().Set("Content-Type", "image/png")
+		_, _ = w.Write(imageBytes)
+	}))
+	t.Cleanup(imageServer.Close)
+
+	fetchSetting := system_setting.GetFetchSetting()
+	require.NotNil(t, fetchSetting)
+	originalFetchSetting := *fetchSetting
+	fetchSetting.EnableSSRFProtection = false
+	t.Cleanup(func() {
+		*fetchSetting = originalFetchSetting
+	})
+	originalMaxFileDownloadMB := rootconstant.MaxFileDownloadMB
+	rootconstant.MaxFileDownloadMB = 1
+	t.Cleanup(func() {
+		rootconstant.MaxFileDownloadMB = originalMaxFileDownloadMB
+	})
+	service.InitHttpClient()
+
+	tests := []struct {
+		name           string
+		responseFormat string
+		wantBase64     string
+		wantDownloads  int32
+	}{
+		{
+			name:           "base64",
+			responseFormat: "b64_json",
+			wantBase64:     base64.StdEncoding.EncodeToString(imageBytes),
+			wantDownloads:  1,
+		},
+		{
+			name:           "url",
+			responseFormat: "url",
+		},
+		{
+			name: "default",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			downloads.Store(0)
+			recorder := httptest.NewRecorder()
+			c, _ := gin.CreateTestContext(recorder)
+			info := &relaycommon.RelayInfo{
+				RelayMode: constant.RelayModeImagesGenerations,
+				StartTime: time.Unix(1, 0),
+				Request: &dto.ImageRequest{
+					ResponseFormat: tt.responseFormat,
+				},
+			}
+			responseBody := fmt.Sprintf(`{"output":{"results":[{"url":%q}]}}`, imageServer.URL)
+			resp := &http.Response{
+				StatusCode: http.StatusOK,
+				Header:     http.Header{},
+				Body:       io.NopCloser(strings.NewReader(responseBody)),
+			}
+
+			newAPIError, usage := aliImageHandler(&Adaptor{IsSyncImageModel: true}, c, resp, info)
+			require.Nil(t, newAPIError)
+			require.NotNil(t, usage)
+
+			var imageResponse dto.ImageResponse
+			require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &imageResponse))
+			require.Len(t, imageResponse.Data, 1)
+			assert.Equal(t, imageServer.URL, imageResponse.Data[0].Url)
+			assert.Equal(t, tt.wantBase64, imageResponse.Data[0].B64Json)
+			assert.Equal(t, tt.wantDownloads, downloads.Load())
+		})
+	}
+}

relay/channel/ali/image.go

@@ -284,7 +284,10 @@ func responseAli2OpenAIImage(c *gin.Context, response *AliResponse, originBody [
 }
 
 func aliImageHandler(a *Adaptor, c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*types.NewAPIError, *dto.Usage) {
-	responseFormat := c.GetString("response_format")
+	responseFormat := ""
+	if imageReq, ok := info.Request.(*dto.ImageRequest); ok {
+		responseFormat = imageReq.ResponseFormat
+	}
 
 	var aliTaskResponse AliResponse
 	responseBody, err := io.ReadAll(resp.Body)