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

Refactor ChromaVectorStore builder API

Browse Source 
The commit restructures the ChromaVectorStore builder pattern to use a no-args constructor
with fluent API for setting the ChromaApi.

This change:

 - Makes builder creation consistent with other vector stores
 - Moves ChromaApi validation to the doValidate method
 - Improves builder API ergonomics

The change requires updating all builder usages to use the new .chromaApi() method instead
of passing it in the constructor.
This commit is contained in:
Soby Chacko
2024-12-10 17:58:19 -05:00
committed by Mark Pollack
parent ad87c122d9
commit e2825c9877
7 changed files with 26 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreAutoConfiguration.java
View File

@@ -86,7 +86,8 @@ public class ChromaVectorStoreAutoConfiguration {
ChromaVectorStoreProperties storeProperties, ObjectProvider<ObservationRegistry> observationRegistry,
ObjectProvider<VectorStoreObservationConvention> customObservationConvention,
BatchingStrategy chromaBatchingStrategy) {
return ChromaVectorStore.builder(chromaApi)
return ChromaVectorStore.builder()
.chromaApi(chromaApi)
.embeddingModel(embeddingModel)
.collectionName(storeProperties.getCollectionName())
.initializeSchema(storeProperties.isInitializeSchema())

17
vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java
View File

@@ -102,7 +102,8 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements
boolean initializeSchema, ObservationRegistry observationRegistry,
VectorStoreObservationConvention customObservationConvention, BatchingStrategy batchingStrategy) {
this(builder(chromaApi).embeddingModel(embeddingModel)
this(builder().chromaApi(chromaApi)
.embeddingModel(embeddingModel)
.collectionName(collectionName)
.initializeSchema(initializeSchema)
.observationRegistry(observationRegistry)
@@ -113,8 +114,11 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements
/**
* @param builder {@link Builder} for chroma vector store
*/
private ChromaVectorStore(ChromaBuilder builder) {
protected ChromaVectorStore(ChromaBuilder builder) {
super(builder);
Assert.notNull(builder.chromaApi, "ChromaApi must not be null");
this.chromaApi = builder.chromaApi;
this.collectionName = builder.collectionName;
this.initializeSchema = builder.initializeSchema;
@@ -151,8 +155,8 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements
}
}
public static ChromaBuilder builder(ChromaApi chromaApi) {
return new ChromaBuilder(chromaApi);
public static ChromaBuilder builder() {
return new ChromaBuilder();
}
@Override
@@ -274,7 +278,7 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements
public static class ChromaBuilder extends AbstractVectorStoreBuilder<ChromaBuilder> {
private final ChromaApi chromaApi;
private ChromaApi chromaApi;
private String collectionName = DEFAULT_COLLECTION_NAME;
@@ -286,9 +290,10 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements
private boolean initializeImmediately = false;
public ChromaBuilder(ChromaApi chromaApi) {
public ChromaBuilder chromaApi(ChromaApi chromaApi) {
Assert.notNull(chromaApi, "ChromaApi must not be null");
this.chromaApi = chromaApi;
return this;
}
/**

3
vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/BasicAuthChromaWhereIT.java
View File

@@ -111,7 +111,8 @@ public class BasicAuthChromaWhereIT {
@Bean
public VectorStore chromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi) {
return ChromaVectorStore.builder(chromaApi)
return ChromaVectorStore.builder()
.chromaApi(chromaApi)
.embeddingModel(embeddingModel)
.collectionName("TestCollection")
.initializeSchema(true)

8
vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaApiIT.java
View File

@@ -207,7 +207,8 @@ public class ChromaApiIT {
assertThat(collection).isNotNull();
assertThat(collection.name()).isEqualTo("test-collection");
ChromaVectorStore store = ChromaVectorStore.builder(this.chromaApi)
ChromaVectorStore store = ChromaVectorStore.builder()
.chromaApi(this.chromaApi)
.embeddingModel(this.embeddingModel)
.collectionName("test-collection")
.initializeImmediately(true)
@@ -219,7 +220,7 @@ public class ChromaApiIT {
@Test
void shouldCreateNewCollectionWhenSchemaInitializationEnabled() {
ChromaVectorStore store = new ChromaVectorStore.ChromaBuilder(this.chromaApi)
ChromaVectorStore store = new ChromaVectorStore.ChromaBuilder().chromaApi(this.chromaApi)
.embeddingModel(this.embeddingModel)
.collectionName("new-collection")
.initializeSchema(true)
@@ -236,7 +237,8 @@ public class ChromaApiIT {
@Test
void shouldFailWhenCollectionDoesNotExist() {
assertThatThrownBy(() -> new ChromaVectorStore.ChromaBuilder(this.chromaApi).embeddingModel(this.embeddingModel)
assertThatThrownBy(() -> new ChromaVectorStore.ChromaBuilder().chromaApi(this.chromaApi)
.embeddingModel(this.embeddingModel)
.collectionName("non-existent")
.initializeSchema(false)
.initializeImmediately(true)

3
vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java
View File

@@ -252,7 +252,8 @@ public class ChromaVectorStoreIT {
@Bean
public VectorStore chromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi) {
return ChromaVectorStore.builder(chromaApi)
return ChromaVectorStore.builder()
.chromaApi(chromaApi)
.embeddingModel(embeddingModel)
.collectionName("TestCollection")
.initializeSchema(true)

3
vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreObservationIT.java
View File

@@ -176,7 +176,8 @@ public class ChromaVectorStoreObservationIT {
@Bean
public VectorStore chromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi,
ObservationRegistry observationRegistry) {
return ChromaVectorStore.builder(chromaApi)
return ChromaVectorStore.builder()
.chromaApi(chromaApi)
.embeddingModel(embeddingModel)
.collectionName("TestCollection")
.initializeSchema(true)

3
vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/TokenSecuredChromaWhereIT.java
View File

@@ -144,7 +144,8 @@ public class TokenSecuredChromaWhereIT {
@Bean
public VectorStore chromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi) {
return ChromaVectorStore.builder(chromaApi)
return ChromaVectorStore.builder()
.chromaApi(chromaApi)
.embeddingModel(embeddingModel)
.collectionName("TestCollection")
.initializeSchema(true)