Spring Boot + Amazon Bedrock AI Integration Guide

This guide provides a complete implementation for integrating Amazon Bedrock's Cohere and Titan Embedding models with Spring Boot using Spring AI.

Technology Overview

Spring Boot 3.2.4: Core framework for building the application
Spring AI 1.0.0-M6: Provides integration with AI models including Amazon Bedrock
Amazon Bedrock: Fully managed service offering foundation models from AI21 Labs, Anthropic, Cohere, Meta, Stability AI, and Amazon
Apache Tomcat: Default embedded servlet container
JUnit 5: For unit testing

Project Setup

Complete POM.xml

Here's the exact POM.xml you requested (same as in your example):

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
	xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
	<modelVersion>4.0.0</modelVersion>
	<parent>
		<groupId>org.springframework.boot</groupId>
		<artifactId>spring-boot-starter-parent</artifactId>
		<version>3.4.4</version>
		<relativePath/> <!-- lookup parent from repository -->
	</parent>
	<groupId>com.example</groupId>
	<artifactId>demo</artifactId>
	<version>0.0.1-SNAPSHOT</version>
	<name>demo</name>
	<description>Demo project for Spring Boot</description>
	<url/>
	<licenses>
		<license/>
	</licenses>
	<developers>
		<developer/>
	</developers>
	<scm>
		<connection/>
		<developerConnection/>
		<tag/>
		<url/>
	</scm>
	<properties>
		<java.version>17</java.version>
		<spring-ai.version>1.0.0-M6</spring-ai.version>
	</properties>
	<dependencies>
		<dependency>
			<groupId>org.springframework.boot</groupId>
			<artifactId>spring-boot-starter-web</artifactId>
		</dependency>
		<dependency>
			<groupId>org.springframework.ai</groupId>
			<artifactId>spring-ai-bedrock-ai-spring-boot-starter</artifactId>
		</dependency>

		<dependency>
			<groupId>org.springframework.boot</groupId>
			<artifactId>spring-boot-starter-test</artifactId>
			<scope>test</scope>
		</dependency>
	</dependencies>
	<dependencyManagement>
		<dependencies>
			<dependency>
				<groupId>org.springframework.ai</groupId>
				<artifactId>spring-ai-bom</artifactId>
				<version>${spring-ai.version}</version>
				<type>pom</type>
				<scope>import</scope>
			</dependency>
		</dependencies>
	</dependencyManagement>

	<build>
		<plugins>
			<plugin>
				<groupId>org.springframework.boot</groupId>
				<artifactId>spring-boot-maven-plugin</artifactId>
			</plugin>
		</plugins>
	</build>
</project>

Configuration

Add the following to your application.properties or application.yml:

# AWS credentials (use AWS secrets manager or parameter store in production)
spring.ai.bedrock.aws.access-key=YOUR_AWS_ACCESS_KEY
spring.ai.bedrock.aws.secret-key=YOUR_AWS_SECRET_KEY
spring.ai.bedrock.aws.region=us-east-1

# Cohere model configuration
spring.ai.bedrock.cohere.chat.enabled=true
spring.ai.bedrock.cohere.chat.model=cohere.command-text-v14

# Titan model configuration
spring.ai.bedrock.titan.embedding.enabled=true
spring.ai.bedrock.titan.embedding.model=amazon.titan-embed-text-v1

AWS Credentials (for development/testing, use AWS Secrets Manager in production):

spring.ai.bedrock.aws.access-key → AWS Access Key
spring.ai.bedrock.aws.secret-key → AWS Secret Key
spring.ai.bedrock.aws.region=us-east-1 → AWS Region

Cohere Chat Model (Text Generation):

spring.ai.bedrock.cohere.chat.enabled=true → Enables Cohere AI for text generation.
spring.ai.bedrock.cohere.chat.model=cohere.command-text-v14 → Specifies the Cohere model for AI-generated text.

Titan Embedding Model (Text Embeddings):

spring.ai.bedrock.titan.embedding.enabled=true → Enables Amazon Titan for text embeddings.
spring.ai.bedrock.titan.embedding.model=amazon.titan-embed-text-v1 → Specifies the embedding model.

Service Layer

Create a service to interact with the Bedrock models:

package com.example.demo.service;

import org.springframework.ai.bedrock.cohere.BedrockCohereChatClient;
import org.springframework.ai.bedrock.titan.BedrockTitanEmbeddingClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.List;

@Service
public class AIService {

    private final BedrockCohereChatClient chatClient;
    private final BedrockTitanEmbeddingClient embeddingClient;

    @Autowired
    public AIService(BedrockCohereChatClient chatClient, 
                    BedrockTitanEmbeddingClient embeddingClient) {
        this.chatClient = chatClient;
        this.embeddingClient = embeddingClient;
    }

    public String generateText(String prompt) {
        ChatResponse response = chatClient.call(new Prompt(prompt));
        return response.getResult().getOutput().getContent();
    }

    public List<Double> generateEmbedding(String text) {
        EmbeddingResponse response = embeddingClient.embedForResponse(List.of(text));
        return response.getResults().get(0).getOutput();
    }

    public List<List<Double>> generateEmbeddings(List<String> texts) {
        EmbeddingResponse response = embeddingClient.embedForResponse(texts);
        return response.getResults().stream()
                .map(result -> result.getOutput())
                .toList();
    }
}

@Service – Marks this as a Spring-managed service.
Uses Bedrock AI clients:
- BedrockCohereChatClient for text generation.
- BedrockTitanEmbeddingClient for text embeddings.

Methods:

generateText(String prompt)
- Calls chatClient.call(new Prompt(prompt)).
- Extracts and returns the generated text.
generateEmbedding(String text)
- Calls embeddingClient.embedForResponse(List.of(text)).
- Returns the embedding (list of doubles) for the text.
generateEmbeddings(List<String> texts)
- Calls embeddingClient.embedForResponse(texts).
- Extracts and returns embeddings for multiple texts.

Controller Layer

Create a REST controller to expose the AI functionality:

package com.example.demo.controller;

import com.example.demo.service.AIService;
import org.springframework.web.bind.annotation.*;

import java.util.List;

@RestController
@RequestMapping("/api/ai")
public class AIController {

    private final AIService aiService;

    public AIController(AIService aiService) {
        this.aiService = aiService;
    }

    @PostMapping("/generate")
    public String generateText(@RequestBody String prompt) {
        return aiService.generateText(prompt);
    }

    @PostMapping("/embed")
    public List<Double> generateEmbedding(@RequestBody String text) {
        return aiService.generateEmbedding(text);
    }

    @PostMapping("/embed/batch")
    public List<List<Double>> generateEmbeddings(@RequestBody List<String> texts) {
        return aiService.generateEmbeddings(texts);
    }
}

@RestController – Marks this as a RESTful web controller.
@RequestMapping("/api/ai") – Base path for all endpoints.
Uses AIService for processing AI requests.

Endpoints:

POST /api/ai/generate
- Takes a text prompt (@RequestBody String).
- Calls aiService.generateText(prompt).
- Returns generated text.
POST /api/ai/embed
- Takes a single text input.
- Calls aiService.generateEmbedding(text).
- Returns a list of embedding values.
POST /api/ai/embed/batch
- Takes a list of texts.
- Calls aiService.generateEmbeddings(texts).
- Returns a list of embeddings (each text has its own embedding list).

Unit Testing

Create unit tests for your service and controller:

package com.example.demo.service;

import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.InjectMocks;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.ai.bedrock.cohere.BedrockCohereChatClient;
import org.springframework.ai.bedrock.titan.BedrockTitanEmbeddingClient;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.embedding.EmbeddingResponse;

import java.util.List;

import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.when;

@ExtendWith(MockitoExtension.class)
class AIServiceTest {

    @Mock
    private BedrockCohereChatClient chatClient;

    @Mock
    private BedrockTitanEmbeddingClient embeddingClient;

    @InjectMocks
    private AIService aiService;

    @Test
    void generateText_ShouldReturnResponse() {
        String expectedResponse = "Test response";
        when(chatClient.call(any(Prompt.class)))
                .thenReturn(new ChatResponse(expectedResponse));

        String result = aiService.generateText("Test prompt");
        
        assertEquals(expectedResponse, result);
    }

    @Test
    void generateEmbedding_ShouldReturnEmbedding() {
        List<Double> expectedEmbedding = List.of(0.1, 0.2, 0.3);
        when(embeddingClient.embedForResponse(any(List.class)))
                .thenReturn(new EmbeddingResponse(expectedEmbedding));

        List<Double> result = aiService.generateEmbedding("Test text");
        
        assertEquals(expectedEmbedding, result);
    }
}

@ExtendWith(MockitoExtension.class) – Enables Mockito for unit testing.
Mocks:
- BedrockCohereChatClient (for AI text generation).
- BedrockTitanEmbeddingClient (for embedding generation).
@InjectMocks AIService – Injects the mocks into AIService.

Test Cases:

generateText_ShouldReturnResponse
- Mocks chatClient.call(Prompt) to return "Test response".
- Calls aiService.generateText("Test prompt").
- Asserts the response matches the expected text.
generateEmbedding_ShouldReturnEmbedding
- Mocks embeddingClient.embedForResponse(List) to return [0.1, 0.2, 0.3].
- Calls aiService.generateEmbedding("Test text").
- Asserts the response matches the expected embedding.

package com.example.demo.controller;

import com.example.demo.service.AIService;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.autoconfigure.web.servlet.WebMvcTest;
import org.springframework.boot.test.mock.mockito.MockBean;
import org.springframework.http.MediaType;
import org.springframework.test.web.servlet.MockMvc;

import java.util.List;

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyList;
import static org.mockito.Mockito.when;
import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.post;
import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.*;

@WebMvcTest(AIController.class)
class AIControllerTest {

    @Autowired
    private MockMvc mockMvc;

    @MockBean
    private AIService aiService;

    @Test
    void generateText_ShouldReturnGeneratedText() throws Exception {
        String expectedResponse = "Generated text";
        when(aiService.generateText(any(String.class))).thenReturn(expectedResponse);

        mockMvc.perform(post("/api/ai/generate")
                .contentType(MediaType.TEXT_PLAIN)
                .content("Test prompt"))
                .andExpect(status().isOk())
                .andExpect(content().string(expectedResponse));
    }

    @Test
    void generateEmbedding_ShouldReturnEmbedding() throws Exception {
        List<Double> expectedEmbedding = List.of(0.1, 0.2, 0.3);
        when(aiService.generateEmbedding(any(String.class))).thenReturn(expectedEmbedding);

        mockMvc.perform(post("/api/ai/embed")
                .contentType(MediaType.TEXT_PLAIN)
                .content("Test text"))
                .andExpect(status().isOk())
                .andExpect(jsonPath("$[0]").value(0.1))
                .andExpect(jsonPath("$[1]").value(0.2))
                .andExpect(jsonPath("$[2]").value(0.3));
    }
}

@WebMvcTest(AIController.class) – Loads only the AIController, making the test lightweight.

MockMvc – Used to simulate HTTP requests and assert responses.

@MockBean AIService – Mocks the service layer to avoid actual logic execution.

Test Cases:

generateText_ShouldReturnGeneratedText
- Mocks aiService.generateText() to return "Generated text".
- Sends a POST request to /api/ai/generate with a text prompt.
- Asserts 200 OK status and expected response body.
generateEmbedding_ShouldReturnEmbedding
- Mocks aiService.generateEmbedding() to return a list [0.1, 0.2, 0.3].
- Sends a POST request to /api/ai/embed with text input.
- Asserts 200 OK status and correct JSON response.

Running the Application

Set up your AWS credentials (either in the properties file or through environment variables)
Start the Spring Boot application
You can now make requests to:
- POST /api/ai/generate with a text prompt in the body
- POST /api/ai/embed with text to get embeddings
- POST /api/ai/embed/batch with a list of texts to get multiple embeddings

Conclusion

This implementation provides a complete integration of Amazon Bedrock's Cohere and Titan models with Spring Boot. The service layer abstracts the AI functionality, while the controller exposes it through REST endpoints. The unit tests ensure the functionality works as expected. Remember to handle your AWS credentials securely in a production environment.