Image Generation

spring:
  ai:
    model:
      image: openai
    openai:
      api-key: ${OPENAI_API_KEY}
      timeout: 150s          # the default is 60 s (see the latency section)
      image:
        model: gpt-image-2.5-flare
        quality: medium
        size: 1024x1024

package com.example.docsassistant;

import java.net.URI;
import java.util.Base64;

import org.springframework.ai.image.Image;
import org.springframework.ai.image.ImageGeneration;
import org.springframework.ai.image.ImageModel;
import org.springframework.ai.image.ImageOptions;
import org.springframework.ai.image.ImageOptionsBuilder;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.ai.image.ImageResponse;
import org.springframework.ai.openai.metadata.OpenAiImageGenerationMetadata;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestClient;

@Service
class IllustrationService {

    record Illustration(String id, String revisedPrompt) {
    }

    private final ImageModel imageModel;
    private final IllustrationStore store;
    private final RestClient downloads = RestClient.create();

    IllustrationService(ImageModel imageModel, IllustrationStore store) {
        this.imageModel = imageModel;
        this.store = store;
    }

    /** One illustration for a docs page, with options every provider understands: size and count. */
    Illustration illustrate(String subject) {
        ImageOptions options = ImageOptionsBuilder.builder()
                .width(1024)
                .height(1024)
                .n(1)
                .build();
        String prompt = "A flat, minimal vector illustration for a technical documentation page about: " + subject
                + ". No text and no logos in the picture.";
        ImageResponse response = this.imageModel.call(new ImagePrompt(prompt, options));

        ImageGeneration generation = response.getResult();
        String id = this.store.save(bytesOf(generation.getOutput()));
        // OpenAI rewrites some prompts before drawing; keep what was really used, next to the picture
        String revised = generation.getMetadata() instanceof OpenAiImageGenerationMetadata openAi
                ? openAi.getRevisedPrompt() : null;
        return new Illustration(id, revised);
    }

    // Providers answer with base64 data or with a link that expires; store the bytes either way.
    private byte[] bytesOf(Image image) {
        if (image.getB64Json() != null) {
            return Base64.getDecoder().decode(image.getB64Json());
        }
        if (image.getUrl() != null) {
            URI url = URI.create(image.getUrl());
            if (!"https".equals(url.getScheme())) {
                throw new IllegalStateException("refusing to download from " + url.getScheme());
            }
            return this.downloads.get().uri(url).retrieve().body(byte[].class);
        }
        throw new IllegalStateException("the provider returned neither base64 data nor a URL");
    }
}

{"prompt":"A flat, minimal vector illustration for a technical documentation page about: chat memory. No text and no logos in the picture.","model":"gpt-image-1-mini","n":1,"size":"1024x1024"}

package com.example.docsassistant;

import java.util.List;

import org.springframework.ai.image.ImageMessage;
import org.springframework.ai.image.ImageModel;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.ai.image.ImageResponse;
import org.springframework.ai.openai.OpenAiImageOptions;
import org.springframework.ai.stabilityai.StyleEnum;
import org.springframework.ai.stabilityai.api.StabilityAiImageOptions;
import org.springframework.stereotype.Component;

/** Provider-specific options ride on the same ImagePrompt as the portable ones. */
@Component
class ImageOptionsExamples {

    private final ImageModel imageModel;

    ImageOptionsExamples(ImageModel imageModel) {
        this.imageModel = imageModel;
    }

    /** With spring.ai.model.image=openai (the default when only the OpenAI starter is present). */
    ImageResponse openAi(String prompt) {
        OpenAiImageOptions options = OpenAiImageOptions.builder()
                .model("gpt-image-2.5-flare")
                .width(1024)
                .height(1024)
                .quality("medium")
                .user("docsassistant-illustrations")   // an opaque id, helps the provider trace abuse; never an email
                .build();
        return this.imageModel.call(new ImagePrompt(prompt, options));
    }

    /** With spring.ai.model.image=stabilityai: weighted messages, a negative weight steers away from a trait. */
    ImageResponse stability(String subject) {
        StabilityAiImageOptions options = StabilityAiImageOptions.builder()
                .model("stable-diffusion-v1-6")
                .width(512)
                .height(512)
                .steps(30)
                .cfgScale(7f)
                .stylePreset(StyleEnum.DIGITAL_ART)
                .build();
        List<ImageMessage> messages = List.of(
                new ImageMessage("A clean technical diagram of " + subject, 1.0f),
                new ImageMessage("blurry, text, watermark", -1.0f));
        return this.imageModel.call(new ImagePrompt(messages, options));
    }
}

{"prompt":"cat","model":"gpt-image-2.5-flare","quality":"medium","size":"1024x1024","user":"docsassistant-illustrations"}

{"text_prompts":[{"text":"A clean technical diagram of chat memory","weight":1.0},{"text":"blurry, text, watermark","weight":-1.0}],
 "height":512,"width":512,"cfg_scale":7.0,"steps":30,"style_preset":"digital-art"}

package com.example.docsassistant;

import java.io.IOException;
import java.io.UncheckedIOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Optional;
import java.util.UUID;

import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

/** Keeps generated pictures on disk under names the server chooses. Swap the body for S3 or a database as needed. */
@Component
class IllustrationStore {

    private final Path directory;

    IllustrationStore(@Value("${docsassistant.images.directory:./generated-images}") Path directory)
            throws IOException {
        this.directory = Files.createDirectories(directory).toAbsolutePath().normalize();
    }

    /** The file name is a random UUID; nothing the model, the provider or the user wrote becomes part of the path. */
    String save(byte[] image) {
        String id = UUID.randomUUID().toString();
        try {
            Files.write(this.directory.resolve(id + ".png"), image);
        }
        catch (IOException ex) {
            throw new UncheckedIOException(ex);
        }
        return id;
    }

    Optional<byte[]> load(String id) {
        UUID parsed;
        try {
            parsed = UUID.fromString(id);   // rejects "../../etc/passwd" before it can reach the file system
        }
        catch (IllegalArgumentException ex) {
            return Optional.empty();
        }
        Path file = this.directory.resolve(parsed + ".png");
        try {
            return Files.exists(file) ? Optional.of(Files.readAllBytes(file)) : Optional.empty();
        }
        catch (IOException ex) {
            throw new UncheckedIOException(ex);
        }
    }
}

package com.example.docsassistant;

import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.security.access.prepost.PreAuthorize;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RestController;

@RestController
class IllustrationController {

    record Request(String subject) {
    }

    private final IllustrationService illustrations;
    private final IllustrationStore store;

    IllustrationController(IllustrationService illustrations, IllustrationStore store) {
        this.illustrations = illustrations;
        this.store = store;
    }

    // Every call costs money and produces content under your name: restrict who can trigger it.
    @PostMapping("/api/images/generate")
    @PreAuthorize("hasRole('MAINTAINER')")
    IllustrationService.Illustration generate(@RequestBody Request request) {
        String subject = request.subject() == null ? "" : request.subject().strip();
        if (subject.isEmpty() || subject.length() > 200) {
            throw new IllegalArgumentException("subject must be 1 to 200 characters");
        }
        return this.illustrations.illustrate(subject);
    }

    @GetMapping(path = "/api/images/{id}", produces = MediaType.IMAGE_PNG_VALUE)
    ResponseEntity<byte[]> fetch(@PathVariable String id) {
        return this.store.load(id).map(ResponseEntity::ok).orElse(ResponseEntity.notFound().build());
    }
}

package com.example.docsassistant;

import static org.assertj.core.api.Assertions.assertThat;

import java.io.IOException;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;

import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
import org.springframework.ai.image.Image;
import org.springframework.ai.image.ImageGeneration;
import org.springframework.ai.image.ImageModel;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.ai.image.ImageResponse;

class IllustrationServiceTest {

    @TempDir
    Path directory;

    @Test
    void storesTheBase64ImageAndAsksForAPortableSize() throws IOException {
        List<ImagePrompt> seen = new ArrayList<>();
        ImageModel fake = prompt -> {
            seen.add(prompt);
            String b64 = Base64.getEncoder().encodeToString("png-bytes".getBytes());
            return new ImageResponse(List.of(new ImageGeneration(new Image(null, b64))));
        };
        var store = new IllustrationStore(this.directory);
        var service = new IllustrationService(fake, store);

        var illustration = service.illustrate("chat memory");

        assertThat(store.load(illustration.id())).contains("png-bytes".getBytes());
        assertThat(illustration.revisedPrompt()).isNull();
        assertThat(seen.get(0).getOptions().getWidth()).isEqualTo(1024);
        assertThat(seen.get(0).getInstructions().get(0).getText()).contains("chat memory");
    }
}

package com.example.docsassistant;

import static org.assertj.core.api.Assertions.assertThat;

import java.io.IOException;
import java.nio.file.Path;

import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;

class IllustrationStoreTest {

    @TempDir
    Path directory;

    @Test
    void roundTripsUnderAServerChosenName() throws IOException {
        var store = new IllustrationStore(this.directory);

        String id = store.save("png".getBytes());

        assertThat(store.load(id)).contains("png".getBytes());
    }

    @Test
    void refusesAnythingThatIsNotAnId() throws IOException {
        var store = new IllustrationStore(this.directory);

        assertThat(store.load("../../etc/passwd")).isEmpty();
        assertThat(store.load("not-a-uuid")).isEmpty();
    }
}