A Guide to Embeddings Model API in Spring AI

1. Overview

Integrating artificial intelligence into an application typically involves working with text data. A critical technique in this domain is the embedding model, which converts textual information into embeddings that applications can process.

In this tutorial, we’ll explore the embeddings model API in Spring AI. This powerful API provides an abstraction that makes it easy to adopt different embedding models with minimal effort and facilitates our application in understanding text.

2. Introduction to Embeddings

To train AI models to learn the semantic meaning of text and images, we generally convert these data types into high-dimensional vector representations that are known as embeddings.

AI models understand the relationships between embeddings by calculating their similarity. When two embeddings have a higher similarity score, it means the contextual meaning of the text they represent is similar.

3. Embedding Model APIs

Spring AI provides a set of APIs to simplify our work with the embedding models. The APIs are interfaces that hide all the implementation details from us.

3.1. EmbeddingModel

An embedding model is a trained machine learning model that converts different kinds of objects, such as paragraphs and images, into a high-dimensional vector space.

There are different models, such as BERT, that are offered by different providers. Spring AI embeddings API provides an interface EmbeddingModel that encapsulates the details of adopting the embedding model:

public interface EmbeddingModel extends Model<EmbeddingRequest, EmbeddingResponse> {
    EmbeddingResponse call(EmbeddingRequest request);

    // constructors & other methods
}

The call() method simply accepts an EmbeddingRequest that contains the data source and sends it to the model provider, returning the EmbeddingResponse that contains the Embedding.

3.2. EmbeddingRequest

EmbeddingRequest contains the payload with a list of texts for converting to embeddings. Besides the texts, we could include additional options that are specifically for our EmbeddingModel:

public class EmbeddingRequest implements ModelRequest<List<String>> {
    private final List<String> inputs;
    private final EmbeddingOptions options;

    // constructors & other methods
}

3.3. EmbeddingResponse

The EmbeddingResponse encapsulates the response from the embedding model provider. It contains a list of Embedding objects and additional metadata, such as token usage:

public class EmbeddingResponse implements ModelResponse<Embedding> {
    private final List<Embedding> embeddings;
    private final EmbeddingResponseMetadata metadata;

    // constructors & other methods
}

3.4. Embedding

Embedding contains the vector representation in a float array. The dimensionality depends on the embedding model we’ve chosen, which normally varies from a few hundred to about thousands:

public class Embedding implements ModelResult<float[]> {
    private final float[] embedding;
    private final Integer index;
    private final EmbeddingResultMetadata metadata;

    // constructors & other methods
}

4. Integration With OpenAI

Spring AI supports OpenAI as one of the embedding model integrations. In this section, we’ll adopt OpenAI and create a Spring service to convert texts into embeddings.

4.1. Maven Dependency

Let us start by adding the following Spring AI Open AI dependency to our pom.xml:

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

4.2. Open AI Configuration

To complete the integration of OpenAI with Spring AI, we’ll need to put the API key for authentication to the OpenAI API:

spring:
  ai:
    openai:
      api-key: "<YOUR-API-KEY>"

4.3. Auto-configuration of the EmbeddingModel

Spring AI is capable of auto-configuring the EmbeddingModel. Let’s enable it by adding an additional property to define the embedding model in application.yml:

spring:
  ai:
    openai:
      embedding:
        options:
          model: "text-embedding-3-small"

This model property configures the embedding model we’re going to use. There are three different models currently provided by OpenAI.

Once we define this embedding model, we simply inject the EmbeddingModel into the Spring Boot service without specifying any OpenAI details. Everything relies on the Spring AI embedding APIs:

@Service
public class EmbeddingService {
    private final EmbeddingModel embeddingModel;

    public EmbeddingService(EmbeddingModel embeddingModel) {
        this.embeddingModel = embeddingModel;
    }

    public EmbeddingResponse getEmbeddings(String... texts) {
        EmbeddingRequest request = new EmbeddingRequest(Arrays.asList(texts), null);
        return embeddingModel.call(request);
    }
}

Auto-configuration provides us with convenience without exposing the actual embedding implementation. This enables us to switch to different implementations easily by merely updating our application.yml.

4.4. Manual Configuration of the EmbeddingModel

Although auto-configuration is convenient, it cannot provide us with flexibility in some cases, such as when our application requires generating embeddings using more than one embedding model or with different embedding model providers.

In this scenario, we manually define the embedding model producer in a configuration class:

@Configuration
public class EmbeddingConfig {
    @Bean
    public OpenAiApi openAiApi(@Value("${spring.ai.openai.api-key}") String apiKey) {
        return OpenAiApi.builder()
          .apiKey(apiKey)
          .build();
    }

    @Bean
    public OpenAiEmbeddingModel openAiEmbeddingModel(OpenAiApi openAiApi) {
        OpenAiEmbeddingOptions options = OpenAiEmbeddingOptions.builder()
          .model("text-embedding-3-small")
          .build();
        return new OpenAiEmbeddingModel(openAiApi, MetadataMode.EMBED, options);
    }
}

From our example, we first created an OpenAI client, OpenAiApi, using the injected API key openAiApiKey. We’ll then create the OpenAI embedding model using the client.

We update the service slightly to inject the OpenAIEmbeddingModel implementation instead of the EmbeddingModel interface:

@Service
public class ManualEmbeddingService {
    private final OpenAiEmbeddingModel openAiEmbeddingModel;

    public ManualEmbeddingService(OpenAiEmbeddingModel openAiEmbeddingModel) {
        this.openAiEmbeddingModel = openAiEmbeddingModel;
    }

    public EmbeddingResponse getEmbeddings(String... texts) {
        EmbeddingRequest request = new EmbeddingRequest(Arrays.asList(texts), null);
        return openAiEmbeddingModel.call(request);
    }
}

5. Testing the Embedding Service

Based on our auto-configuration service implementation in the previous section, we expose a REST endpoint that allows us to test the embedding service:

@RestController
public class EmbeddingController {
    private final EmbeddingService embeddingService;

    public EmbeddingController(EmbeddingService embeddingService) {
        this.embeddingService = embeddingService;
    }

    @PostMapping("/embeddings")
    public ResponseEntity<EmbeddingResponse> getEmbeddings(@RequestBody String text) {
        EmbeddingResponse response = embeddingService.getEmbeddings(text);
        return ResponseEntity.ok(response);
    }
}

Let’s make a request with texts in the body to this endpoint via curl:

$ curl -X POST http://localhost:8080/embeddings -H "Content-Type: text/plain" -d "Hello world"

We get the following response:

{
    "metadata": {
        "model": "text-embedding-3-small",
        "usage": {
            "promptTokens": 2,
            "completionTokens": 0,
            "totalTokens": 2,
            "nativeUsage": {
                "prompt_tokens": 48,
                "total_tokens": 48
            }
        },
        "empty": true
    },
    "result": {
        "index": 0,
        "metadata": {
            "modalityType": "TEXT",
            "documentId": "",
            "mimeType": {
                "type": "text",
                "subtype": "plain",
                "parameters": {},
                "charset": null,
                "concrete": true,
                "wildcardSubtype": false,
                "subtypeSuffix": null,
                "wildcardType": false
            },
            "documentData": null
        },
        "output": [
            -0.0020785425,
            -0.049085874,
            ...
       ]
    }
}

Notably, this isn’t a full response, as it’s very long. We’ve trimmed it to describe the two main top-level nodes in the JSON: metadata and result.

metadata provides information about the model used and resource usage during the embedding conversion. model indicates the OpenAI model we’ve picked, and totalTokens reveals the number of tokens consumed by the conversion.

result holds the embedding result. output in the result contains a float array, which is the embedding converted by the embedding model from our provided text.

6. Conclusion

The embeddings model API in Spring AI provides the abstraction layer and support for model providers like OpenAI, enabling us to incorporate it into our Java applications.

In this article, we adopted a reference embedding model, OpenAI, with auto-configuration for simplicity and a manual configuration for flexibility. The embeddings API provides the capability of converting texts into embedding vectors.