|
19 | 19 | import java.util.Objects; |
20 | 20 |
|
21 | 21 | import org.springframework.ai.model.ModelResponse; |
| 22 | +import org.springframework.util.CollectionUtils; |
22 | 23 |
|
| 24 | +/** |
| 25 | + * The image completion (e.g. imageGeneration) response returned by an AI provider. |
| 26 | + * |
| 27 | + * @author Mark Pollack |
| 28 | + * @author Christian Tzolov |
| 29 | + * @author Hyunjoon Choi |
| 30 | + */ |
23 | 31 | public class ImageResponse implements ModelResponse<ImageGeneration> { |
24 | 32 |
|
25 | 33 | private final ImageResponseMetadata imageResponseMetadata; |
26 | 34 |
|
| 35 | + /** |
| 36 | + * List of generate images returned by the AI provider. |
| 37 | + */ |
27 | 38 | private final List<ImageGeneration> imageGenerations; |
28 | 39 |
|
| 40 | + /** |
| 41 | + * Construct a new {@link ImageResponse} instance without metadata. |
| 42 | + * @param generations the {@link List} of {@link ImageGeneration} returned by the AI |
| 43 | + * provider. |
| 44 | + */ |
29 | 45 | public ImageResponse(List<ImageGeneration> generations) { |
30 | 46 | this(generations, ImageResponseMetadata.NULL); |
31 | 47 | } |
32 | 48 |
|
| 49 | + /** |
| 50 | + * Construct a new {@link ImageResponse} instance. |
| 51 | + * @param generations the {@link List} of {@link ImageGeneration} returned by the AI |
| 52 | + * provider. |
| 53 | + * @param imageResponseMetadata {@link ImageResponseMetadata} containing information |
| 54 | + * about the use of the AI provider's API. |
| 55 | + */ |
33 | 56 | public ImageResponse(List<ImageGeneration> generations, ImageResponseMetadata imageResponseMetadata) { |
34 | 57 | this.imageResponseMetadata = imageResponseMetadata; |
35 | 58 | this.imageGenerations = List.copyOf(generations); |
36 | 59 | } |
37 | 60 |
|
| 61 | + /** |
| 62 | + * The {@link List} of {@link ImageGeneration generated outputs}. |
| 63 | + * <p> |
| 64 | + * It is a {@link List} of {@link List lists} because the Prompt could request |
| 65 | + * multiple output {@link ImageGeneration generations}. |
| 66 | + * @return the {@link List} of {@link ImageGeneration generated outputs}. |
| 67 | + */ |
38 | 68 | @Override |
39 | | - public ImageGeneration getResult() { |
40 | | - return imageGenerations.get(0); |
| 69 | + public List<ImageGeneration> getResults() { |
| 70 | + return imageGenerations; |
41 | 71 | } |
42 | 72 |
|
| 73 | + /** |
| 74 | + * @return Returns the first {@link ImageGeneration} in the generations list. |
| 75 | + */ |
43 | 76 | @Override |
44 | | - public List<ImageGeneration> getResults() { |
45 | | - return imageGenerations; |
| 77 | + public ImageGeneration getResult() { |
| 78 | + if (CollectionUtils.isEmpty(this.imageGenerations)) { |
| 79 | + return null; |
| 80 | + } |
| 81 | + return imageGenerations.get(0); |
46 | 82 | } |
47 | 83 |
|
| 84 | + /** |
| 85 | + * @return Returns {@link ImageResponseMetadata} containing information about the use |
| 86 | + * of the AI provider's API. |
| 87 | + */ |
48 | 88 | @Override |
49 | 89 | public ImageResponseMetadata getMetadata() { |
50 | 90 | return imageResponseMetadata; |
51 | 91 | } |
52 | 92 |
|
53 | 93 | @Override |
54 | 94 | public String toString() { |
55 | | - return "ImageResponse{" + "imageResponseMetadata=" + imageResponseMetadata + ", imageGenerations=" |
56 | | - + imageGenerations + '}'; |
| 95 | + return "ImageResponse [" + "imageResponseMetadata=" + imageResponseMetadata + ", imageGenerations=" |
| 96 | + + imageGenerations + "]"; |
57 | 97 | } |
58 | 98 |
|
59 | 99 | @Override |
|
0 commit comments