Welcome to the fourth lesson of our course on building an image generation service in Java! So far, we've built several key components of our application: the PromptManager for formatting user inputs, the ImageManager for storing and processing images, and the ImageGeneratorService that connects to Google's Gemini API to generate images.

Now, we're ready to add another important layer to our application architecture: the controller. In Java web applications, controllers play a crucial role in the Model-View-Controller (MVC) pattern. They act as intermediaries between the service layer (which contains our business logic) and the presentation layer (which handles user interactions, such as HTTP requests and responses).

The controller's primary responsibilities include:

• Receiving and validating input from the user interface (such as HTTP requests)
• Calling the appropriate service methods with validated inputs
• Handling errors that might occur during processing
• Formatting responses in a consistent way before sending them back to the user

In our image generation application, the controller will receive text prompts and an aspectRatio from users, validate them, pass them to our ImageGeneratorService, and then format the responses (either successful image data or error messages) before sending them back.

By adding this controller layer, we're further improving the separation of concerns in our application. The service layer can focus purely on business logic (generating images), while the controller handles the specifics of HTTP requests and responses. This makes our code more maintainable, testable, and easier to extend in the future.

Let's start by creating our controller class. We'll create a new file called ImageGeneratorController.java in the controllers package. This controller will depend on our ImageGeneratorService to perform the actual image generation.

Here's how we'll set up the basic structure of our controller:

package com.codesignal.controllers;

import com.codesignal.services.ImageGeneratorService;

public class ImageGeneratorController {
    private ImageGeneratorService imageGeneratorService;

    public ImageGeneratorController() {
        this.imageGeneratorService = new ImageGeneratorService();
    }
}

In this code, we're importing our ImageGeneratorService class that we created in the previous lesson.

The controller's constructor initializes an instance of the ImageGeneratorService. This establishes the dependency between our controller and service layers.

Our controller will be responsible for two main operations:

1. Generating a new image based on user input and aspect ratio
2. Retrieving all previously generated images

For each of these operations, we'll create a dedicated method in our controller class. These methods will handle input validation, error handling, and response formatting, ensuring that our API provides a consistent interface to clients.

Now, let's implement the method for generating images. This method will receive user input and aspect ratio, validate them, call the service to generate an image, and format the response.

We'll define a simple Response class to encapsulate the response data and status code:

package com.codesignal.models;

public class Response {
    private Object body;
    private int statusCode;

    public Response(Object body, int statusCode) {
        this.body = body;
        this.statusCode = statusCode;
    }

    public Object getBody() {
        return body;
    }

    public int getStatusCode() {
        return statusCode;
    }
}

Now, let's implement the generateImage method in the controller:

package com.codesignal.controllers;

import com.codesignal.models.Response;
import com.codesignal.services.ImageGeneratorService;

import java.util.List;
import java.util.Map;

public class ImageGeneratorController {
    private ImageGeneratorService imageGeneratorService;

    public ImageGeneratorController() {
        this.imageGeneratorService = new ImageGeneratorService();
    }

    public Response generateImage(String userInput, String aspectRatio) {
        if (userInput == null || userInput.trim().isEmpty()) {
            return new Response(
                Map.of("error", "Missing input"),
                400
            );
        }

        String finalAspectRatio = (aspectRatio == null || aspectRatio.trim().isEmpty())
            ? "16:9"
            : aspectRatio;

        try {
            String base64Image = imageGeneratorService.generateImage(userInput, finalAspectRatio);
            return new Response(
                Map.of("image", base64Image),
                200
            );
        } catch (Exception e) {
            return new Response(
                Map.of("error", e.getMessage()),
                500
            );
        }
    }
}

Let's break down what's happening in this method:

• First, we check if the userInput is empty or null. If it is, we return an error response with a 400 status code (Bad Request), indicating that the client provided invalid input.
• Next, we normalize the aspectRatio. If it is missing or blank, we default to "16:9".
• If the input is valid, we enter a try-catch block. Within the try block, we call the generateImage method of our service, passing both the userInput and the final aspectRatio.
• If the image generation is successful, we return a Map with an "image" key containing the base64 string. This response will have a status code of 200 (OK).
• If an exception occurs during image generation, we catch it in the catch block and return an error response with a 500 status code (Internal Server Error).

Notice how we're using HTTP status codes to indicate the nature of the response. This is a standard practice in RESTful APIs:

• 200: Success
• 400: Client error (invalid input)
• 500: Server error (something went wrong on our end)

Next, let's implement the method for retrieving all previously generated images:

public Response getImages() {
    List<Map<String, Object>> images = imageGeneratorService.getAllImages();
    return new Response(
        Map.of("images", images),
        200
    );
}

This method is simpler than the previous one because it doesn't require input validation. We simply call the getAllImages method of our service, which returns a List of all stored images.

The response is a Map with an "images" key containing the list of image objects. Each image object includes an ID, the prompt used to generate it, and the base64-encoded image data.

With these two methods, our controller provides a complete interface for clients to generate images and retrieve previously generated ones. The controller abstracts away the details of how images are generated and stored, presenting a clean, consistent API to clients.

Now that we've implemented our controller, let's validate it to make sure it works correctly. We'll create a simple test class called Main.java to use the new controller:

package com.codesignal;

import com.codesignal.controllers.ImageGeneratorController;
import com.codesignal.models.Response;

import java.util.Map;

public class Main {
    public static void main(String[] args) {
        ImageGeneratorController controller = new ImageGeneratorController();

        try {
            Response response = controller.generateImage("Luxury Tech Conference 2025", "16:9");
            var body = response.getBody();
            var bodyJson = (Map<?, ?>) body;

            if (body instanceof Map && bodyJson.containsKey("error")) {
                System.out.println("Error: " + bodyJson.get("error"));
            } else {
                System.out.println("Generated Image:");
                System.out.println(bodyJson.get("image"));
            }
        } catch (Exception e) {
            System.out.println("Error generating image: " + e.getMessage());
        }

        Response allImagesResponse = controller.getImages();
        System.out.println("\nAll Images via Controller:");
        System.out.println(allImagesResponse.getBody());
    }
}

In this test script, we're creating an instance of our ImageGeneratorController and using it to perform two operations:

• First, we call the generateImage method with a sample prompt and aspect ratio. We check if the response contains an "error" key, which would indicate that something went wrong. If there's no error, we print the base64-encoded image data.
• Then, we call the getImages method to retrieve all stored images. We print the response body, which includes all previously generated images.

When running this script with a valid Gemini API key, you would see output similar to:

Generated Image:
/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkSEw8UHRofHh0a...

All Images via Controller:
{images=[{id=0, prompt=You are the lead graphic designer for 'Eventify Co.'..., image_base64=/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkSEw8UHRofHh0a...}]}

Let's also test the error handling by trying to generate an image with an empty prompt:

Response emptyResponse = controller.generateImage("", "16:9");
System.out.println("\nEmpty Input Test:");
System.out.println("Status Code: " + emptyResponse.getStatusCode());
System.out.println("Response: " + emptyResponse.getBody());

This would produce output like:

Empty Input Test:
Status Code: 400
Response: {error=Missing input}

This confirms that our controller correctly validates input and returns appropriate error responses with status codes.

In this lesson, we've built the ImageGeneratorController, which serves as an intermediary between our service layer and future web endpoints. This controller handles input validation, aspect ratio defaults, error management, and response formatting, providing a clean, consistent interface for clients to interact with our image generation service.

Let's review what we've learned:

1. We understood the role of controllers in the MVC architecture and how they fit into our application.
2. We created a controller class that depends on our ImageGeneratorService.
3. We implemented methods for generating images and retrieving previously generated ones.
4. We added input validation and error handling to ensure robust operation.
5. We passed the selected aspectRatio through to the service layer.
6. We tested our controller to verify that it works correctly.

The ImageGeneratorController is a crucial component of our application architecture. It abstracts away the details of how images are generated and stored, presenting a clean, consistent API to clients. This separation of concerns makes our code more maintainable, testable, and easier to extend.

In the next lesson, we'll integrate our controller with web endpoints to create a complete web API for image generation. We'll define endpoints that clients can call to generate images and retrieve previously generated ones.