Fix Issue #211 : Improved Embedding Performance by Handling Base64 Encoding

[yoshioterada]

Overview

This commit includes the fix described in Issue #211.

• Addressed the issue where Base64 encoding could not be handled.
• Improved performance by using Base64 encoding by default.

Detail

This pull request introduces several changes to the Embedding class and related components in the openai-java-core package. The primary goal is to enhance the handling of embedding vectors by supporting both float lists and Base64-encoded strings. The most important changes include the introduction of the EmbeddingValue class, modifications to the Embedding class to use EmbeddingValue, and updates to the deserialization logic.

Enhancements to embedding handling:

• openai-java-core/src/main/kotlin/com/openai/models/embeddings/Embedding.kt: Modified the Embedding class to use EmbeddingValue instead of List<Double> for embedding vectors. This includes changes to the constructor, builder, and relevant methods. [1] [2] [3] [4] [5] [6]

Introduction of EmbeddingValue class:

• openai-java-core/src/main/kotlin/com/openai/models/embeddings/EmbeddingValue.kt: Added a new class EmbeddingValue to represent embedding vectors, which can be either a list of floats or a Base64-encoded string. This class includes methods for converting between these representations.

Deserialization improvements:

• openai-java-core/src/main/kotlin/com/openai/models/embeddings/EmbeddingValueDeserializer.kt: Introduced a custom deserializer EmbeddingValueDeserializer to handle the deserialization of EmbeddingValue objects from JSON, supporting both float arrays and Base64 strings.

Default encoding format:

• openai-java-core/src/main/kotlin/com/openai/models/embeddings/EmbeddingCreateParams.kt: Set the default EncodingFormat to BASE64 for performance improvements.

Test updates:

• openai-java-core/src/test/kotlin/com/openai/models/embeddings/CreateEmbeddingResponseTest.kt and openai-java-core/src/test/kotlin/com/openai/models/embeddings/EmbeddingTest.kt: Updated test cases to accommodate the changes in the Embedding class and the introduction of EmbeddingValue. [1] [2] [3]

This code will run look like following Java code.

This PR code will run with look like following code style.

package com.openai.example;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
import com.openai.azure.AzureOpenAIServiceVersion;
import com.openai.azure.credential.AzureApiKeyCredential;
import com.openai.client.OpenAIClient;
import com.openai.client.OpenAIClientAsync;
import com.openai.client.okhttp.OpenAIOkHttpClient;
import com.openai.client.okhttp.OpenAIOkHttpClientAsync;
import com.openai.models.embeddings.CreateEmbeddingResponse;
import com.openai.models.embeddings.EmbeddingCreateParams;
import com.openai.services.blocking.EmbeddingService;
import com.openai.models.embeddings.EmbeddingModel;

public final class EmbeddingsExample {
    private EmbeddingsExample() {}

    private static final String AZURE_OPENAI_ENDPOINT = "https://$AOAI.openai.azure.com";
    private static final String AZURE_OPENAI_KEY = $AOAI_KEY;

    private static OpenAIClient client;

    public static void main(String[] args) {
        client = OpenAIOkHttpClient.builder().baseUrl(AZURE_OPENAI_ENDPOINT)
                .credential(AzureApiKeyCredential.create(AZURE_OPENAI_KEY))
                .azureServiceVersion(AzureOpenAIServiceVersion.getV2024_02_15_PREVIEW()).build();
        EmbeddingsExample example = new EmbeddingsExample();
        example.basicSample();
        example.multipleDataSample();
        example.asyncSample();
    }

    // Basic Sample
    public void basicSample() {
        EmbeddingService embeddings = client.embeddings();
        String singlePoem = "In the quiet night, stars whisper secrets, dreams take flight.";

        // No specified format
        EmbeddingCreateParams embeddingCreateParams = EmbeddingCreateParams.builder()
                .input(singlePoem).model(EmbeddingModel.TEXT_EMBEDDING_3_SMALL.asString()).build();
        embeddings.create(embeddingCreateParams).data().forEach(embedding -> {
            // It will show both base64 and float embedding
            System.out.println("EMBEDDEING (Default non-format toString()) -------------"
                    + embedding.toString());
        });
        System.out.println("------------------------------------------------");

        // Specified FloatEmbedding format
        EmbeddingCreateParams embeddingCreateParams2 = EmbeddingCreateParams.builder()
                .input(singlePoem).model(EmbeddingModel.TEXT_EMBEDDING_3_SMALL.asString())
                .encodingFormat(EmbeddingCreateParams.EncodingFormat.FLOAT).build();
        embeddings.create(embeddingCreateParams).data().forEach(embedding -> {
            embedding.embedding().getFloatEmbedding().ifPresent(emb -> {
                System.out.println("EMBEDDEING FLOAT (Float Embedding) -------------" + emb);
            });
            embedding.embedding().getBase64Embedding().ifPresent(emb -> {
                System.out.println("EMBEDDEING BASE64  (Float Embedding) -------------" + emb);
            });
            System.out.println("------------------------------------------------");
        });

        // Specified Base64Embedding format
        EmbeddingCreateParams embeddingCreateParams3 = EmbeddingCreateParams.builder()
                .input(singlePoem).model(EmbeddingModel.TEXT_EMBEDDING_3_SMALL.asString())
                .encodingFormat(EmbeddingCreateParams.EncodingFormat.BASE64).build();
        embeddings.create(embeddingCreateParams).data().forEach(embedding -> {
            embedding.embedding().getBase64Embedding().ifPresent(emb -> {
                System.out.println("EMBEDDEING BASE64 (Base64 Embedding) -------------" + emb);
            });
            embedding.embedding().getFloatEmbedding().ifPresent(emb -> {
                System.out.println("EMBEDDEING FLOAT (Base64 Embedding)  -------------" + emb);
            });
            System.out.println("------------------------------------------------");

        });
    }

    // Multiple Data Sample
    public void multipleDataSample() {
        EmbeddingService embeddings = client.embeddings();

        getPoems().forEach(poem -> {
            System.out.println("POEM (START) -------------" + poem);

            EmbeddingCreateParams embeddingCreateParams = EmbeddingCreateParams.builder()
                    .input(poem).model(EmbeddingModel.TEXT_EMBEDDING_3_SMALL.asString()).build();
            embeddings.create(embeddingCreateParams).data().forEach(embedding -> {
                embedding.embedding().getFloatEmbedding().ifPresent(emb -> {
                    System.out.println("EMBEDDEING Float (by Dfault) -------------" + emb);
                });
            });
            System.out.println("POEM (END) -------------" + poem);
        });
    }

    // Async Sample
    public void asyncSample() {
        CountDownLatch latch = new CountDownLatch(1);
        try {
            OpenAIClientAsync asyncClient = OpenAIOkHttpClientAsync.builder()
                    .baseUrl(AZURE_OPENAI_ENDPOINT)
                    .credential(AzureApiKeyCredential.create(AZURE_OPENAI_KEY))
                    .azureServiceVersion(AzureOpenAIServiceVersion.getV2024_02_15_PREVIEW())
                    .build();
            CompletableFuture<CreateEmbeddingResponse> completableFuture = asyncClient.embeddings()
                    .create(EmbeddingCreateParams.builder()
                            .input("The quick brown fox jumped over the lazy dog")
                            .model(EmbeddingModel.TEXT_EMBEDDING_3_SMALL)
                            .encodingFormat(EmbeddingCreateParams.EncodingFormat.FLOAT)
                            .user("user-1234").build());

            completableFuture.thenAccept(response -> {
                response.validate();
                response.data().forEach(embedding -> {
                    System.out.println("EMBEDDEING (Async) -------------" + embedding.toString());
                    latch.countDown();
                });
            }).exceptionally(ex -> {
                System.err.println("Error: " + ex.getMessage());
                latch.countDown();
                return null;
            });
            latch.await();
            System.out.println("Latch count down completed");
            System.exit(0);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }

    private List<String> getPoems() {
        List<String> poems = new ArrayList<>();
        poems.add("In the quiet night, stars whisper secrets, dreams take flight.");
        poems.add("Beneath the moon's glow, shadows dance, hearts begin to know.");
        poems.add("Waves crash on the shore, time stands still, love forevermore.");
        poems.add("Autumn leaves fall, painting the ground, nature's final call.");
        poems.add("Morning dew glistens, a new day dawns, hope always listens.");
        poems.add("Mountains stand tall, silent guardians, witnessing it all.");
        poems.add("In a field of green, flowers bloom bright, a serene scene.");
        poems.add("Winter's chill bites, fireside warmth, cozy, long nights.");
        poems.add("Spring's gentle breeze, life awakens, hearts find ease.");
        poems.add("Sunset hues blend, day meets night, a perfect end.");
        return poems;
    }
}

[yoshioterada]

Dear @TomerAberbach san, If you have time, please review my PR please?

[TomerAberbach]

Some high-level thoughts:

1. It would be nice if end-users didn't have to pass EmbeddingValue. Prior to this PR a user could simply pass a Double, which was simpler. Can we preserve the Double signatures via overloads? We should probably also add associated String overloads for convenience.

2. Currently, all constructed params and response classes in the SDK are immutable. The EmbeddingValue class added in this PR is not immutable. It should be immutable to keep the same contract as before.

3. We should not expose the constructor of EmbeddingValue. We should make a public API similar to the other union classes (i.e. one-ofs) we have in the SDK. Tool is an example of a class that can be "one of" a few different things. You don't have to include all the convenience methods in Tool, but the general concept of having ofA, ofB, a, and b methods for a class that can be "a" or "b" is nice and we should copy it.