Skip to content

Commit

Permalink
Add comments and prevent IndexOutOfBounds in ImageResponse
Browse files Browse the repository at this point in the history
* Add comments in ImageResponse like ChatResponse
* Prevent IndexOutOfBounds in getResult function
* Update toString for formatting like ChatResponse
  • Loading branch information
devholic22 authored and markpollack committed May 21, 2024
1 parent ef5a3fa commit a46e014
Showing 1 changed file with 46 additions and 6 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -19,41 +19,81 @@
import java.util.Objects;

import org.springframework.ai.model.ModelResponse;
import org.springframework.util.CollectionUtils;

/**
* The image completion (e.g. imageGeneration) response returned by an AI provider.
*
* @author Mark Pollack
* @author Christian Tzolov
* @author Hyunjoon Choi
*/
public class ImageResponse implements ModelResponse<ImageGeneration> {

private final ImageResponseMetadata imageResponseMetadata;

/**
* List of generate images returned by the AI provider.
*/
private final List<ImageGeneration> imageGenerations;

/**
* Construct a new {@link ImageResponse} instance without metadata.
* @param generations the {@link List} of {@link ImageGeneration} returned by the AI
* provider.
*/
public ImageResponse(List<ImageGeneration> generations) {
this(generations, ImageResponseMetadata.NULL);
}

/**
* Construct a new {@link ImageResponse} instance.
* @param generations the {@link List} of {@link ImageGeneration} returned by the AI
* provider.
* @param imageResponseMetadata {@link ImageResponseMetadata} containing information
* about the use of the AI provider's API.
*/
public ImageResponse(List<ImageGeneration> generations, ImageResponseMetadata imageResponseMetadata) {
this.imageResponseMetadata = imageResponseMetadata;
this.imageGenerations = List.copyOf(generations);
}

/**
* The {@link List} of {@link ImageGeneration generated outputs}.
* <p>
* It is a {@link List} of {@link List lists} because the Prompt could request
* multiple output {@link ImageGeneration generations}.
* @return the {@link List} of {@link ImageGeneration generated outputs}.
*/
@Override
public ImageGeneration getResult() {
return imageGenerations.get(0);
public List<ImageGeneration> getResults() {
return imageGenerations;
}

/**
* @return Returns the first {@link ImageGeneration} in the generations list.
*/
@Override
public List<ImageGeneration> getResults() {
return imageGenerations;
public ImageGeneration getResult() {
if (CollectionUtils.isEmpty(this.imageGenerations)) {
return null;
}
return imageGenerations.get(0);
}

/**
* @return Returns {@link ImageResponseMetadata} containing information about the use
* of the AI provider's API.
*/
@Override
public ImageResponseMetadata getMetadata() {
return imageResponseMetadata;
}

@Override
public String toString() {
return "ImageResponse{" + "imageResponseMetadata=" + imageResponseMetadata + ", imageGenerations="
+ imageGenerations + '}';
return "ImageResponse [" + "imageResponseMetadata=" + imageResponseMetadata + ", imageGenerations="
+ imageGenerations + "]";
}

@Override
Expand Down

0 comments on commit a46e014

Please sign in to comment.