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

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

@@ -11,12 +11,14 @@ 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.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 @@ public class AzureOpenAiEmbeddingClient implements EmbeddingClient {
 		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 @@ public class AzureOpenAiEmbeddingClient implements EmbeddingClient {
 		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 @@ public class OpenAiEmbeddingClient implements EmbeddingClient {
 		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 @@ class EmbeddingIT {
 		assertThat(embeddingResponse.getMetadata()).containsEntry("total-tokens", 2L);
 		assertThat(embeddingResponse.getMetadata()).containsEntry("prompt-tokens", 2L);
 
+		assertThat(embeddingClient.dimensions()).isEqualTo(1536);
 	}
 
 }