Exploring Embeddings API With Java and Spring AI

Hi community! 

This is my second article in a series of introductions to Spring AI. You may find the first one, where I explained how to generate images using Spring AI and OpenAI DALL-E 3 models, here. Today, we will create simple applications using embeddings API and Spring AI. 

In this article, I’ll skip the explanation of some basic Spring concepts like bean management, starters, etc, as the main goal of this article is to discover Spring AI capabilities. For the same reason, I’ll not create detailed instructions on generating the OpenAI API Key. In case you don’t have one, follow the links in Step 0, which should give you enough context on how to create one.

The code I will share in this article will also be available in the GitHub repo. You may find this repo useful because to make this article shorter, I’ll not paste here some precalculated values and simple POJOs.

What Are Embeddings?

Before we start code implementation, let's discuss what embeddings are.

In the Spring AI documentation, we can find the following definition of embeddings:

Embeddings are numerical representations of text, images, or videos that capture relationships between inputs.

Embeddings convert text, images, and video into arrays of floating-point numbers called vectors. These vectors are designed to capture the meaning of the text, images, and videos. The length of the embedding array is called the vector’s dimensionality.

Key points we should pay attention to:

1. Numerical representation of text (also applicable for images and videos, but let’s focus just on texts in this article)
2. Embeddings are vectors. And as every vector has coordinates for each dimension it exists, we should think about embeddings as a coordinate of our input in “Text Universe”

As with every other vector, we can find the distance between two embeddings. The closer the two embeddings are to each other, the more similar their context. We will use this approach in our application. 

Determining the Scope of Our Future Application

Let’s imagine that we have an online shop with different electronics. Every single item has its ID and description. We need to create a module that will receive users' input describing the item the user wants to find or buy and return five of the most relevant products to this query.

We will achieve this goal using embeddings. The following are steps we need to implement:

1. We will fetch embeddings (vector representation) of our existing products and store them. I’ll not show this step in this article, because it will be similar to one we will explore later. But you can find precalculated embeddings to use in your code in the GitHub repo I previously shared.
2. We will call the embeddings API for each user input.
3. We will be comparing user input embeddings with precalculated embeddings of our item description. We will leverage the Cosine Similarity approach to find the closest vectors.

Implementation

Step 0: Generate an Open API Key 

If you don’t have an active OpenAI API key, do the following steps:

1. Create an account on the OpenAI signup page 
2. Generate the token on the API Keys page

Step 1: Set Up a Project

To quickly generate a project template with all necessary dependencies, one may use https://start.spring.io/.

In my example, I’ll be using Java 17 and Spring Boot 3.4.1. Also, we need to include the following dependency:

OpenAI

This dependency provides us with smooth integration with OpenAI just by writing a couple lines of code and a few lines of configurations.

After clicking generate, open the downloaded files in the IDE you are working on and validate that all necessary dependencies exist in pom.xml.

<dependency>
  <groupId>org.springframework.ai</groupId>
  <artifactId>spring-ai-openai-spring-boot-starter</artifactId>
</dependency>

<dependencyManagement>
  <dependencies>
     <dependency>
        <groupId>org.springframework.ai</groupId>
        <artifactId>spring-ai-bom</artifactId>
        <version>1.0.0-M4</version>
        <type>pom</type>
        <scope>import</scope>
     </dependency>
  </dependencies>
</dependencyManagement>

At the moment of writing this article, Spring AI version 1.0.0-M4 has not yet been published in the central Maven repository and is only available in the Spring repository. That’s why we need to add a link to that repo in our pom.xml as well:

<repositories>
  <repository>
     <id>spring-milestones</id>
     <name>Spring Milestones</name>
     <url>https://repo.spring.io/milestone</url>
     <snapshots>
        <enabled>false</enabled>
     </snapshots>
  </repository>
</repositories>

Step 2: Set Up the Configuration File

As a next step, we need to configure our property file. By default, Spring uses application.yaml or application.properties file. In this example, I’m using the YAML format. You may reformat code into .properties if you feel more comfortable working with this format.

Here are all the configs we need to add to the application.yaml file:

spring:
 application:
   name: aiembeddings
 ai:
   openai:
     api-key: [your OpenAI api key]
     embedding:
       options:
         model: text-embedding-ada-002

Model

The model to use. We will be using text-embedding-ada-002. There are other options: text-embedding-3-large, text-embedding-3-small. You may learn more about the differences in models in OpenAI docs.

As the main purpose of this article is to show the ease of Spring AI integration with embedding models, we will not go deeper into other configurations. You may find more config options in the Spring docs.

Step 3: Create Resource Files

Let’s create two files in the resource folder. 

The first one is the JSON-formatted “database” of items in our shop. Every Item will have the following parameters: Id, Name, and Description. I named this file samples.json and saved it in the resource folder. 

[
 {
   "id": 1,
   "name": "Smartphone A",
   "description": "5G smartphone with AMOLED display and 128GB storage."
 },
 {
   "id": 2,
   "name": "Smartphone B",
   "description": "4G smartphone with IPS screen and 64GB storage."
 },
 {
   "id": 3,
   "name": "Wireless Headphones",
   "description": "Bluetooth headphones with active noise cancellation."
 },
 {
   "id": 4,
   "name": "Smartwatch X",
   "description": "Fitness smartwatch with heart rate monitor and AMOLED display."
 },
 {
   "id": 5,
   "name": "Tablet Pro",
   "description": "10-inch tablet with 4GB RAM and 64GB storage."
 },
 {
   "id": 6,
   "name": "Bluetooth Speaker",
   "description": "Portable speaker with 12-hour battery life and waterproof design."
 },
 {
   "id": 7,
   "name": "Gaming Laptop",
   "description": "High-performance laptop with RTX 3060 GPU and 16GB RAM."
 },
 {
   "id": 8,
   "name": "External SSD",
   "description": "1TB external SSD with USB-C for fast data transfer."
 },
 {
   "id": 9,
   "name": "4K Monitor",
   "description": "27-inch monitor with 4K resolution and HDR support."
 },
 {
   "id": 10,
   "name": "Wireless Mouse",
   "description": "Ergonomic wireless mouse with adjustable DPI."
 }
]

The second one is a list of embeddings of the product description. I executed embeddings API in a separate application and saved responses for every single product into a separate file, embeddings.json. I’ll not share the whole file here, as it will make the article unreadable, but you still can download it from the GitHub repo of this project I shared at the beginning of the article.

Step 4: Create Embeddings Service

Now, let’s create the main service of our application -> embedding service.

To integrate our application with the embeddings API, we need to autowire EmbeddingModel. We have already configured OpenAI embeddings in the application.yaml. Spring Boot will automatically create and configure the instance (Bean) of EmbeddingModel. 

To fetch embeddings for a particular String or text, we just need to write one line of code:

EmbeddingResponse embeddingResponse = embeddingModel.embedForResponse(List.of(text));

Let’s see what the whole service looks like:

@Service
public class EmbeddingsService {

   private static List<Product> productList = new ArrayList<>();
   private static Map<Integer, float[]> embeddings = new HashMap<>();

   @Autowired
   private EmbeddingModel embeddingModel;

   @Autowired
   private SimilarityCalculator similarityCalculator;

   @PostConstruct
   public void initProducts() throws IOException {
       ObjectMapper objectMapper = new ObjectMapper();
       InputStream inputStream = getClass().getClassLoader().getResourceAsStream("samples.json");
       if (inputStream != null) {
           // map JSON into List<Product>
           productList = objectMapper.readValue(inputStream, new TypeReference<List<Product>>() {
           });
           System.out.println("Products loaded: List size = " + productList.size());

       } else {
           System.out.println("File samples.json not found in resources.");
       }
       embeddings = loadEmbeddingsFromFile();
   }

   public Map<Integer, float[]> loadEmbeddingsFromFile() {
       try {
           InputStream inputStream = getClass().getClassLoader().getResourceAsStream("embeddings.json");
           ObjectMapper objectMapper = new ObjectMapper();
           return objectMapper.readValue(inputStream, new TypeReference<Map<Integer, float[]>>() {
           });
       } catch (Exception e) {
           System.err.println("Error loading embeddings from file: " + e.getMessage());
           return null;
       }
   }

   public void getSimilarProducts(String query) {
       EmbeddingResponse embeddingResponse = embeddingModel.embedForResponse(List.of(query));
       List<ProductSimilarity> topSimilarProducts = similarityCalculator.findTopSimilarProducts(embeddingResponse.getResult().getOutput(),
               embeddings,
               productList,
               5);
       for (ProductSimilarity ps : topSimilarProducts) {
           System.out.printf("Product ID: %d, Name: %s, Description: %s, Similarity: %.4f%n",
                   ps.getProduct().getId(),
                   ps.getProduct().getName(),
                   ps.getProduct().getDescription(),
                   ps.getSimilarity());
       }
   }
}

Let’s deep dive into this code:

1. In the @postconstruct method, we are loading our resources into collections. The list of Products reads our products from samples.json. The product is a POJO with ID, name, and description fields. We also load precalculated embeddings of our products from another file  embeddings.json. We will need these embeddings later when we look for the most similar product.
2. The most important method in our service is getSimilarProducts which will receive user queries, fetch its embeddings using embeddingModel, and calculate similarities with our existing products. We will take a closer look at similarityCalculator.findTopSimilarProductsa little bit later in this article. After receiving a list of similarities, we will print the top N similar products in the following format:
   • Product ID, Name, Description, Similarity (a number between 0 and 1)

To calculate similarities, we introduced SimilarityCalculator Service. Let’s take a deeper look at its implementation.

@Service
public class SimilarityCalculator {

   public float calculateCosineSimilarity(float[] vectorA, float[] vectorB) {
       float dotProduct = 0.0f;
       float normA = 0.0f;
       float normB = 0.0f;

       for (int i = 0; i < vectorA.length; i++) {
           dotProduct += vectorA[i] * vectorB[i];
           normA += Math.pow(vectorA[i], 2);
           normB += Math.pow(vectorB[i], 2);
       }

       return (float) (dotProduct / (Math.sqrt(normA) * Math.sqrt(normB)));
   }

   public List<ProductSimilarity> findTopSimilarProducts(
           float[] queryEmbedding,
           Map<Integer, float[]> embeddings,
           List<Product> products,
           int topN) {

       List<ProductSimilarity> similarities = new ArrayList<>();

       for (Product product : products) {
           float[] productEmbedding = embeddings.get(product.getId());
           if (productEmbedding != null) {
               float similarity = calculateCosineSimilarity(queryEmbedding, productEmbedding);
               similarities.add(new ProductSimilarity(product, similarity));
           }
       }

       return similarities.stream()
               .sorted((p1, p2) -> Double.compare(p2.getSimilarity(), p1.getSimilarity()))
               .limit(topN)
               .toList();
   }
}

1. ProductSimilarity is a POJO class containing Product and similarity fields. You can find the code for this class in the GitHub repo.
2. calculateCosineSimilarity is the method used to find the most similar descriptions to user queries. Cosine similarity is one of the most popular ways to measure the similarity between embeddings. Explaining the exact workings of cosine similarity is beyond the scope of this article.
3. findTopSimilarProducts is a method called from our embedding service. It calculates similarities with all products, sorts them, and returns the top N products with the highest similarity.

Step 5: Execute Application

We will execute this application directly from the code, without using REST controllers or making API calls. You may use an approach similar to the one I used in the first article, when I created a separate endpoint to make execution smoother.

@SpringBootApplication
public class AiEmbeddingsApplication {
  public static void main(String[] args) {
     ConfigurableApplicationContext run = new SpringApplicationBuilder(AiEmbeddingsApplication.class)
           .web(WebApplicationType.NONE)
           .run(args);
     run.getBean(EmbeddingsService.class).getSimilarProducts("5G Phone. IPS");
  }
}

We are executing our code in the last line of the method, fetching the bean from the context, and executing the getSimilarProducts method with a provided query. 

In my query, I’ve included three keywords: 5G, Phone, and IPS. We should receive quite a high similarity with Product 1 and Product 2. Both products are smartphones, but Product 1 is a 5G smartphone, while Product 2 has an IPS screen. 

To start our application, we need to run the following command:

• mvn spring-boot:run

In a couple of seconds after executing, we may see the following result in the console:

Product ID: 2, Name: Smartphone B, Description: 4G smartphone with IPS screen and 64GB storage., Similarity: 0,9129
Product ID: 1, Name: Smartphone A, Description: 5G smartphone with AMOLED display and 128GB storage., Similarity: 0,8843
Product ID: 9, Name: 4K Monitor, Description: 27-inch monitor with 4K resolution and HDR support., Similarity: 0,8156
Product ID: 5, Name: Tablet Pro, Description: 10-inch tablet with 4GB RAM and 64GB storage., Similarity: 0,8156
Product ID: 4, Name: Smartwatch X, Description: Fitness smartwatch with heart rate monitor and AMOLED display., Similarity: 0,8037

We can see that Product 2 and Product 1 have the biggest similarities. What’s also interesting is that as we included the IPS keyword, our top five similar products all are products with displays. 

Conclusion

Spring AI is a great tool that helps developers smoothly integrate with different AI models. At the moment of writing this article, Spring AI supports 10 embedding models, including but not limited to Ollama and Amazon Bedrock.

I hope you found this article helpful and that it will inspire you to explore Spring AI deeper.