From 50a9a92c55b31e807e22993ca630d9ce4011b061 Mon Sep 17 00:00:00 2001 From: Soby Chacko Date: Tue, 10 Dec 2024 17:58:19 -0500 Subject: [PATCH 1/3] refactor ChromaVectorStore builder API 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. --- .../AbstractVectorStoreBuilder.java | 3 +++ .../ChromaVectorStoreAutoConfiguration.java | 3 ++- .../chroma/vectorstore/ChromaVectorStore.java | 25 ++++++++++++++----- .../vectorstore/BasicAuthChromaWhereIT.java | 3 ++- .../ai/chroma/vectorstore/ChromaApiIT.java | 8 +++--- .../vectorstore/ChromaVectorStoreIT.java | 3 ++- .../ChromaVectorStoreObservationIT.java | 3 ++- .../TokenSecuredChromaWhereIT.java | 3 ++- 8 files changed, 37 insertions(+), 14 deletions(-) diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/AbstractVectorStoreBuilder.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/AbstractVectorStoreBuilder.java index 0affbdcc8b0..b241b1e9102 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/AbstractVectorStoreBuilder.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/AbstractVectorStoreBuilder.java @@ -85,6 +85,9 @@ public T embeddingModel(EmbeddingModel embeddingModel) { protected void validate() { Assert.notNull(this.embeddingModel, "EmbeddingModel must be configured"); + doValidate(); } + protected abstract void doValidate(); + } diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreAutoConfiguration.java index 9f0b8493c64..50f75891882 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreAutoConfiguration.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreAutoConfiguration.java @@ -86,7 +86,8 @@ public ChromaVectorStore vectorStore(EmbeddingModel embeddingModel, ChromaApi ch ChromaVectorStoreProperties storeProperties, ObjectProvider observationRegistry, ObjectProvider customObservationConvention, BatchingStrategy chromaBatchingStrategy) { - return ChromaVectorStore.builder(chromaApi) + return ChromaVectorStore.builder() + .chromaApi(chromaApi) .embeddingModel(embeddingModel) .collectionName(storeProperties.getCollectionName()) .initializeSchema(storeProperties.isInitializeSchema()) diff --git a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java index f1827162fe1..4a46a0cc975 100644 --- a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java +++ b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java @@ -102,7 +102,8 @@ public ChromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi, Str 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,14 @@ public ChromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi, Str /** * @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"); + Assert.notNull(builder.batchingStrategy, "BatchingStrategy must not be null"); + Assert.notNull(builder.filterExpressionConverter, "FilterExpressionConverter must not be null"); + Assert.hasText(builder.collectionName, "Collection name must not be empty"); + this.chromaApi = builder.chromaApi; this.collectionName = builder.collectionName; this.initializeSchema = builder.initializeSchema; @@ -151,8 +158,8 @@ public void afterPropertiesSet() throws Exception { } } - public static ChromaBuilder builder(ChromaApi chromaApi) { - return new ChromaBuilder(chromaApi); + public static ChromaBuilder builder() { + return new ChromaBuilder(); } @Override @@ -274,7 +281,7 @@ public VectorStoreObservationContext.Builder createObservationContextBuilder(Str public static class ChromaBuilder extends AbstractVectorStoreBuilder { - private final ChromaApi chromaApi; + private ChromaApi chromaApi; private String collectionName = DEFAULT_COLLECTION_NAME; @@ -286,9 +293,10 @@ public static class ChromaBuilder extends AbstractVectorStoreBuilder 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) diff --git a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java index 74ab85f5edb..c48c57a779b 100644 --- a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java +++ b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java @@ -252,7 +252,8 @@ public ChromaApi chromaApi(RestClient.Builder builder) { @Bean public VectorStore chromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi) { - return ChromaVectorStore.builder(chromaApi) + return ChromaVectorStore.builder() + .chromaApi(chromaApi) .embeddingModel(embeddingModel) .collectionName("TestCollection") .initializeSchema(true) diff --git a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreObservationIT.java b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreObservationIT.java index 34846632efe..22157fee4fe 100644 --- a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreObservationIT.java +++ b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreObservationIT.java @@ -176,7 +176,8 @@ public ChromaApi chromaApi(RestClient.Builder builder) { @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) diff --git a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/TokenSecuredChromaWhereIT.java b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/TokenSecuredChromaWhereIT.java index bd8220d3ab7..88b5f56b8d0 100644 --- a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/TokenSecuredChromaWhereIT.java +++ b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/TokenSecuredChromaWhereIT.java @@ -144,7 +144,8 @@ public ChromaApi chromaApi(RestClient.Builder builder) { @Bean public VectorStore chromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi) { - return ChromaVectorStore.builder(chromaApi) + return ChromaVectorStore.builder() + .chromaApi(chromaApi) .embeddingModel(embeddingModel) .collectionName("TestCollection") .initializeSchema(true) From ad4fcd963532e3fd061e066cbee04f5f0a30c087 Mon Sep 17 00:00:00 2001 From: Soby Chacko Date: Tue, 10 Dec 2024 18:08:05 -0500 Subject: [PATCH 2/3] Addressing PR review --- .../ai/vectorstore/AbstractVectorStoreBuilder.java | 3 --- .../ai/chroma/vectorstore/ChromaVectorStore.java | 5 ----- 2 files changed, 8 deletions(-) diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/AbstractVectorStoreBuilder.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/AbstractVectorStoreBuilder.java index b241b1e9102..0affbdcc8b0 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/AbstractVectorStoreBuilder.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/AbstractVectorStoreBuilder.java @@ -85,9 +85,6 @@ public T embeddingModel(EmbeddingModel embeddingModel) { protected void validate() { Assert.notNull(this.embeddingModel, "EmbeddingModel must be configured"); - doValidate(); } - protected abstract void doValidate(); - } diff --git a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java index 4a46a0cc975..8524e434062 100644 --- a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java +++ b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java @@ -365,11 +365,6 @@ public ChromaVectorStore build() { return new ChromaVectorStore(this); } - @Override - protected void doValidate() { - Assert.notNull(this.chromaApi, "ChromaApi must not be null"); - } - } } From 1e816a5b5fb37de6bf0d2f57363015206aa65067 Mon Sep 17 00:00:00 2001 From: Soby Chacko Date: Tue, 10 Dec 2024 18:10:19 -0500 Subject: [PATCH 3/3] cleanup --- .../ai/chroma/vectorstore/ChromaVectorStore.java | 3 --- 1 file changed, 3 deletions(-) diff --git a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java index 8524e434062..80898ac17ce 100644 --- a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java +++ b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java @@ -118,9 +118,6 @@ protected ChromaVectorStore(ChromaBuilder builder) { super(builder); Assert.notNull(builder.chromaApi, "ChromaApi must not be null"); - Assert.notNull(builder.batchingStrategy, "BatchingStrategy must not be null"); - Assert.notNull(builder.filterExpressionConverter, "FilterExpressionConverter must not be null"); - Assert.hasText(builder.collectionName, "Collection name must not be empty"); this.chromaApi = builder.chromaApi; this.collectionName = builder.collectionName;