Add embedding model dimensions retrieval

 - Add dimension method to the EmbeddingClient interface. Default dimension implementation uses the embed method to produce results and counts the result dimensions. - Add EmbeddingUtil#dimensions utilities that look up the model dimensions from a pre-defined (static) file. If the requested model is unknown, fallback to the default behaviour. - Override the dimensions method in the OpenAiEmbeddingClient and AzureOpenAiEmbeddingClient to implement local caching. - Add unit and IT tests. Resolves #28

Author: tzolov

spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/embedding/AzureOpenAiEmbeddingClient.java

@@ -11,12 +11,14 @@
 import org.springframework.ai.embedding.Embedding;
 import org.springframework.ai.embedding.EmbeddingClient;
 import org.springframework.ai.embedding.EmbeddingResponse;
+import org.springframework.ai.embedding.EmbeddingUtil;
 import org.springframework.util.Assert;
 
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
 
 public class AzureOpenAiEmbeddingClient implements EmbeddingClient {
@@ -27,6 +29,8 @@ public class AzureOpenAiEmbeddingClient implements EmbeddingClient {
 
 private final String model;
 
+private final AtomicInteger embeddingDimensions = new AtomicInteger(-1);
+
 public AzureOpenAiEmbeddingClient(OpenAIClient azureOpenAiClient) {
 this(azureOpenAiClient, "text-embedding-ada-002");
 }
@@ -89,8 +93,6 @@ private Map<String, Object> generateMetadata(String model, EmbeddingsUsage embed
 Map<String, Object> metadata = new HashMap<>();
 metadata.put("model", model);
 metadata.put("prompt-tokens", embeddingsUsage.getPromptTokens());
-// NOTE, not in API of AzureAI - metadata.put("completion-tokens",
-// embeddingsUsage.getCompletionTokens());
 metadata.put("total-tokens", embeddingsUsage.getTotalTokens());
 return metadata;
 }
@@ -106,4 +108,12 @@ private List<Embedding> generateEmbeddingList(List<EmbeddingItem> nativeData) {
 return data;
 }
 
+@Override
+public int dimensions() {
+if (this.embeddingDimensions.get() < 0) {
+this.embeddingDimensions.set(EmbeddingUtil.dimensions(this, this.model));
+}
+return this.embeddingDimensions.get();
+}
+
 }

spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingClient.java

@@ -14,4 +14,8 @@ public interface EmbeddingClient {
 
 EmbeddingResponse embedForResponse(List<String> texts);
 
+default int dimensions() {
+return embed("Test String").size();
+}
+
 }

spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingUtil.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright 2023-2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.ai.embedding;
+
+import java.io.IOException;
+import java.util.Map;
+import java.util.Properties;
+import java.util.stream.Collectors;
+
+import org.springframework.core.io.DefaultResourceLoader;
+
+/**
+ * @author Christian Tzolov
+ */
+public class EmbeddingUtil {
+
+private static Map<String, Integer> KNOWN_EMBEDDING_DIMENSIONS = loadKnownModelDimensions();
+
+/**
+ * Return the dimension of the requested embedding model name. If the model name is
+ * unknown uses the EmbeddingClient to perform a dummy EmbeddingClient#embed and count
+ * the response dimensions.
+ * @param embeddingClient Fall-back client to determine, empirically the dimensions.
+ * @param modelName Embedding model name to retrieve the dimensions for.
+ * @return Returns the embedding dimensions for the modelName.
+ */
+public static int dimensions(EmbeddingClient embeddingClient, String modelName) {
+
+if (KNOWN_EMBEDDING_DIMENSIONS.containsKey(modelName)) {
+// Retrieve the dimension from a pre-configured file.
+return KNOWN_EMBEDDING_DIMENSIONS.get(modelName);
+}
+else {
+// Determine the dimensions empirically.
+// Generate an embedding and count the dimension size;
+return embeddingClient.embed("Test String").size();
+}
+}
+
+private static Map<String, Integer> loadKnownModelDimensions() {
+try {
+Properties properties = new Properties();
+properties.load(new DefaultResourceLoader()
+.getResource("classpath:/embedding/embedding-model-dimensions.properties")
+.getInputStream());
+return properties.entrySet()
+.stream()
+.collect(Collectors.toMap(e -> e.getKey().toString(), e -> Integer.parseInt(e.getValue().toString())));
+}
+catch (IOException e) {
+throw new RuntimeException(e);
+}
+}
+
+}

spring-ai-core/src/main/resources/embedding/embedding-model-dimensions.properties

@@ -0,0 +1,19 @@
+# Map of embedding model names and their dimesions
+# OpenAI
+text-embedding-ada-002=1536
+text-similarity-ada-001=1024
+text-similarity-babbage-001=2048
+text-similarity-curie-001=4096
+text-similarity-davinci-001=12288
+text-search-ada-doc-001=1024
+text-search-ada-query-001=1024
+text-search-babbage-doc-001=2048
+text-search-babbage-query-001=2048
+text-search-curie-doc-001=4096
+text-search-curie-query-001=4096
+text-search-davinci-doc-001=12288
+text-search-davinci-query-001=12288
+code-search-ada-code-001=1024
+code-search-ada-text-001=1024
+code-search-babbage-code-001=2048
+code-search-babbage-text-001=2048

spring-ai-core/src/test/java/org/springframework/ai/embedding/EmbeddingUtilTest.java

@@ -0,0 +1,89 @@
+/*
+ * Copyright 2023-2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.ai.embedding;
+
+import java.util.List;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvFileSource;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+import org.springframework.ai.document.Document;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/**
+ * @author Christian Tzolov
+ */
+@ExtendWith(MockitoExtension.class)
+public class EmbeddingUtilTest {
+
+@Mock
+private EmbeddingClient embeddingClient;
+
+@Test
+public void testDefaultMethodImplementation() {
+
+EmbeddingClient dummy = new EmbeddingClient() {
+
+@Override
+public List<Double> embed(String text) {
+return List.of(0.1, 0.1, 0.1);
+}
+
+@Override
+public List<Double> embed(Document document) {
+throw new UnsupportedOperationException("Unimplemented method 'embed'");
+}
+
+@Override
+public List<List<Double>> embed(List<String> texts) {
+throw new UnsupportedOperationException("Unimplemented method 'embed'");
+}
+
+@Override
+public EmbeddingResponse embedForResponse(List<String> texts) {
+throw new UnsupportedOperationException("Unimplemented method 'embedForResponse'");
+}
+};
+
+assertThat(dummy.dimensions()).isEqualTo(3);
+}
+
+@ParameterizedTest
+@CsvFileSource(resources = "/embedding/embedding-model-dimensions.properties", numLinesToSkip = 1, delimiter = '=')
+public void testKnownEmbeddingModelDimensions(String model, String dimension) {
+assertThat(EmbeddingUtil.dimensions(embeddingClient, model)).isEqualTo(Integer.valueOf(dimension));
+verify(embeddingClient, never()).embed(any(String.class));
+verify(embeddingClient, never()).embed(any(Document.class));
+}
+
+@Test
+public void testUnknownModelDimension() {
+when(embeddingClient.embed(eq("Test String"))).thenReturn(List.of(0.1, 0.1, 0.1));
+assertThat(EmbeddingUtil.dimensions(embeddingClient, "unknown_model")).isEqualTo(3);
+}
+
+}

spring-ai-openai/src/main/java/org/springframework/ai/openai/embedding/OpenAiEmbeddingClient.java

@@ -1,22 +1,25 @@
 package org.springframework.ai.openai.embedding;
 
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+
 import com.theokanning.openai.Usage;
 import com.theokanning.openai.embedding.EmbeddingRequest;
 import com.theokanning.openai.service.OpenAiService;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+
 import org.springframework.ai.document.Document;
 import org.springframework.ai.embedding.Embedding;
 import org.springframework.ai.embedding.EmbeddingClient;
 import org.springframework.ai.embedding.EmbeddingResponse;
+import org.springframework.ai.embedding.EmbeddingUtil;
 import org.springframework.util.Assert;
 
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.stream.Collectors;
-
 public class OpenAiEmbeddingClient implements EmbeddingClient {
 
 private static final Logger logger = LoggerFactory.getLogger(OpenAiEmbeddingClient.class);
@@ -25,6 +28,8 @@ public class OpenAiEmbeddingClient implements EmbeddingClient {
 
 private final String model;
 
+private final AtomicInteger embeddingDimensions = new AtomicInteger(-1);
+
 public OpenAiEmbeddingClient(OpenAiService openAiService) {
 this(openAiService, "text-embedding-ada-002");
 }
@@ -95,4 +100,12 @@ private Map<String, Object> generateMetadata(String model, Usage usage) {
 return metadata;
 }
 
+@Override
+public int dimensions() {
+if (this.embeddingDimensions.get() < 0) {
+this.embeddingDimensions.set(EmbeddingUtil.dimensions(this, this.model));
+}
+return this.embeddingDimensions.get();
+}
+
 }

spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java

@@ -28,6 +28,7 @@ void simpleEmbedding() {
 assertThat(embeddingResponse.getMetadata()).containsEntry("total-tokens", 2L);
 assertThat(embeddingResponse.getMetadata()).containsEntry("prompt-tokens", 2L);
 
+assertThat(embeddingClient.dimensions()).isEqualTo(1536);
 }
 
 }