Taking Spring AI for A Spin: A Developer's Playground

Spring has long been the Swiss Army knife in a Java developer's toolkit. From managing dependencies to simplifying web app development, the Spring ecosystem has been making our lives easier for decades. Now, with the AI revolution in full swing, Spring AI has entered the chat, bringing machine learning capabilities directly into the Spring framework we know and love.

Let's dive into Spring AI, get our hands dirty, and see what this exciting new addition to the Spring family can do!

What is Spring AI?

Spring AI is an extension of the Spring framework designed to simplify the integration of AI capabilities into your applications. Think of it as doing for AI what Spring Boot did for application setup, taking away the boilerplate and letting you focus on building cool stuff. 

The project aims to provide abstractions across different AI providers (like OpenAI, Azure, AWS Bedrock, etc.), meaning you can switch between them without rewriting your application logic the Spring way. 

Setting Up Your First Spring AI Project

Let's start by creating a simple Spring Boot application with Spring AI, leveraging the power of Java 21. We'll use Spring Initializr to set up our project.

@SpringBootApplication
public class Main {
    public static void main(String[] args) {
        SpringApplication.run(Main.class, args);
}

Next, we need to add Spring AI dependencies to our pom.xml

<!-- Set Java 21 as the target -->
<properties>
    <java.version>21</java.version>
    <spring-ai.version>0.8.0</spring-ai.version>
</properties>

<!-- Spring AI Core -->
<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-core</artifactId>
    <version>${spring-ai.version}</version>
</dependency>

<!-- Choose your AI provider - lets go with OpenAI for this example -->
<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-openai-spring-boot-starter</artifactId>
    <version>${spring-ai.version}</version>
</dependency>

Configuring The AI Provider

Before we can make our first AI call, we need to set up our API credentials. Let's add them to our application.properties. 

spring.ai.openai.api-key=${OPENAI_API_KEY}
spring.ai.openai.model=gpt-4

Your First AI Interaction with Spring AI

Now, let's create a service that will interact with the AI, leveraging Java 21's text blocks and virtual threads:

@Service
public class AIAssistantService {
    
    private final OpenAiChatClient chatClient;
    
    public AIAssistantService(OpenAiChatClient chatClient) {
        this.chatClient = chatClient;
    }
    
    public String askQuestion(String question) {
        Prompt prompt = new Prompt(question);
        
        // Using virtual threads for non-blocking operations (Java 21 feature)
        return Thread.ofVirtual().name("ai-request-" + System.currentTimeMillis())
            .callableTask(() -> {
                ChatResponse response = chatClient.call(prompt);
                return response.getResult().getOutput().getContent();
            })
            .start()
            .join();
    }
}

And a simple REST controller to expose this functionality:

@RestController
@RequestMapping("/ai")
public class AIController {
    
    private final AIAssistantService assistantService;
    
    public AIController(AIAssistantService assistantService) {
        this.assistantService = assistantService;
    }
    
    @GetMapping("/ask")
    public String askAI(@RequestParam String question) {
        return assistantService.askQuestion(question);
    }
}

Just like that, you have a web endpoint that can answer any question using AI! Try hitting http://localhost:8080/ai/ask?question=What is Spring Framework?

Switching Between AI Models Using Environment Variables

One of Spring AI's most powerful features is its ability to dynamically switch between models or providers based on environment configurations. Let's implement a flexible service that uses environment variables to determine which AI provider and model to use:

@Service
public class DynamicModelService {
    
    private static final Logger log = LoggerFactory.getLogger(DynamicModelService.class);
    
    // Inject all potential AI clients
    private final OpenAiChatClient openAiClient;
    private final BedrockChatClient bedrockClient;
    private final VertexAiChatClient vertexAiClient;
    
    @Value("${app.ai.provider:openai}")
    private String activeProvider;
    
    @Value("${app.ai.model.openai:gpt-4}")
    private String openAiModel;
    
    @Value("${app.ai.model.bedrock:anthropic.claude-3-sonnet-20240229-v1:0}")
    private String bedrockModel;
    
    @Value("${app.ai.model.vertex:gemini-pro}")
    private String vertexModel;
    
    public DynamicModelService(
            OpenAiChatClient openAiClient,
            BedrockChatClient bedrockClient,
            VertexAiChatClient vertexAiClient) {
        this.openAiClient = openAiClient;
        this.bedrockClient = bedrockClient;
        this.vertexAiClient = vertexAiClient;
        
        log.info("Dynamic Model Service initialized with provider: {}", activeProvider);
        log.info("OpenAI model: {}", openAiModel);
        log.info("Bedrock model: {}", bedrockModel);
        log.info("Vertex model: {}", vertexModel);
    }
    
    public record ModelResponse(String content, String provider, String model, long processingTimeMs) {}
    
    public ModelResponse processPrompt(String prompt) {
        long startTime = System.currentTimeMillis();
        String content;
        String provider = activeProvider.toLowerCase();
        String model;
        
        switch(provider) {
            case "openai" -> {
                model = openAiModel;
                content = openAiClient.call(
                    ChatOptions.builder().withModel(model).build(),
                    new Prompt(prompt)
                ).getResult().getOutput().getContent();
            }
            case "bedrock" -> {
                model = bedrockModel;
                content = bedrockClient.call(
                    ChatOptions.builder().withModel(model).build(),
                    new Prompt(prompt)
                ).getResult().getOutput().getContent();
            }
            case "vertex" -> {
                model = vertexModel;
                content = vertexAiClient.call(
                    ChatOptions.builder().withModel(model).build(),
                    new Prompt(prompt)
                ).getResult().getOutput().getContent();
            }
            default -> throw new IllegalArgumentException("Unknown provider: " + provider);
        }
        
        long processingTime = System.currentTimeMillis() - startTime;
        return new ModelResponse(content, provider, model, processingTime);
    }
}

Now let's create a controller to use this service:

@RestController
@RequestMapping("/ai/dynamic")
public class DynamicModelController {
    
    private final DynamicModelService modelService;
    
    public DynamicModelController(DynamicModelService modelService) {
        this.modelService = modelService;
    }
    
    @GetMapping("/ask")
    public DynamicModelService.ModelResponse askDynamic(@RequestParam String prompt) {
        return modelService.processPrompt(prompt);
    }
}

Environment-Based Configuration

To make our dynamic model switching work, we need to configure our application.properties or application.yml file. Here's how it would look:

# Base AI provider configurations
spring:
  ai:
    openai:
      api-key: ${OPENAI_API_KEY}
    bedrock:
      aws:
        region: ${AWS_REGION:us-west-2}
        access-key: ${AWS_ACCESS_KEY}
        secret-key: ${AWS_SECRET_KEY}
    vertex:
      ai:
        project-id: ${GCP_PROJECT_ID}
        location: ${GCP_LOCATION:us-central1}

# Application-specific AI settings
app:
  ai:
    provider: ${AI_PROVIDER:openai}
    model:
      openai: ${OPENAI_MODEL:gpt-4}
      bedrock: ${BEDROCK_MODEL:anthropic.claude-3-sonnet-20240229-v1:0}
      vertex: ${VERTEX_MODEL:gemini-pro}

Runtime Provider Switching

With this setup, you can easily switch between providers by changing environment variables. For example:

# Switch to OpenAI
export AI_PROVIDER=openai
export OPENAI_MODEL=gpt-4-0125-preview

# OR switch to AWS Bedrock
export AI_PROVIDER=bedrock
export BEDROCK_MODEL=anthropic.claude-3-sonnet-20240229-v1:0

# OR switch to Google Vertex AI
export AI_PROVIDER=vertex
export VERTEX_MODEL=gemini-pro-vision

Now, when you restart your Spring application (or if you're using Spring Cloud Config for dynamic property updates), it will use the configured provider and model!

Creating Environment-Specific Profiles

Another powerful approach is to leverage Spring Boot profiles. Let's add profile-specific configurations:

@Configuration
public class AIModelConfiguration {
    
    @Bean
    @Profile("openai")
    public AIModelResolver openAiModelResolver(OpenAiChatClient openAiClient) {
        return prompt -> openAiClient.call(new Prompt(prompt))
            .getResult().getOutput().getContent();
    }
    
    @Bean
    @Profile("bedrock")
    public AIModelResolver bedrockModelResolver(BedrockChatClient bedrockClient) {
        return prompt -> bedrockClient.call(new Prompt(prompt))
            .getResult().getOutput().getContent();
    }
    
    @Bean
    @Profile("vertex")
    public AIModelResolver vertexModelResolver(VertexAiChatClient vertexAiClient) {
        return prompt -> vertexAiClient.call(new Prompt(prompt))
            .getResult().getOutput().getContent();
    }
    
    // Simple interface to abstract AI resolution
    public interface AIModelResolver {
        String resolvePrompt(String prompt);
    }
}

Now you can switch providers simply by changing the active profile:

# Start with OpenAI
./mvnw spring-boot:run -Dspring-boot.run.profiles=openai

# Or with AWS Bedrock
./mvnw spring-boot:run -Dspring-boot.run.profiles=bedrock

# Or with Google Vertex AI
./mvnw spring-boot:run -Dspring-boot.run.profiles=vertex

Advanced Features

Spring AI shines when you start using its more advanced features

Prompt Templates

Let's explore prompt templates, which allow you to create structured interactions with AI models:

@Service
public class RecipeService {
    
    private final OpenAiChatClient chatClient;
    
    public RecipeService(OpenAiChatClient chatClient) {
        this.chatClient = chatClient;
    }
    
    public String generateRecipe(String ingredients, String dietaryRestrictions) {
        PromptTemplate template = new PromptTemplate("""
            Create a recipe using these ingredients: {ingredients}.
            The recipe must follow these dietary restrictions: {restrictions}.
            Format the output as a title, followed by ingredients list, then step-by-step instructions.
            """);
            
        Map<String, Object> variables = new HashMap<>();
        variables.put("ingredients", ingredients);
        variables.put("restrictions", dietaryRestrictions);
        
        Prompt prompt = template.create(variables);
        ChatResponse response = chatClient.call(prompt);
        
        return response.getResult().getOutput().getContent();
    }
}

Working with Embeddings

One of the most powerful features of Spring AI is its support for vector embeddings, which enable semantic search and similarity-based recommendations:

@Service
public class DocumentSearchService {
    
    private final OpenAiEmbeddingClient embeddingClient;
    private final VectorStore vectorStore;
    
    // Constructor injection...
    
    public void indexDocument(String id, String content) {
        List<Document> documents = List.of(new Document(id, content));
        vectorStore.add(embeddingClient.embed(documents));
    }
    
    public List<Document> searchSimilarDocuments(String query, int limit) {
        Embedding queryEmbedding = embeddingClient.embed(query);
        return vectorStore.similaritySearch(queryEmbedding, limit);
    }
}

Handling AI Rate Limits with Spring Retry

When working with AI APIs, you'll inevitably hit rate limits. Spring Retry to the rescue:

@Service
@Retryable(maxAttempts = 3, backoff = @Backoff(delay = 2000))
public class ResilientAIService {
    
    private final OpenAiChatClient chatClient;
    
    // Constructor...
    
    public String generateContent(String prompt) {
        try {
            return chatClient.call(new Prompt(prompt))
                .getResult().getOutput().getContent();
        } catch (Exception e) {
            log.warn("AI service temporarily unavailable, retrying...", e);
            throw e; // Let Spring Retry handle the retry
        }
    }
    
    @Recover
    public String fallbackMethod(Exception e, String prompt) {
        log.error("All retries failed for prompt: " + prompt, e);
        return "I'm sorry, I'm unable to generate content at the moment.";
    }
}

Streaming Responses

For a more interactive experience, Spring AI supports streaming responses, now with Java 21's virtual threads for better scalability:

@GetMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
public Flux<String> streamResponse(@RequestParam String prompt) {
    // Using Java 21's Structured Concurrency
    try (var scope = StructuredTaskScope.ShutdownOnSuccess<Flux<String>>()) {
        var future = scope.fork(() -> 
            chatClient.stream(new Prompt(prompt))
                .map(chunk -> chunk.getResult().getOutput().getContent())
        );
        
        scope.join();
        return future.resultNow();
    } catch (Exception e) {
        log.error("Error streaming AI response", e);
        return Flux.error(e);
    }
}

This allows you to display AI responses as they're being generated, similar to how ChatGPT works!

Using Aspect-Oriented Programming To Observe Interactions with AI

One of Spring's superpowers is AOP, and it works beautifully with Spring AI. Let's create an aspect to log all our AI interactions:

@Aspect
@Component
public class AILoggingAspect {
    
    private final Logger log = LoggerFactory.getLogger(AILoggingAspect.class);
    
    @Around("execution(* org.springframework.ai.client.AiClient.*(..)) && args(prompt,..)")
    public Object logAIInteraction(ProceedingJoinPoint pjp, Prompt prompt) throws Throwable {
        log.info("AI Request: {}", prompt.getContents());
        long startTime = System.currentTimeMillis();
        
        Object result = pjp.proceed();
        
        long duration = System.currentTimeMillis() - startTime;
        log.info("AI Response received in {}ms", duration);
        
        return result;
    }
}

To summarize

Spring AI is a powerful addition to the Spring ecosystem that makes integrating AI capabilities into your applications remarkably straightforward. True to Spring's philosophy, it abstracts away the complexities, letting you focus on building features rather than wrestling with AI provider APIs.

The environment-based model switching capability is particularly powerful, allowing you to:

• Switch providers based on deployment environments (dev, staging, production)
• A/B test different AI models to find the best fit for your use case
• Create fallback strategies when primary providers are unavailable
• Optimize for cost by using different models based on workload
• Support multi-tenant systems where different customers might use different AI providers

Combined with Java 21's virtual threads and structured concurrency, Spring AI applications can efficiently handle multiple concurrent AI requests without blocking valuable server resources.

Whether you're building a chatbot, implementing semantic search, or creating content generation tools, Spring AI provides the structure and simplicity that Java developers have come to expect from Spring projects.

‍

‍

‍