fix: WanxImageModel Image response list IndexOutOfBound (#307)

* fix: WanxImageModel Image response list IndexOutOfBound
wanx2.1-t2i-turbo model will return error code whrn prompt contains illegal info which casued IndexOutOfBound Exception when processing respons.

* make format

---------

Co-authored-by: Martin7-1 <[email protected]>

Author: Flynn-zhangyg

models/langchain4j-community-dashscope/src/main/java/dev/langchain4j/community/model/dashscope/WanxImageModel.java

@@ -88,43 +88,26 @@ public WanxImageModel(
 
     @Override
     public Response<Image> generate(String prompt) {
-        ImageSynthesisParam.ImageSynthesisParamBuilder<?, ?> builder =
-                requestBuilder(prompt).n(1);
-
-        try {
-            imageSynthesisParamCustomizer.accept(builder);
-            ImageSynthesisResult result = imageSynthesis.call(builder.build());
-            return Response.from(imagesFrom(result).get(0));
-        } catch (NoApiKeyException e) {
-            throw new IllegalArgumentException(e);
-        }
+        return Response.from(doGenerate(requestBuilder(prompt).n(1)).get(0));
     }
 
     @Override
     public Response<List<Image>> generate(String prompt, int n) {
-        ImageSynthesisParam.ImageSynthesisParamBuilder<?, ?> builder =
-                requestBuilder(prompt).n(n);
-
-        try {
-            imageSynthesisParamCustomizer.accept(builder);
-            ImageSynthesisResult result = imageSynthesis.call(builder.build());
-            return Response.from(imagesFrom(result));
-        } catch (NoApiKeyException e) {
-            throw new IllegalArgumentException(e);
-        }
+        return Response.from(doGenerate(requestBuilder(prompt).n(n)));
     }
 
     @Override
     public Response<Image> edit(Image image, String prompt) {
         String imageUrl = imageUrl(image, modelName, apiKey);
-
         ImageSynthesisParam.ImageSynthesisParamBuilder<?, ?> builder =
                 requestBuilder(prompt).refImage(imageUrl).n(1);
-
         if (imageUrl.startsWith("oss://")) {
             builder.header("X-DashScope-OssResourceResolve", "enable");
         }
+        return Response.from(doGenerate(builder).get(0));
+    }
 
+    private List<Image> doGenerate(ImageSynthesisParam.ImageSynthesisParamBuilder<?, ?> builder) {
         try {
             imageSynthesisParamCustomizer.accept(builder);
             ImageSynthesisResult result = imageSynthesis.call(builder.build());
@@ -135,7 +118,7 @@ public Response<Image> edit(Image image, String prompt) {
                         String.format("[%s] %s: %s", output.getTaskStatus(), output.getCode(), output.getMessage());
                 throw new IllegalStateException(errorMessage);
             }
-            return Response.from(images.get(0));
+            return images;
         } catch (NoApiKeyException e) {
             throw new IllegalArgumentException(e);
         }

models/langchain4j-community-dashscope/src/test/java/dev/langchain4j/community/model/dashscope/WanxImageModelIT.java

@@ -3,10 +3,12 @@
 import static dev.langchain4j.community.model.dashscope.QwenTestHelper.apiKey;
 import static dev.langchain4j.community.model.dashscope.QwenTestHelper.multimodalImageData;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 import dev.langchain4j.data.image.Image;
 import dev.langchain4j.model.output.Response;
 import java.net.URI;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
@@ -135,4 +137,17 @@ void simple_image_edition_works_by_data(String modelName) {
         log.info("Your remote image is here: {}", remoteImage);
         assertThat(remoteImage).isNotNull();
     }
+
+    // This case only occurs when using model WANX2_1_T2I_TURBO.
+    @Test
+    void simple_image_generation_return_error_msg_when_prompt_illegal() {
+        WanxImageModel model = WanxImageModel.builder()
+                .apiKey(apiKey())
+                .modelName(WanxModelName.WANX2_1_T2I_TURBO)
+                .build();
+
+        assertThatThrownBy(() -> model.generate("Draw a Donald Duck."))
+                .isInstanceOf(IllegalStateException.class)
+                .hasMessageContaining("IPInfringementSuspect");
+    }
 }