Skip to content

Commit 8a1dcfc

Browse files
sobychackomarkpollack
sobychacko
authored and
markpollack
committed
Conditionally enable invidividual vector store implementations
- Introduce top level properties to support conditional logic - The properties will have the format like "spring.ai.vectorstore.type=<vectorstore provider>" to enable specific vectore store implementation - By default, these will be enabled when no specific properties are set. To disable, set any value other than the provider name for example, none" Signed-off-by: Soby Chacko <soby.chacko@broadcom.com>
1 parent 6b12ff5 commit 8a1dcfc

File tree

38 files changed

+661
-3
lines changed

38 files changed

+661
-3
lines changed

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-azure-cosmos-db/src/main/java/org/springframework/ai/vectorstore/cosmosdb/autoconfigure/CosmosDBVectorStoreAutoConfiguration.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,12 +23,14 @@
2323
import org.springframework.ai.embedding.BatchingStrategy;
2424
import org.springframework.ai.embedding.EmbeddingModel;
2525
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
26+
import org.springframework.ai.vectorstore.SpringAIVectorStoreTypes;
2627
import org.springframework.ai.vectorstore.cosmosdb.CosmosDBVectorStore;
2728
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
2829
import org.springframework.beans.factory.ObjectProvider;
2930
import org.springframework.boot.autoconfigure.AutoConfiguration;
3031
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
3132
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
33+
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
3234
import org.springframework.boot.context.properties.EnableConfigurationProperties;
3335
import org.springframework.context.annotation.Bean;
3436
import java.util.List;
@@ -43,6 +45,8 @@
4345
@AutoConfiguration
4446
@ConditionalOnClass({ CosmosDBVectorStore.class, EmbeddingModel.class, CosmosAsyncClient.class })
4547
@EnableConfigurationProperties(CosmosDBVectorStoreProperties.class)
48+
@ConditionalOnProperty(name = SpringAIVectorStoreTypes.TYPE, havingValue = SpringAIVectorStoreTypes.AZURE_COSMOS_DB,
49+
matchIfMissing = true)
4650
public class CosmosDBVectorStoreAutoConfiguration {
4751

4852
String endpoint;

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-azure-cosmos-db/src/test/java/org/springframework/ai/vectorstore/cosmosdb/autoconfigure/CosmosDBVectorStoreAutoConfigurationIT.java

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import org.springframework.ai.transformers.TransformersEmbeddingModel;
3232
import org.springframework.ai.vectorstore.SearchRequest;
3333
import org.springframework.ai.vectorstore.VectorStore;
34+
import org.springframework.ai.vectorstore.cosmosdb.CosmosDBVectorStore;
3435
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
3536
import org.springframework.boot.autoconfigure.AutoConfigurations;
3637
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
@@ -169,6 +170,33 @@ void testSimilaritySearchWithFilter() {
169170
assertThat(results4).isEmpty();
170171
}
171172

173+
@Test
174+
public void autoConfigurationDisabledWhenTypeIsNone() {
175+
this.contextRunner.withPropertyValues("spring.ai.vectorstore.type=none").run(context -> {
176+
assertThat(context.getBeansOfType(CosmosDBVectorStoreProperties.class)).isEmpty();
177+
assertThat(context.getBeansOfType(CosmosDBVectorStore.class)).isEmpty();
178+
assertThat(context.getBeansOfType(VectorStore.class)).isEmpty();
179+
});
180+
}
181+
182+
@Test
183+
public void autoConfigurationEnabledByDefault() {
184+
this.contextRunner.run(context -> {
185+
assertThat(context.getBeansOfType(CosmosDBVectorStoreProperties.class)).isNotEmpty();
186+
assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty();
187+
assertThat(context.getBean(VectorStore.class)).isInstanceOf(CosmosDBVectorStore.class);
188+
});
189+
}
190+
191+
@Test
192+
public void autoConfigurationEnabledWhenTypeIsAzureCosmosDB() {
193+
this.contextRunner.withPropertyValues("spring.ai.vectorstore.type=azure-cosmmos-db").run(context -> {
194+
assertThat(context.getBeansOfType(CosmosDBVectorStoreProperties.class)).isNotEmpty();
195+
assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty();
196+
assertThat(context.getBean(VectorStore.class)).isInstanceOf(CosmosDBVectorStore.class);
197+
});
198+
}
199+
172200
@Configuration(proxyBeanMethods = false)
173201
static class Config {
174202

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-azure/src/main/java/org/springframework/ai/vectorstore/azure/autoconfigure/AzureVectorStoreAutoConfiguration.java

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import org.springframework.ai.embedding.BatchingStrategy;
2929
import org.springframework.ai.embedding.EmbeddingModel;
3030
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
31+
import org.springframework.ai.vectorstore.SpringAIVectorStoreTypes;
3132
import org.springframework.ai.vectorstore.azure.AzureVectorStore;
3233
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
3334
import org.springframework.beans.factory.ObjectProvider;
@@ -47,7 +48,8 @@
4748
@AutoConfiguration
4849
@ConditionalOnClass({ EmbeddingModel.class, SearchIndexClient.class, AzureVectorStore.class })
4950
@EnableConfigurationProperties({ AzureVectorStoreProperties.class })
50-
@ConditionalOnProperty(prefix = "spring.ai.vectorstore.azure", value = { "url", "api-key", "index-name" })
51+
@ConditionalOnProperty(name = SpringAIVectorStoreTypes.TYPE, havingValue = SpringAIVectorStoreTypes.AZURE,
52+
matchIfMissing = true)
5153
public class AzureVectorStoreAutoConfiguration {
5254

5355
private static final String APPLICATION_ID = "spring-ai";

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-azure/src/test/java/org/springframework/ai/vectorstore/azure/autoconfigure/AzureVectorStoreAutoConfigurationIT.java

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,33 @@ public void addAndSearchTest() {
146146
});
147147
}
148148

149+
@Test
150+
public void autoConfigurationDisabledWhenTypeIsNone() {
151+
this.contextRunner.withPropertyValues("spring.ai.vectorstore.type=none").run(context -> {
152+
assertThat(context.getBeansOfType(AzureVectorStoreProperties.class)).isEmpty();
153+
assertThat(context.getBeansOfType(AzureVectorStore.class)).isEmpty();
154+
assertThat(context.getBeansOfType(VectorStore.class)).isEmpty();
155+
});
156+
}
157+
158+
@Test
159+
public void autoConfigurationEnabledByDefault() {
160+
this.contextRunner.run(context -> {
161+
assertThat(context.getBeansOfType(AzureVectorStoreProperties.class)).isNotEmpty();
162+
assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty();
163+
assertThat(context.getBean(VectorStore.class)).isInstanceOf(AzureVectorStore.class);
164+
});
165+
}
166+
167+
@Test
168+
public void autoConfigurationEnabledWhenTypeIsAzure() {
169+
this.contextRunner.withPropertyValues("spring.ai.vectorstore.type=azure").run(context -> {
170+
assertThat(context.getBeansOfType(AzureVectorStoreProperties.class)).isNotEmpty();
171+
assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty();
172+
assertThat(context.getBean(VectorStore.class)).isInstanceOf(AzureVectorStore.class);
173+
});
174+
}
175+
149176
@Configuration(proxyBeanMethods = false)
150177
static class Config {
151178

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-cassandra/src/main/java/org/springframework/ai/vectorstore/cassandra/autoconfigure/CassandraVectorStoreAutoConfiguration.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import org.springframework.ai.embedding.BatchingStrategy;
2626
import org.springframework.ai.embedding.EmbeddingModel;
2727
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
28+
import org.springframework.ai.vectorstore.SpringAIVectorStoreTypes;
2829
import org.springframework.ai.vectorstore.cassandra.CassandraVectorStore;
2930
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
3031
import org.springframework.beans.factory.ObjectProvider;
@@ -33,6 +34,7 @@
3334
import org.springframework.boot.autoconfigure.cassandra.DriverConfigLoaderBuilderCustomizer;
3435
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
3536
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
37+
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
3638
import org.springframework.boot.context.properties.EnableConfigurationProperties;
3739
import org.springframework.context.annotation.Bean;
3840

@@ -47,6 +49,8 @@
4749
@AutoConfiguration(after = CassandraAutoConfiguration.class)
4850
@ConditionalOnClass({ CassandraVectorStore.class, CqlSession.class })
4951
@EnableConfigurationProperties(CassandraVectorStoreProperties.class)
52+
@ConditionalOnProperty(name = SpringAIVectorStoreTypes.TYPE, havingValue = SpringAIVectorStoreTypes.CASSANDRA,
53+
matchIfMissing = true)
5054
public class CassandraVectorStoreAutoConfiguration {
5155

5256
@Bean

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-cassandra/src/test/java/org/springframework/ai/vectorstore/cassandra/autoconfigure/CassandraVectorStoreAutoConfigurationIT.java

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import java.util.Map;
2121

2222
import io.micrometer.observation.tck.TestObservationRegistry;
23+
import org.junit.jupiter.api.Disabled;
2324
import org.junit.jupiter.api.Test;
2425
import org.testcontainers.containers.CassandraContainer;
2526
import org.testcontainers.junit.jupiter.Container;
@@ -34,6 +35,7 @@
3435
import org.springframework.ai.transformers.TransformersEmbeddingModel;
3536
import org.springframework.ai.vectorstore.SearchRequest;
3637
import org.springframework.ai.vectorstore.VectorStore;
38+
import org.springframework.ai.vectorstore.cassandra.CassandraVectorStore;
3739
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
3840
import org.springframework.boot.autoconfigure.AutoConfigurations;
3941
import org.springframework.boot.autoconfigure.cassandra.CassandraAutoConfiguration;
@@ -111,6 +113,42 @@ void addAndSearch() {
111113
});
112114
}
113115

116+
@Test
117+
public void autoConfigurationDisabledWhenTypeIsNone() {
118+
this.contextRunner.withPropertyValues("spring.ai.vectorstore.type=none").run(context -> {
119+
assertThat(context.getBeansOfType(CassandraVectorStoreProperties.class)).isEmpty();
120+
assertThat(context.getBeansOfType(CassandraVectorStore.class)).isEmpty();
121+
assertThat(context.getBeansOfType(VectorStore.class)).isEmpty();
122+
});
123+
}
124+
125+
@Test
126+
public void autoConfigurationEnabledByDefault() {
127+
this.contextRunner.withPropertyValues("spring.cassandra.contactPoints=" + getContactPointHost())
128+
.withPropertyValues("spring.cassandra.port=" + getContactPointPort())
129+
.withPropertyValues("spring.cassandra.localDatacenter=" + cassandraContainer.getLocalDatacenter())
130+
.withPropertyValues("spring.ai.vectorstore.cassandra.fixedThreadPoolExecutorSize=8")
131+
.run(context -> {
132+
assertThat(context.getBeansOfType(CassandraVectorStoreProperties.class)).isNotEmpty();
133+
assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty();
134+
assertThat(context.getBean(VectorStore.class)).isInstanceOf(CassandraVectorStore.class);
135+
});
136+
}
137+
138+
@Test
139+
public void autoConfigurationEnabledWhenTypeIsCassandra() {
140+
this.contextRunner.withPropertyValues("spring.ai.vectorstore.type=cassandra")
141+
.withPropertyValues("spring.cassandra.contactPoints=" + getContactPointHost())
142+
.withPropertyValues("spring.cassandra.port=" + getContactPointPort())
143+
.withPropertyValues("spring.cassandra.localDatacenter=" + cassandraContainer.getLocalDatacenter())
144+
.withPropertyValues("spring.ai.vectorstore.cassandra.fixedThreadPoolExecutorSize=8")
145+
.run(context -> {
146+
assertThat(context.getBeansOfType(CassandraVectorStoreProperties.class)).isNotEmpty();
147+
assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty();
148+
assertThat(context.getBean(VectorStore.class)).isInstanceOf(CassandraVectorStore.class);
149+
});
150+
}
151+
114152
private String getContactPointHost() {
115153
return cassandraContainer.getContactPoint().getHostString();
116154
}

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-chroma/src/main/java/org/springframework/ai/vectorstore/chroma/autoconfigure/ChromaVectorStoreAutoConfiguration.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,11 +24,13 @@
2424
import org.springframework.ai.embedding.EmbeddingModel;
2525
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
2626
import org.springframework.ai.chroma.vectorstore.ChromaVectorStore;
27+
import org.springframework.ai.vectorstore.SpringAIVectorStoreTypes;
2728
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
2829
import org.springframework.beans.factory.ObjectProvider;
2930
import org.springframework.boot.autoconfigure.AutoConfiguration;
3031
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
3132
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
33+
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
3234
import org.springframework.boot.context.properties.EnableConfigurationProperties;
3335
import org.springframework.context.annotation.Bean;
3436
import org.springframework.util.StringUtils;
@@ -45,6 +47,8 @@
4547
@AutoConfiguration
4648
@ConditionalOnClass({ EmbeddingModel.class, RestClient.class, ChromaVectorStore.class, ObjectMapper.class })
4749
@EnableConfigurationProperties({ ChromaApiProperties.class, ChromaVectorStoreProperties.class })
50+
@ConditionalOnProperty(name = SpringAIVectorStoreTypes.TYPE, havingValue = SpringAIVectorStoreTypes.CHROMA,
51+
matchIfMissing = true)
4852
public class ChromaVectorStoreAutoConfiguration {
4953

5054
@Bean

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-chroma/src/test/java/org/springframework/ai/vectorstore/chroma/autoconfigure/ChromaVectorStoreAutoConfigurationIT.java

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import org.testcontainers.junit.jupiter.Container;
2929
import org.testcontainers.junit.jupiter.Testcontainers;
3030

31+
import org.springframework.ai.chroma.vectorstore.ChromaVectorStore;
3132
import org.springframework.ai.document.Document;
3233
import org.springframework.ai.embedding.EmbeddingModel;
3334
import org.springframework.ai.observation.conventions.VectorStoreProvider;
@@ -145,6 +146,38 @@ public void throwExceptionOnMissingCollectionAndDisabledInitializedSchema() {
145146
"Collection TestCollection doesn't exist and won't be created as the initializeSchema is set to false."));
146147
}
147148

149+
@Test
150+
public void autoConfigurationDisabledWhenTypeIsNone() {
151+
this.contextRunner.withPropertyValues("spring.ai.vectorstore.type=none").run(context -> {
152+
assertThat(context.getBeansOfType(ChromaVectorStoreProperties.class)).isEmpty();
153+
assertThat(context.getBeansOfType(ChromaVectorStore.class)).isEmpty();
154+
assertThat(context.getBeansOfType(VectorStore.class)).isEmpty();
155+
});
156+
}
157+
158+
@Test
159+
@Disabled
160+
public void autoConfigurationEnabledByDefault() {
161+
this.contextRunner.withPropertyValues("spring.ai.vectorstore.chroma.initializeSchema=true").run(context -> {
162+
assertThat(context.getBeansOfType(ChromaVectorStoreProperties.class)).isNotEmpty();
163+
assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty();
164+
assertThat(context.getBean(VectorStore.class)).isInstanceOf(ChromaVectorStore.class);
165+
});
166+
}
167+
168+
@Test
169+
@Disabled
170+
public void autoConfigurationEnabledWhenTypeIsChroma() {
171+
this.contextRunner
172+
.withPropertyValues("spring.ai.vectorstore.type=chroma",
173+
"spring.ai.vectorstore.chroma.initializeSchema=true")
174+
.run(context -> {
175+
assertThat(context.getBeansOfType(ChromaVectorStoreProperties.class)).isNotEmpty();
176+
assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty();
177+
assertThat(context.getBean(VectorStore.class)).isInstanceOf(ChromaVectorStore.class);
178+
});
179+
}
180+
148181
@Configuration(proxyBeanMethods = false)
149182
static class Config {
150183

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-elasticsearch/src/main/java/org/springframework/ai/vectorstore/elasticsearch/autoconfigure/ElasticsearchVectorStoreAutoConfiguration.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,13 +22,15 @@
2222
import org.springframework.ai.embedding.BatchingStrategy;
2323
import org.springframework.ai.embedding.EmbeddingModel;
2424
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
25+
import org.springframework.ai.vectorstore.SpringAIVectorStoreTypes;
2526
import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStore;
2627
import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStoreOptions;
2728
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
2829
import org.springframework.beans.factory.ObjectProvider;
2930
import org.springframework.boot.autoconfigure.AutoConfiguration;
3031
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
3132
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
33+
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
3234
import org.springframework.boot.autoconfigure.elasticsearch.ElasticsearchRestClientAutoConfiguration;
3335
import org.springframework.boot.context.properties.EnableConfigurationProperties;
3436
import org.springframework.context.annotation.Bean;
@@ -47,6 +49,8 @@
4749
@AutoConfiguration(after = ElasticsearchRestClientAutoConfiguration.class)
4850
@ConditionalOnClass({ ElasticsearchVectorStore.class, EmbeddingModel.class, RestClient.class })
4951
@EnableConfigurationProperties(ElasticsearchVectorStoreProperties.class)
52+
@ConditionalOnProperty(name = SpringAIVectorStoreTypes.TYPE, havingValue = SpringAIVectorStoreTypes.ELASTICSEARCH,
53+
matchIfMissing = true)
5054
public class ElasticsearchVectorStoreAutoConfiguration {
5155

5256
@Bean

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-elasticsearch/src/test/java/org/springframework/ai/vectorstore/elasticsearch/autoconfigure/ElasticsearchVectorStoreAutoConfigurationIT.java

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
import org.springframework.ai.document.Document;
3535
import org.springframework.ai.observation.conventions.VectorStoreProvider;
3636
import org.springframework.ai.test.vectorstore.ObservationTestUtil;
37+
import org.springframework.ai.vectorstore.VectorStore;
3738
import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStore;
3839
import org.springframework.ai.vectorstore.SearchRequest;
3940
import org.springframework.ai.vectorstore.elasticsearch.SimilarityFunction;
@@ -149,6 +150,33 @@ public void propertiesTest() {
149150
});
150151
}
151152

153+
@Test
154+
public void autoConfigurationDisabledWhenTypeIsNone() {
155+
this.contextRunner.withPropertyValues("spring.ai.vectorstore.type=none").run(context -> {
156+
assertThat(context.getBeansOfType(ElasticsearchVectorStoreProperties.class)).isEmpty();
157+
assertThat(context.getBeansOfType(ElasticsearchVectorStore.class)).isEmpty();
158+
assertThat(context.getBeansOfType(VectorStore.class)).isEmpty();
159+
});
160+
}
161+
162+
@Test
163+
public void autoConfigurationEnabledByDefault() {
164+
this.contextRunner.run(context -> {
165+
assertThat(context.getBeansOfType(ElasticsearchVectorStoreProperties.class)).isNotEmpty();
166+
assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty();
167+
assertThat(context.getBean(VectorStore.class)).isInstanceOf(ElasticsearchVectorStore.class);
168+
});
169+
}
170+
171+
@Test
172+
public void autoConfigurationEnabledWhenTypeIsElasticsearch() {
173+
this.contextRunner.withPropertyValues("spring.ai.vectorstore.type=elasticsearch").run(context -> {
174+
assertThat(context.getBeansOfType(ElasticsearchVectorStoreProperties.class)).isNotEmpty();
175+
assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty();
176+
assertThat(context.getBean(VectorStore.class)).isInstanceOf(ElasticsearchVectorStore.class);
177+
});
178+
}
179+
152180
private String getText(String uri) {
153181
var resource = new DefaultResourceLoader().getResource(uri);
154182
try {

0 commit comments

Comments
 (0)