Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Dynamic embedding dimensions resolution

Browse Source 
Leverage #28 to allow the vector stores to resolve the embedding dimensions dynamically.
 The explicitly set dimensions (if set) precedence over other configurations.
 If the embedding dimensions are not explicitly set, the embeddingClient is used
 to determine  them dynamically. If the client fails the it falls back to 1536.
This commit is contained in:
Christian Tzolov
2023-09-28 08:57:09 +02:00
parent cdaac4c532
commit 65a6ffcffb
8 changed files with 237 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,5 +1,11 @@
package org.springframework.ai.azure.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 com.azure.ai.openai.OpenAIClient;
import com.azure.ai.openai.models.EmbeddingItem;
import com.azure.ai.openai.models.Embeddings;
@@ -7,6 +13,7 @@ import com.azure.ai.openai.models.EmbeddingsOptions;
import com.azure.ai.openai.models.EmbeddingsUsage;
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;
@@ -14,13 +21,6 @@ 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 {
private static final Logger logger = LoggerFactory.getLogger(AzureOpenAiEmbeddingClient.class);

2
spring-ai-core/src/test/java/org/springframework/ai/embedding/EmbeddingUtilTest.java → spring-ai-core/src/test/java/org/springframework/ai/embedding/EmbeddingUtilTests.java
View File

@@ -38,7 +38,7 @@ import static org.mockito.Mockito.when;
* @author Christian Tzolov
*/
@ExtendWith(MockitoExtension.class)
public class EmbeddingUtilTest {
public class EmbeddingUtilTests {
@Mock
private EmbeddingClient embeddingClient;

23
vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java
View File

@@ -67,6 +67,8 @@ public class MilvusVectorStore implements VectorStore, SmartLifecycle {
public static final int OPENAI_EMBEDDING_DIMENSION_SIZE = 1536;
public static final int INVALID_EMBEDDING_DIMENSION = -1;
public static final String DEFAULT_DATABASE_NAME = "default";
public static final String DEFAULT_COLLECTION_NAME = "vector_store";
@@ -140,7 +142,7 @@ public class MilvusVectorStore implements VectorStore, SmartLifecycle {
private String collectionName = DEFAULT_COLLECTION_NAME;
private int embeddingDimension = OPENAI_EMBEDDING_DIMENSION_SIZE;
private int embeddingDimension = INVALID_EMBEDDING_DIMENSION;
private IndexType indexType = IndexType.IVF_FLAT;
@@ -442,7 +444,7 @@ public class MilvusVectorStore implements VectorStore, SmartLifecycle {
FieldType embeddingFieldType = FieldType.newBuilder()
.withName(EMBEDDING_FIELD_NAME)
.withDataType(DataType.FloatVector)
.withDimension(this.config.embeddingDimension)
.withDimension(this.embeddingDimensions())
.build();
CreateCollectionParam createCollectionReq = CreateCollectionParam.newBuilder()
@@ -495,6 +497,23 @@ public class MilvusVectorStore implements VectorStore, SmartLifecycle {
}
}
int embeddingDimensions() {
if (this.config.embeddingDimension != INVALID_EMBEDDING_DIMENSION) {
return this.config.embeddingDimension;
}
try {
int embeddingDimensions = this.embeddingClient.dimensions();
if (embeddingDimensions > 0) {
return embeddingDimensions;
}
}
catch (Exception e) {
logger.warn("Failed to obtain the embedding dimensions from the embedding client and fall backs to default:"
+ this.config.embeddingDimension, e);
}
return OPENAI_EMBEDDING_DIMENSION_SIZE;
}
// used by the test as well
void dropCollection() {

87
vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusEmbeddingDimensionsTests.java Normal file
View File

@@ -0,0 +1,87 @@
/*
* 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.vectorstore;
import io.milvus.client.MilvusServiceClient;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.vectorstore.MilvusVectorStore.MilvusVectorStoreConfig;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.only;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
/**
* @author Christian Tzolov
*/
@ExtendWith(MockitoExtension.class)
public class MilvusEmbeddingDimensionsTests {
@Mock
private EmbeddingClient embeddingClient;
@Mock
private MilvusServiceClient milvusClient;
@Test
public void explicitlySetDimensions() {
final int explicitDimensions = 696;
MilvusVectorStoreConfig config = MilvusVectorStoreConfig.builder()
.withEmbeddingDimension(explicitDimensions)
.build();
var dim = new MilvusVectorStore(milvusClient, embeddingClient, config).embeddingDimensions();
assertThat(dim).isEqualTo(explicitDimensions);
verify(embeddingClient, never()).dimensions();
}
@Test
public void embeddingClientDimensions() {
when(embeddingClient.dimensions()).thenReturn(969);
MilvusVectorStoreConfig config = MilvusVectorStoreConfig.builder().build();
var dim = new MilvusVectorStore(milvusClient, embeddingClient, config).embeddingDimensions();
assertThat(dim).isEqualTo(969);
verify(embeddingClient, only()).dimensions();
}
@Test
public void fallBackToDefaultDimensions() {
when(embeddingClient.dimensions()).thenThrow(new RuntimeException());
var dim = new MilvusVectorStore(milvusClient, embeddingClient,
MilvusVectorStoreConfig.builder().build())
.embeddingDimensions();
assertThat(dim).isEqualTo(MilvusVectorStore.OPENAI_EMBEDDING_DIMENSION_SIZE);
verify(embeddingClient, only()).dimensions();
}
}

9
vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java
View File

@@ -102,7 +102,7 @@ public class MilvusVectorStoreIT {
public void addAndSearchTest(String metricType) {
contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class))
.withPropertyValues("spring.ai.vectorstore.milvus.metricType=" + metricType)
.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType)
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
@@ -135,7 +135,7 @@ public class MilvusVectorStoreIT {
public void documentUpdateTest(String metricType) {
contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class))
.withPropertyValues("spring.ai.vectorstore.milvus.metricType=" + metricType)
.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType)
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
@@ -181,7 +181,7 @@ public class MilvusVectorStoreIT {
public void searchThresholdTest(String metricType) {
contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class))
.withPropertyValues("spring.ai.vectorstore.milvus.metricType=" + metricType)
.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType)
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
@@ -217,7 +217,7 @@ public class MilvusVectorStoreIT {
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestApplication {
@Value("${spring.ai.vectorstore.milvus.metricType}")
@Value("${test.spring.ai.vectorstore.milvus.metricType}")
private MetricType metricType;
@Bean
@@ -227,7 +227,6 @@ public class MilvusVectorStoreIT {
.withDatabaseName("default")
.withIndexType(IndexType.IVF_FLAT)
.withMetricType(metricType)
.withEmbeddingDimension(MilvusVectorStore.OPENAI_EMBEDDING_DIMENSION_SIZE)
.build();
return new MilvusVectorStore(milvusClient, embeddingClient, config);
}

36
vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java
View File

@@ -29,6 +29,8 @@ import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.pgvector.PGvector;
import org.postgresql.util.PGobject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
@@ -45,8 +47,12 @@ import org.springframework.lang.Nullable;
*/
public class PgVectorStore implements VectorStore, SmartLifecycle {
private static final Logger logger = LoggerFactory.getLogger(PgVectorStore.class);
public static final int OPENAI_EMBEDDING_DIMENSION_SIZE = 1536;
public static final int INVALID_EMBEDDING_DIMENSION = -1;
public static final String VECTOR_TABLE_NAME = "vector_store";
private final JdbcTemplate jdbcTemplate;
@@ -176,8 +182,13 @@ public class PgVectorStore implements VectorStore, SmartLifecycle {
}
public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingClient embeddingClient) {
this(jdbcTemplate, embeddingClient, OPENAI_EMBEDDING_DIMENSION_SIZE,
PgVectorStore.PgDistanceType.CosineDistance, false, PgIndexType.NONE);
this(jdbcTemplate, embeddingClient, INVALID_EMBEDDING_DIMENSION, PgVectorStore.PgDistanceType.CosineDistance,
false, PgIndexType.NONE);
}
public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingClient embeddingClient, int dimensions) {
this(jdbcTemplate, embeddingClient, dimensions, PgVectorStore.PgDistanceType.CosineDistance, false,
PgIndexType.NONE);
}
public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingClient embeddingClient, int dimensions,
@@ -299,7 +310,7 @@ public class PgVectorStore implements VectorStore, SmartLifecycle {
this.jdbcTemplate.execute("CREATE TABLE IF NOT EXISTS " + VECTOR_TABLE_NAME + " ( "
+ "id uuid DEFAULT uuid_generate_v4 () PRIMARY KEY, " + "content text, " + "metadata json, "
+ "embedding vector(" + this.dimensions + "))");
+ "embedding vector(" + this.embeddingDimensions() + "))");
if (this.createIndexMethod != PgIndexType.NONE) {
this.jdbcTemplate.execute("CREATE INDEX ON " + VECTOR_TABLE_NAME + " USING " + this.createIndexMethod
@@ -309,6 +320,25 @@ public class PgVectorStore implements VectorStore, SmartLifecycle {
this.isRunning.set(true);
}
int embeddingDimensions() {
// The manually set dimensions have precedence over the computed one.
if (this.dimensions > 0) {
return this.dimensions;
}
try {
int embeddingDimensions = this.embeddingClient.dimensions();
if (embeddingDimensions > 0) {
return embeddingDimensions;
}
}
catch (Exception e) {
logger.warn("Failed to obtain the embedding dimensions from the embedding client and fall backs to default:"
+ OPENAI_EMBEDDING_DIMENSION_SIZE, e);
}
return OPENAI_EMBEDDING_DIMENSION_SIZE;
}
@Override
public void stop() {
// Remove existing VectorStoreTable

78
vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorEmbeddingDimensionsTests.java Normal file
View File

@@ -0,0 +1,78 @@
/*
* 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.vectorstore;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.jdbc.core.JdbcTemplate;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.only;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
/**
* @author Christian Tzolov
*/
@ExtendWith(MockitoExtension.class)
public class PgVectorEmbeddingDimensionsTests {
@Mock
private EmbeddingClient embeddingClient;
@Mock
private JdbcTemplate jdbcTemplate;
@Test
public void explicitlySetDimensions() {
final int explicitDimensions = 696;
var dim = new PgVectorStore(jdbcTemplate, embeddingClient, explicitDimensions).embeddingDimensions();
assertThat(dim).isEqualTo(explicitDimensions);
verify(embeddingClient, never()).dimensions();
}
@Test
public void embeddingClientDimensions() {
when(embeddingClient.dimensions()).thenReturn(969);
var dim = new PgVectorStore(jdbcTemplate, embeddingClient).embeddingDimensions();
assertThat(dim).isEqualTo(969);
verify(embeddingClient, only()).dimensions();
}
@Test
public void fallBackToDefaultDimensions() {
when(embeddingClient.dimensions()).thenThrow(new RuntimeException());
var dim = new PgVectorStore(jdbcTemplate, embeddingClient).embeddingDimensions();
assertThat(dim).isEqualTo(PgVectorStore.OPENAI_EMBEDDING_DIMENSION_SIZE);
verify(embeddingClient, only()).dimensions();
}
}

12
vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java
View File

@@ -72,7 +72,7 @@ public class PgVectorStoreIT {
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class)
.withPropertyValues("spring.ai.openai.apiKey=" + System.getenv("OPENAI_API_KEY"),
"spring.ai.vectorstore.pgvector.distanceType=CosineDistance",
"test.spring.ai.vectorstore.pgvector.distanceType=CosineDistance",
// JdbcTemplate configuration
String.format("app.datasource.url=jdbc:postgresql://localhost:%d/%s",
@@ -84,7 +84,7 @@ public class PgVectorStoreIT {
@ValueSource(strings = { "CosineDistance", "EuclideanDistance", "NegativeInnerProduct" })
public void addAndSearchTest(String distanceType) {
contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class))
.withPropertyValues("spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
@@ -114,7 +114,7 @@ public class PgVectorStoreIT {
public void documentUpdateTest(String distanceType) {
contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class))
.withPropertyValues("spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
@@ -153,7 +153,7 @@ public class PgVectorStoreIT {
public void searchThresholdTest(String distanceType) {
contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class))
.withPropertyValues("spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
@@ -210,12 +210,12 @@ public class PgVectorStoreIT {
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestApplication {
@Value("${spring.ai.vectorstore.pgvector.distanceType}")
@Value("${test.spring.ai.vectorstore.pgvector.distanceType}")
PgVectorStore.PgDistanceType distanceType;
@Bean
public VectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingClient embeddingClient) {
return new PgVectorStore(jdbcTemplate, embeddingClient, PgVectorStore.OPENAI_EMBEDDING_DIMENSION_SIZE,
return new PgVectorStore(jdbcTemplate, embeddingClient, PgVectorStore.INVALID_EMBEDDING_DIMENSION,
distanceType, true, PgIndexType.HNSW);
}