diff --git a/relay/channel/ali/adaptor_test.go b/relay/channel/ali/adaptor_test.go index 08bc959acd..ec44c5968a 100644 --- a/relay/channel/ali/adaptor_test.go +++ b/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()) + }) + } +} diff --git a/relay/channel/ali/image.go b/relay/channel/ali/image.go index 6913fa346a..a2828ac47b 100644 --- a/relay/channel/ali/image.go +++ b/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)