From 92773f17c041b2c140aa13039abe32fe62310c94 Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Thu, 21 Sep 2023 11:50:25 +0200 Subject: [PATCH] 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 --- .../embedding/AzureOpenAiEmbeddingClient.java | 14 ++- .../ai/embedding/EmbeddingClient.java | 4 + .../ai/embedding/EmbeddingUtil.java | 69 ++++++++++++++ .../embedding-model-dimensions.properties | 19 ++++ .../ai/embedding/EmbeddingUtilTest.java | 89 +++++++++++++++++++ .../embedding/OpenAiEmbeddingClient.java | 25 ++++-- .../ai/openai/embedding/EmbeddingIT.java | 1 + 7 files changed, 213 insertions(+), 8 deletions(-) create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingUtil.java create mode 100644 spring-ai-core/src/main/resources/embedding/embedding-model-dimensions.properties create mode 100644 spring-ai-core/src/test/java/org/springframework/ai/embedding/EmbeddingUtilTest.java diff --git a/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/embedding/AzureOpenAiEmbeddingClient.java b/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/embedding/AzureOpenAiEmbeddingClient.java index 1a02c412c..f936c3c7b 100644 --- a/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/embedding/AzureOpenAiEmbeddingClient.java +++ b/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 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(); + } + } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingClient.java b/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingClient.java index cbb44a5bc..64bb019d5 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingClient.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingClient.java @@ -14,4 +14,8 @@ public interface EmbeddingClient { EmbeddingResponse embedForResponse(List texts); + default int dimensions() { + return embed("Test String").size(); + } + } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingUtil.java b/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingUtil.java new file mode 100644 index 000000000..7cd1e7582 --- /dev/null +++ b/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 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 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); + } + } + +} diff --git a/spring-ai-core/src/main/resources/embedding/embedding-model-dimensions.properties b/spring-ai-core/src/main/resources/embedding/embedding-model-dimensions.properties new file mode 100644 index 000000000..13792315f --- /dev/null +++ b/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 diff --git a/spring-ai-core/src/test/java/org/springframework/ai/embedding/EmbeddingUtilTest.java b/spring-ai-core/src/test/java/org/springframework/ai/embedding/EmbeddingUtilTest.java new file mode 100644 index 000000000..7a26f2851 --- /dev/null +++ b/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 embed(String text) { + return List.of(0.1, 0.1, 0.1); + } + + @Override + public List embed(Document document) { + throw new UnsupportedOperationException("Unimplemented method 'embed'"); + } + + @Override + public List> embed(List texts) { + throw new UnsupportedOperationException("Unimplemented method 'embed'"); + } + + @Override + public EmbeddingResponse embedForResponse(List 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); + } + +} diff --git a/spring-ai-openai/src/main/java/org/springframework/ai/openai/embedding/OpenAiEmbeddingClient.java b/spring-ai-openai/src/main/java/org/springframework/ai/openai/embedding/OpenAiEmbeddingClient.java index 955d0ec75..a7ac0652c 100644 --- a/spring-ai-openai/src/main/java/org/springframework/ai/openai/embedding/OpenAiEmbeddingClient.java +++ b/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(); + } + } diff --git a/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java b/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java index 5011134b3..368af4a57 100644 --- a/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java +++ b/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); } }