From d16665dad9c0a0316508cebaf35d7e4d50de5d19 Mon Sep 17 00:00:00 2001 From: Soby Chacko Date: Mon, 25 Nov 2024 19:05:02 -0500 Subject: [PATCH] Add VectorStore builder hierarchy This refactoring introduces a consistent builder pattern across vector store implementations to standardize configuration and initialization, while also moving ChromaVectorStore to a dedicated chroma package. Key changes: - Add VectorStore.Builder interface and AbstractVectorStoreBuilder to establish a common builder hierarchy - Move ChromaVectorStore and related classes from vectorstore to org.springframework.ai.chroma.vectorstore package - Migrate ChromaVectorStore to builder pattern as the first implementation - Add null-safety annotations and parameter validation - Deprecate direct constructors in favor of builder API - Update all tests and documentation to reflect new structure The builder pattern provides several benefits: - Consistent configuration across all vector store implementations - Better validation of required parameters - More flexible initialization order - Clearer separation of concerns between configuration and usage - Improved discoverability of options through method chaining --- .../AbstractVectorStoreBuilder.java | 91 ++++++++++ .../ai/vectorstore/VectorStore.java | 41 +++++ .../AbstractObservationVectorStore.java | 25 ++- .../VectorStoreObservationContext.java | 2 +- .../ChromaVectorStoreAutoConfiguration.java | 15 +- .../chroma/ChromaVectorStoreProperties.java | 2 +- .../vectorstore/azure/AzureVectorStore.java | 2 +- .../ai/vectorstore/CassandraVectorStore.java | 2 +- .../chroma/{ => vectorstore}/ChromaApi.java | 8 +- .../ChromaFilterExpressionConverter.java | 2 +- .../vectorstore/ChromaVectorStore.java | 171 +++++++++++------- .../ai/chroma/vectorstore/package-info.java | 25 +++ .../ai/{ => chroma}/ChromaImage.java | 2 +- .../vectorstore/BasicAuthChromaWhereIT.java | 13 +- .../chroma/{ => vectorstore}/ChromaApiIT.java | 30 +-- .../vectorstore/ChromaVectorStoreIT.java | 13 +- .../ChromaVectorStoreObservationIT.java | 21 ++- .../TokenSecuredChromaWhereIT.java | 13 +- .../vectorstore/ElasticsearchVectorStore.java | 2 +- .../ai/vectorstore/HanaCloudVectorStore.java | 2 +- .../ai/vectorstore/OpenSearchVectorStore.java | 2 +- .../ai/vectorstore/OracleVectorStore.java | 2 +- 22 files changed, 367 insertions(+), 119 deletions(-) create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/vectorstore/AbstractVectorStoreBuilder.java rename vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/{ => vectorstore}/ChromaApi.java (97%) rename vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/{ => chroma}/vectorstore/ChromaFilterExpressionConverter.java (95%) rename vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/{ => chroma}/vectorstore/ChromaVectorStore.java (72%) create mode 100644 vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/package-info.java rename vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/{ => chroma}/ChromaImage.java (95%) rename vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/{ => chroma}/vectorstore/BasicAuthChromaWhereIT.java (91%) rename vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/{ => vectorstore}/ChromaApiIT.java (90%) rename vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/{ => chroma}/vectorstore/ChromaVectorStoreIT.java (96%) rename vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/{ => chroma}/vectorstore/ChromaVectorStoreObservationIT.java (91%) rename vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/{ => chroma}/vectorstore/TokenSecuredChromaWhereIT.java (93%) 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 new file mode 100644 index 000000000..65912b3a4 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/AbstractVectorStoreBuilder.java @@ -0,0 +1,91 @@ +/* + * Copyright 2023-2024 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.micrometer.observation.ObservationRegistry; + +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention; +import org.springframework.lang.Nullable; +import org.springframework.util.Assert; + +/** + * Abstract base builder implementing common builder functionality for + * {@link VectorStore}. Provides default implementations for observation-related settings. + * + * @param the concrete builder type, enabling method chaining with the correct return + * type + */ +public abstract class AbstractVectorStoreBuilder> + implements VectorStore.Builder { + + protected EmbeddingModel embeddingModel; + + protected ObservationRegistry observationRegistry = ObservationRegistry.NOOP; + + @Nullable + protected VectorStoreObservationConvention customObservationConvention; + + public EmbeddingModel getEmbeddingModel() { + return this.embeddingModel; + } + + public ObservationRegistry getObservationRegistry() { + return this.observationRegistry; + } + + @Nullable + public VectorStoreObservationConvention getCustomObservationConvention() { + return this.customObservationConvention; + } + + /** + * Returns this builder cast to the concrete builder type. Used internally to enable + * proper method chaining in subclasses. + * @return this builder cast to the concrete type + */ + @SuppressWarnings("unchecked") + protected T self() { + return (T) this; + } + + @Override + public T observationRegistry(ObservationRegistry observationRegistry) { + Assert.notNull(observationRegistry, "ObservationRegistry must not be null"); + this.observationRegistry = observationRegistry; + return self(); + } + + @Override + public T customObservationConvention(VectorStoreObservationConvention convention) { + Assert.notNull(convention, "VectorStoreObservationConvention must not be null"); + this.customObservationConvention = convention; + return self(); + } + + @Override + public T embeddingModel(EmbeddingModel embeddingModel) { + Assert.notNull(embeddingModel, "EmbeddingModel must not be null"); + this.embeddingModel = embeddingModel; + return self(); + } + + protected void validate() { + Assert.notNull(this.embeddingModel, "EmbeddingModel must be configured"); + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java index 0752edd08..faea243a7 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java @@ -19,8 +19,14 @@ package org.springframework.ai.vectorstore; import java.util.List; import java.util.Optional; +import io.micrometer.observation.ObservationRegistry; + import org.springframework.ai.document.Document; import org.springframework.ai.document.DocumentWriter; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention; +import org.springframework.lang.Nullable; /** * The {@code VectorStore} interface defines the operations for managing and querying @@ -75,4 +81,39 @@ public interface VectorStore extends DocumentWriter { return this.similaritySearch(SearchRequest.query(query)); } + /** + * Builder interface for creating VectorStore instances. Implements a fluent builder + * pattern for configuring observation-related settings. + * + * @param the concrete builder type, enabling method chaining with the correct + * return type + */ + interface Builder> { + + T embeddingModel(EmbeddingModel embeddingModel); + + /** + * Sets the registry for collecting observations and metrics. Defaults to + * {@link ObservationRegistry#NOOP} if not specified. + * @param observationRegistry the registry to use for observations + * @return the builder instance for method chaining + */ + T observationRegistry(ObservationRegistry observationRegistry); + + /** + * Sets a custom convention for creating observations. If not specified, + * {@link DefaultVectorStoreObservationConvention} will be used. + * @param convention the custom observation convention to use + * @return the builder instance for method chaining + */ + T customObservationConvention(VectorStoreObservationConvention convention); + + /** + * Builds and returns a new VectorStore instance with the configured settings. + * @return a new VectorStore instance + */ + VectorStore build(); + + } + } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/AbstractObservationVectorStore.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/AbstractObservationVectorStore.java index 20bc6de87..17a15a5e5 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/AbstractObservationVectorStore.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/AbstractObservationVectorStore.java @@ -22,6 +22,8 @@ import java.util.Optional; import io.micrometer.observation.ObservationRegistry; import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; import org.springframework.lang.Nullable; @@ -31,6 +33,7 @@ import org.springframework.lang.Nullable; * capabilities. * * @author Christian Tzolov + * @author Soby Chacko * @since 1.0.0 */ public abstract class AbstractObservationVectorStore implements VectorStore { @@ -42,17 +45,37 @@ public abstract class AbstractObservationVectorStore implements VectorStore { @Nullable private final VectorStoreObservationConvention customObservationConvention; + @Nullable + protected final EmbeddingModel embeddingModel; + /** * Create a new {@link AbstractObservationVectorStore} instance. * @param observationRegistry the observation registry to use * @param customObservationConvention the custom observation convention to use */ + @Deprecated(since = "1.0.0-M5", forRemoval = true) public AbstractObservationVectorStore(ObservationRegistry observationRegistry, - VectorStoreObservationConvention customObservationConvention) { + @Nullable VectorStoreObservationConvention customObservationConvention) { + this(null, observationRegistry, customObservationConvention); + } + + private AbstractObservationVectorStore(@Nullable EmbeddingModel embeddingModel, + ObservationRegistry observationRegistry, + @Nullable VectorStoreObservationConvention customObservationConvention) { + this.embeddingModel = embeddingModel; this.observationRegistry = observationRegistry; this.customObservationConvention = customObservationConvention; } + /** + * Creates a new AbstractObservationVectorStore instance with the specified builder + * settings. Initializes observation-related components and the embedding model. + * @param builder the builder containing configuration settings + */ + public AbstractObservationVectorStore(AbstractVectorStoreBuilder builder) { + this(builder.getEmbeddingModel(), builder.getObservationRegistry(), builder.getCustomObservationConvention()); + } + /** * Create a new {@link AbstractObservationVectorStore} instance. * @param documents the documents to add diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationContext.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationContext.java index cf5454629..832d3355f 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationContext.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationContext.java @@ -194,7 +194,7 @@ public class VectorStoreObservationContext extends Observation.Context { return this; } - public Builder withFieldName(String fieldName) { + public Builder withFieldName(@Nullable String fieldName) { this.context.setFieldName(fieldName); return this; } 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 de3ba67f8..9f0b8493c 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 @@ -19,11 +19,11 @@ package org.springframework.ai.autoconfigure.vectorstore.chroma; import com.fasterxml.jackson.databind.ObjectMapper; import io.micrometer.observation.ObservationRegistry; -import org.springframework.ai.chroma.ChromaApi; +import org.springframework.ai.chroma.vectorstore.ChromaApi; import org.springframework.ai.embedding.BatchingStrategy; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.embedding.TokenCountBatchingStrategy; -import org.springframework.ai.vectorstore.ChromaVectorStore; +import org.springframework.ai.chroma.vectorstore.ChromaVectorStore; import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention; import org.springframework.beans.factory.ObjectProvider; import org.springframework.boot.autoconfigure.AutoConfiguration; @@ -86,9 +86,14 @@ public class ChromaVectorStoreAutoConfiguration { ChromaVectorStoreProperties storeProperties, ObjectProvider observationRegistry, ObjectProvider customObservationConvention, BatchingStrategy chromaBatchingStrategy) { - return new ChromaVectorStore(embeddingModel, chromaApi, storeProperties.getCollectionName(), - storeProperties.isInitializeSchema(), observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP), - customObservationConvention.getIfAvailable(() -> null), chromaBatchingStrategy); + return ChromaVectorStore.builder(chromaApi) + .embeddingModel(embeddingModel) + .collectionName(storeProperties.getCollectionName()) + .initializeSchema(storeProperties.isInitializeSchema()) + .observationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP)) + .customObservationConvention(customObservationConvention.getIfAvailable(() -> null)) + .batchingStrategy(chromaBatchingStrategy) + .build(); } static class PropertiesChromaConnectionDetails implements ChromaConnectionDetails { diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreProperties.java index 2ce5e6dea..2a54f449a 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreProperties.java @@ -17,7 +17,7 @@ package org.springframework.ai.autoconfigure.vectorstore.chroma; import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties; -import org.springframework.ai.vectorstore.ChromaVectorStore; +import org.springframework.ai.chroma.vectorstore.ChromaVectorStore; import org.springframework.boot.context.properties.ConfigurationProperties; /** diff --git a/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java b/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java index 788352d51..f48f7dca6 100644 --- a/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java +++ b/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java @@ -387,7 +387,7 @@ public class AzureVectorStore extends AbstractObservationVectorStore implements } @Override - public Builder createObservationContextBuilder(String operationName) { + public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) { return VectorStoreObservationContext.builder(VectorStoreProvider.AZURE.value(), operationName) .withCollectionName(this.indexName) diff --git a/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java index df8467fa7..109541f25 100644 --- a/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java +++ b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java @@ -381,7 +381,7 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme } @Override - public Builder createObservationContextBuilder(String operationName) { + public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) { return VectorStoreObservationContext.builder(VectorStoreProvider.CASSANDRA.value(), operationName) .withCollectionName(this.conf.schema.table()) .withDimensions(this.embeddingModel.dimensions()) diff --git a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/ChromaApi.java b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaApi.java similarity index 97% rename from vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/ChromaApi.java rename to vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaApi.java index 4abc485dc..c67a969c8 100644 --- a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/ChromaApi.java +++ b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaApi.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.chroma; +package org.springframework.ai.chroma.vectorstore; import java.util.ArrayList; import java.util.HashMap; @@ -28,7 +28,7 @@ import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; -import org.springframework.ai.chroma.ChromaApi.QueryRequest.Include; +import org.springframework.ai.chroma.vectorstore.ChromaApi.QueryRequest.Include; import org.springframework.http.HttpHeaders; import org.springframework.http.MediaType; import org.springframework.http.client.SimpleClientHttpRequestFactory; @@ -49,10 +49,10 @@ import org.springframework.web.client.RestClient; public class ChromaApi { // Regular expression pattern that looks for a message inside the ValueError(...). - private static Pattern VALUE_ERROR_PATTERN = Pattern.compile("ValueError\\('([^']*)'\\)"); + private static final Pattern VALUE_ERROR_PATTERN = Pattern.compile("ValueError\\('([^']*)'\\)"); // Regular expression pattern that looks for a message. - private static Pattern MESSAGE_ERROR_PATTERN = Pattern.compile("\"message\":\"(.*?)\""); + private static final Pattern MESSAGE_ERROR_PATTERN = Pattern.compile("\"message\":\"(.*?)\""); private final ObjectMapper objectMapper; diff --git a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaFilterExpressionConverter.java b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaFilterExpressionConverter.java similarity index 95% rename from vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaFilterExpressionConverter.java rename to vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaFilterExpressionConverter.java index 526fc30cf..83534f8b1 100644 --- a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaFilterExpressionConverter.java +++ b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaFilterExpressionConverter.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.chroma.vectorstore; import org.springframework.ai.vectorstore.filter.Filter; import org.springframework.ai.vectorstore.filter.converter.PineconeFilterExpressionConverter; diff --git a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaVectorStore.java b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java similarity index 72% rename from vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaVectorStore.java rename to vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java index 75e9c4808..003cc9386 100644 --- a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaVectorStore.java +++ b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.chroma.vectorstore; import java.util.ArrayList; import java.util.HashMap; @@ -27,10 +27,9 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.json.JsonMapper; import io.micrometer.observation.ObservationRegistry; -import org.springframework.ai.chroma.ChromaApi; -import org.springframework.ai.chroma.ChromaApi.AddEmbeddingsRequest; -import org.springframework.ai.chroma.ChromaApi.DeleteEmbeddingsRequest; -import org.springframework.ai.chroma.ChromaApi.Embedding; +import org.springframework.ai.chroma.vectorstore.ChromaApi.AddEmbeddingsRequest; +import org.springframework.ai.chroma.vectorstore.ChromaApi.DeleteEmbeddingsRequest; +import org.springframework.ai.chroma.vectorstore.ChromaApi.Embedding; import org.springframework.ai.document.Document; import org.springframework.ai.document.DocumentMetadata; import org.springframework.ai.embedding.BatchingStrategy; @@ -39,12 +38,16 @@ import org.springframework.ai.embedding.EmbeddingOptionsBuilder; import org.springframework.ai.embedding.TokenCountBatchingStrategy; import org.springframework.ai.observation.conventions.VectorStoreProvider; import org.springframework.ai.util.JacksonUtils; +import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore; import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext; import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention; import org.springframework.beans.factory.InitializingBean; import org.springframework.lang.NonNull; +import org.springframework.lang.Nullable; import org.springframework.util.Assert; import org.springframework.util.CollectionUtils; @@ -65,14 +68,13 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements public static final String DEFAULT_COLLECTION_NAME = "SpringAiCollection"; - private final EmbeddingModel embeddingModel; - private final ChromaApi chromaApi; private final String collectionName; private FilterExpressionConverter filterExpressionConverter; + @Nullable private String collectionId; private final boolean initializeSchema; @@ -83,34 +85,36 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements private boolean initialized = false; + @Deprecated(since = "1.0.0-M5", forRemoval = true) public ChromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi, boolean initializeSchema) { this(embeddingModel, chromaApi, DEFAULT_COLLECTION_NAME, initializeSchema); } + @Deprecated(since = "1.0.0-M5", forRemoval = true) public ChromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi, String collectionName, boolean initializeSchema) { this(embeddingModel, chromaApi, collectionName, initializeSchema, ObservationRegistry.NOOP, null, new TokenCountBatchingStrategy()); } + @Deprecated(since = "1.0.0-M5", forRemoval = true) public ChromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi, String collectionName, boolean initializeSchema, ObservationRegistry observationRegistry, VectorStoreObservationConvention customObservationConvention, BatchingStrategy batchingStrategy) { - super(observationRegistry, customObservationConvention); - - this.embeddingModel = embeddingModel; - this.chromaApi = chromaApi; - this.collectionName = collectionName; - this.initializeSchema = initializeSchema; - this.filterExpressionConverter = new ChromaFilterExpressionConverter(); - this.batchingStrategy = batchingStrategy; - this.objectMapper = JsonMapper.builder().addModules(JacksonUtils.instantiateAvailableModules()).build(); + this(builder(chromaApi).embeddingModel(embeddingModel) + .collectionName(collectionName) + .initializeSchema(initializeSchema) + .observationRegistry(observationRegistry) + .customObservationConvention(customObservationConvention) + .batchingStrategy(batchingStrategy)); } - private ChromaVectorStore(Builder builder) { - super(builder.observationRegistry, builder.customObservationConvention); - this.embeddingModel = builder.embeddingModel; + /** + * @param builder {@link Builder} for chroma vector store + */ + private ChromaVectorStore(ChromaBuilder builder) { + super(builder); this.chromaApi = builder.chromaApi; this.collectionName = builder.collectionName; this.initializeSchema = builder.initializeSchema; @@ -128,9 +132,27 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements } } - public void setFilterExpressionConverter(FilterExpressionConverter filterExpressionConverter) { - Assert.notNull(filterExpressionConverter, "FilterExpressionConverter should not be null."); - this.filterExpressionConverter = filterExpressionConverter; + @Override + public void afterPropertiesSet() throws Exception { + if (!this.initialized) { + var collection = this.chromaApi.getCollection(this.collectionName); + if (collection == null) { + if (this.initializeSchema) { + collection = this.chromaApi + .createCollection(new ChromaApi.CreateCollectionRequest(this.collectionName)); + } + else { + throw new RuntimeException("Collection " + this.collectionName + + " doesn't exist and won't be created as the initializeSchema is set to false."); + } + } + this.collectionId = collection.id(); + this.initialized = true; + } + } + + public static ChromaBuilder builder(ChromaApi chromaApi) { + return new ChromaBuilder(chromaApi); } @Override @@ -217,44 +239,40 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements } } + /** + * @deprecated not used currently anywhere + */ + @Deprecated(forRemoval = true) public String getCollectionName() { return this.collectionName; } + /** + * @deprecated only used in tests + */ + @Deprecated(forRemoval = true) + @Nullable public String getCollectionId() { return this.collectionId; } - @Override - public void afterPropertiesSet() throws Exception { - if (!this.initialized) { - var collection = this.chromaApi.getCollection(this.collectionName); - if (collection == null) { - if (this.initializeSchema) { - collection = this.chromaApi - .createCollection(new ChromaApi.CreateCollectionRequest(this.collectionName)); - } - else { - throw new RuntimeException("Collection " + this.collectionName - + " doesn't exist and won't be created as the initializeSchema is set to false."); - } - } - this.collectionId = collection.id(); - this.initialized = true; - } + /** + * @deprecated in favor the builder method + */ + @Deprecated(forRemoval = true) + public void setFilterExpressionConverter(FilterExpressionConverter filterExpressionConverter) { + Assert.notNull(filterExpressionConverter, "FilterExpressionConverter should not be null."); + this.filterExpressionConverter = filterExpressionConverter; } @Override - public @NonNull VectorStoreObservationContext.Builder createObservationContextBuilder( - @NonNull String operationName) { + public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) { return VectorStoreObservationContext.builder(VectorStoreProvider.CHROMA.value(), operationName) .withDimensions(this.embeddingModel.dimensions()) .withCollectionName(this.collectionName + ":" + this.collectionId); } - public static class Builder { - - private final EmbeddingModel embeddingModel; + public static class ChromaBuilder extends AbstractVectorStoreBuilder { private final ChromaApi chromaApi; @@ -262,57 +280,80 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements private boolean initializeSchema = false; - private ObservationRegistry observationRegistry = ObservationRegistry.NOOP; - - private VectorStoreObservationConvention customObservationConvention = null; - private BatchingStrategy batchingStrategy = new TokenCountBatchingStrategy(); private FilterExpressionConverter filterExpressionConverter = new ChromaFilterExpressionConverter(); private boolean initializeImmediately = false; - public Builder(EmbeddingModel embeddingModel, ChromaApi chromaApi) { - this.embeddingModel = embeddingModel; + public ChromaBuilder(ChromaApi chromaApi) { + Assert.notNull(chromaApi, "ChromaApi must not be null"); this.chromaApi = chromaApi; } - public Builder collectionName(String collectionName) { + /** + * Sets the collection name. + * @param collectionName the name of the collection + * @return the builder instance + * @throws IllegalArgumentException if collectionName is null or empty + */ + public ChromaBuilder collectionName(String collectionName) { + Assert.hasText(collectionName, "collectionName must not be null or empty"); this.collectionName = collectionName; return this; } - public Builder initializeSchema(boolean initializeSchema) { + /** + * Sets whether to initialize the schema. + * @param initializeSchema true to initialize schema, false otherwise + * @return the builder instance + */ + public ChromaBuilder initializeSchema(boolean initializeSchema) { this.initializeSchema = initializeSchema; return this; } - public Builder observationRegistry(ObservationRegistry observationRegistry) { - this.observationRegistry = observationRegistry; - return this; - } - - public Builder customObservationConvention(VectorStoreObservationConvention convention) { - this.customObservationConvention = convention; - return this; - } - - public Builder batchingStrategy(BatchingStrategy batchingStrategy) { + /** + * Sets the batching strategy. + * @param batchingStrategy the batching strategy to use + * @return the builder instance + * @throws IllegalArgumentException if batchingStrategy is null + */ + public ChromaBuilder batchingStrategy(BatchingStrategy batchingStrategy) { + Assert.notNull(batchingStrategy, "batchingStrategy must not be null"); this.batchingStrategy = batchingStrategy; return this; } - public Builder filterExpressionConverter(FilterExpressionConverter converter) { + /** + * Sets the filter expression converter. + * @param converter the filter expression converter to use + * @return the builder instance + * @throws IllegalArgumentException if converter is null + */ + public ChromaBuilder filterExpressionConverter(FilterExpressionConverter converter) { + Assert.notNull(converter, "filterExpressionConverter must not be null"); this.filterExpressionConverter = converter; return this; } - public Builder initializeImmediately(boolean initialize) { + /** + * Sets whether to initialize immediately. + * @param initialize true to initialize immediately, false otherwise + * @return the builder instance + */ + public ChromaBuilder initializeImmediately(boolean initialize) { this.initializeImmediately = initialize; return this; } + /** + * Builds the {@link ChromaVectorStore} instance. + * @return a new ChromaVectorStore instance + * @throws IllegalStateException if the builder is in an invalid state + */ public ChromaVectorStore build() { + validate(); return new ChromaVectorStore(this); } diff --git a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/package-info.java b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/package-info.java new file mode 100644 index 000000000..f6b3ce6c3 --- /dev/null +++ b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/package-info.java @@ -0,0 +1,25 @@ +/* + * Copyright 2023-2024 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. + */ + +/** + * Provides the API for embedding observations. + */ +@NonNullApi +@NonNullFields +package org.springframework.ai.chroma.vectorstore; + +import org.springframework.lang.NonNullApi; +import org.springframework.lang.NonNullFields; diff --git a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/ChromaImage.java b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/ChromaImage.java similarity index 95% rename from vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/ChromaImage.java rename to vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/ChromaImage.java index c9961e5af..a34845426 100644 --- a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/ChromaImage.java +++ b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/ChromaImage.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai; +package org.springframework.ai.chroma; import org.testcontainers.utility.DockerImageName; diff --git a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/BasicAuthChromaWhereIT.java b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/BasicAuthChromaWhereIT.java similarity index 91% rename from vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/BasicAuthChromaWhereIT.java rename to vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/BasicAuthChromaWhereIT.java index d9cbc560d..6444ff795 100644 --- a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/BasicAuthChromaWhereIT.java +++ b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/BasicAuthChromaWhereIT.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.chroma.vectorstore; import java.util.List; import java.util.Map; @@ -26,12 +26,13 @@ import org.testcontainers.junit.jupiter.Container; import org.testcontainers.junit.jupiter.Testcontainers; import org.testcontainers.utility.MountableFile; -import org.springframework.ai.ChromaImage; -import org.springframework.ai.chroma.ChromaApi; +import org.springframework.ai.chroma.ChromaImage; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.openai.OpenAiEmbeddingModel; import org.springframework.ai.openai.api.OpenAiApi; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; import org.springframework.boot.SpringBootConfiguration; import org.springframework.boot.test.context.runner.ApplicationContextRunner; import org.springframework.context.annotation.Bean; @@ -110,7 +111,11 @@ public class BasicAuthChromaWhereIT { @Bean public VectorStore chromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi) { - return new ChromaVectorStore(embeddingModel, chromaApi, "TestCollection", true); + return ChromaVectorStore.builder(chromaApi) + .embeddingModel(embeddingModel) + .collectionName("TestCollection") + .initializeSchema(true) + .build(); } @Bean diff --git a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/ChromaApiIT.java b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaApiIT.java similarity index 90% rename from vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/ChromaApiIT.java rename to vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaApiIT.java index 0b01ae354..79527ec8b 100644 --- a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/ChromaApiIT.java +++ b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaApiIT.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.chroma; +package org.springframework.ai.chroma.vectorstore; import java.util.Collections; import java.util.List; @@ -26,15 +26,14 @@ import org.testcontainers.chromadb.ChromaDBContainer; import org.testcontainers.junit.jupiter.Container; import org.testcontainers.junit.jupiter.Testcontainers; -import org.springframework.ai.ChromaImage; -import org.springframework.ai.chroma.ChromaApi.AddEmbeddingsRequest; -import org.springframework.ai.chroma.ChromaApi.Collection; -import org.springframework.ai.chroma.ChromaApi.GetEmbeddingsRequest; -import org.springframework.ai.chroma.ChromaApi.QueryRequest; +import org.springframework.ai.chroma.ChromaImage; +import org.springframework.ai.chroma.vectorstore.ChromaApi.AddEmbeddingsRequest; +import org.springframework.ai.chroma.vectorstore.ChromaApi.Collection; +import org.springframework.ai.chroma.vectorstore.ChromaApi.GetEmbeddingsRequest; +import org.springframework.ai.chroma.vectorstore.ChromaApi.QueryRequest; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.transformers.TransformersEmbeddingModel; -import org.springframework.ai.vectorstore.ChromaVectorStore; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.boot.SpringBootConfiguration; import org.springframework.boot.test.context.SpringBootTest; @@ -208,7 +207,8 @@ public class ChromaApiIT { assertThat(collection).isNotNull(); assertThat(collection.name()).isEqualTo("test-collection"); - ChromaVectorStore store = new ChromaVectorStore.Builder(this.embeddingModel, this.chromaApi) + ChromaVectorStore store = ChromaVectorStore.builder(this.chromaApi) + .embeddingModel(this.embeddingModel) .collectionName("test-collection") .initializeImmediately(true) .build(); @@ -219,7 +219,8 @@ public class ChromaApiIT { @Test void shouldCreateNewCollectionWhenSchemaInitializationEnabled() { - ChromaVectorStore store = new ChromaVectorStore.Builder(this.embeddingModel, this.chromaApi) + ChromaVectorStore store = new ChromaVectorStore.ChromaBuilder(this.chromaApi) + .embeddingModel(this.embeddingModel) .collectionName("new-collection") .initializeSchema(true) .initializeImmediately(true) @@ -235,12 +236,11 @@ public class ChromaApiIT { @Test void shouldFailWhenCollectionDoesNotExist() { - assertThatThrownBy( - () -> new ChromaVectorStore.Builder(this.embeddingModel, this.chromaApi).collectionName("non-existent") - .initializeSchema(false) - .initializeImmediately(true) - .build()) - .isInstanceOf(IllegalStateException.class) + assertThatThrownBy(() -> new ChromaVectorStore.ChromaBuilder(this.chromaApi).embeddingModel(this.embeddingModel) + .collectionName("non-existent") + .initializeSchema(false) + .initializeImmediately(true) + .build()).isInstanceOf(IllegalStateException.class) .hasMessage("Failed to initialize ChromaVectorStore") .hasCauseInstanceOf(RuntimeException.class) .hasRootCauseMessage( diff --git a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/ChromaVectorStoreIT.java b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java similarity index 96% rename from vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/ChromaVectorStoreIT.java rename to vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java index 5015b1d79..ee9689950 100644 --- a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/ChromaVectorStoreIT.java +++ b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.chroma.vectorstore; import java.util.Collections; import java.util.List; @@ -29,12 +29,13 @@ import org.testcontainers.chromadb.ChromaDBContainer; import org.testcontainers.junit.jupiter.Container; import org.testcontainers.junit.jupiter.Testcontainers; -import org.springframework.ai.ChromaImage; -import org.springframework.ai.chroma.ChromaApi; +import org.springframework.ai.chroma.ChromaImage; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.openai.OpenAiEmbeddingModel; import org.springframework.ai.openai.api.OpenAiApi; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; import org.springframework.boot.SpringBootConfiguration; import org.springframework.boot.test.context.runner.ApplicationContextRunner; import org.springframework.context.annotation.Bean; @@ -251,7 +252,11 @@ public class ChromaVectorStoreIT { @Bean public VectorStore chromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi) { - return new ChromaVectorStore(embeddingModel, chromaApi, "TestCollection", true); + return ChromaVectorStore.builder(chromaApi) + .embeddingModel(embeddingModel) + .collectionName("TestCollection") + .initializeSchema(true) + .build(); } @Bean diff --git a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/ChromaVectorStoreObservationIT.java b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreObservationIT.java similarity index 91% rename from vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/ChromaVectorStoreObservationIT.java rename to vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreObservationIT.java index 72fc9238e..34846632e 100644 --- a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/ChromaVectorStoreObservationIT.java +++ b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreObservationIT.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.chroma.vectorstore; import java.io.IOException; import java.nio.charset.StandardCharsets; @@ -30,8 +30,7 @@ import org.testcontainers.chromadb.ChromaDBContainer; import org.testcontainers.junit.jupiter.Container; import org.testcontainers.junit.jupiter.Testcontainers; -import org.springframework.ai.ChromaImage; -import org.springframework.ai.chroma.ChromaApi; +import org.springframework.ai.chroma.ChromaImage; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.embedding.TokenCountBatchingStrategy; @@ -39,6 +38,8 @@ import org.springframework.ai.observation.conventions.SpringAiKind; import org.springframework.ai.observation.conventions.VectorStoreProvider; import org.springframework.ai.openai.OpenAiEmbeddingModel; import org.springframework.ai.openai.api.OpenAiApi; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention; import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.HighCardinalityKeyNames; import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.LowCardinalityKeyNames; @@ -48,6 +49,7 @@ import org.springframework.boot.test.context.runner.ApplicationContextRunner; import org.springframework.context.annotation.Bean; import org.springframework.core.io.DefaultResourceLoader; import org.springframework.http.client.SimpleClientHttpRequestFactory; +import org.springframework.test.util.ReflectionTestUtils; import org.springframework.web.client.RestClient; import static org.assertj.core.api.Assertions.assertThat; @@ -105,7 +107,7 @@ public class ChromaVectorStoreObservationIT { .doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.DB_VECTOR_QUERY_CONTENT.asString()) .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_VECTOR_DIMENSION_COUNT.asString(), "1536") .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_COLLECTION_NAME.asString(), - "TestCollection:" + vectorStore.getCollectionId()) + "TestCollection:" + ReflectionTestUtils.getField(vectorStore, "collectionId")) .doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.DB_NAMESPACE.asString()) .doesNotHaveHighCardinalityKeyValueWithKey( HighCardinalityKeyNames.DB_SEARCH_SIMILARITY_METRIC.asString()) @@ -138,7 +140,7 @@ public class ChromaVectorStoreObservationIT { "What is Great Depression") .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_VECTOR_DIMENSION_COUNT.asString(), "1536") .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_COLLECTION_NAME.asString(), - "TestCollection:" + vectorStore.getCollectionId()) + "TestCollection:" + ReflectionTestUtils.getField(vectorStore, "collectionId")) .doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.DB_NAMESPACE.asString()) .doesNotHaveHighCardinalityKeyValueWithKey( HighCardinalityKeyNames.DB_SEARCH_SIMILARITY_METRIC.asString()) @@ -174,8 +176,13 @@ public class ChromaVectorStoreObservationIT { @Bean public VectorStore chromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi, ObservationRegistry observationRegistry) { - return new ChromaVectorStore(embeddingModel, chromaApi, "TestCollection", true, observationRegistry, null, - new TokenCountBatchingStrategy()); + return ChromaVectorStore.builder(chromaApi) + .embeddingModel(embeddingModel) + .collectionName("TestCollection") + .initializeSchema(true) + .observationRegistry(observationRegistry) + .batchingStrategy(new TokenCountBatchingStrategy()) + .build(); } @Bean diff --git a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/TokenSecuredChromaWhereIT.java b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/TokenSecuredChromaWhereIT.java similarity index 93% rename from vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/TokenSecuredChromaWhereIT.java rename to vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/TokenSecuredChromaWhereIT.java index 2aea2f6a3..bd8220d3a 100644 --- a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/TokenSecuredChromaWhereIT.java +++ b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/TokenSecuredChromaWhereIT.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.chroma.vectorstore; import java.util.List; import java.util.Map; @@ -25,12 +25,13 @@ import org.testcontainers.chromadb.ChromaDBContainer; import org.testcontainers.junit.jupiter.Container; import org.testcontainers.junit.jupiter.Testcontainers; -import org.springframework.ai.ChromaImage; -import org.springframework.ai.chroma.ChromaApi; +import org.springframework.ai.chroma.ChromaImage; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.openai.OpenAiEmbeddingModel; import org.springframework.ai.openai.api.OpenAiApi; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; import org.springframework.boot.SpringBootConfiguration; import org.springframework.boot.test.context.runner.ApplicationContextRunner; import org.springframework.context.annotation.Bean; @@ -143,7 +144,11 @@ public class TokenSecuredChromaWhereIT { @Bean public VectorStore chromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi) { - return new ChromaVectorStore(embeddingModel, chromaApi, "TestCollection", true); + return ChromaVectorStore.builder(chromaApi) + .embeddingModel(embeddingModel) + .collectionName("TestCollection") + .initializeSchema(true) + .build(); } @Bean diff --git a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java index f96feb9a7..2e4cef70c 100644 --- a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java +++ b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java @@ -272,7 +272,7 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp } @Override - public Builder createObservationContextBuilder(String operationName) { + public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) { return VectorStoreObservationContext.builder(VectorStoreProvider.ELASTICSEARCH.value(), operationName) .withCollectionName(this.options.getIndexName()) .withDimensions(this.embeddingModel.dimensions()) diff --git a/vector-stores/spring-ai-hanadb-store/src/main/java/org/springframework/ai/vectorstore/HanaCloudVectorStore.java b/vector-stores/spring-ai-hanadb-store/src/main/java/org/springframework/ai/vectorstore/HanaCloudVectorStore.java index a0d4ddc9c..2aa7f984e 100644 --- a/vector-stores/spring-ai-hanadb-store/src/main/java/org/springframework/ai/vectorstore/HanaCloudVectorStore.java +++ b/vector-stores/spring-ai-hanadb-store/src/main/java/org/springframework/ai/vectorstore/HanaCloudVectorStore.java @@ -172,7 +172,7 @@ public class HanaCloudVectorStore extends AbstractObservationVectorStore { } @Override - public Builder createObservationContextBuilder(String operationName) { + public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) { return VectorStoreObservationContext.builder(VectorStoreProvider.HANA.value(), operationName) .withDimensions(this.embeddingModel.dimensions()) diff --git a/vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/OpenSearchVectorStore.java b/vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/OpenSearchVectorStore.java index 66a8e9e16..a5fc66fa4 100644 --- a/vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/OpenSearchVectorStore.java +++ b/vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/OpenSearchVectorStore.java @@ -274,7 +274,7 @@ public class OpenSearchVectorStore extends AbstractObservationVectorStore implem } @Override - public Builder createObservationContextBuilder(String operationName) { + public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) { return VectorStoreObservationContext.builder(VectorStoreProvider.OPENSEARCH.value(), operationName) .withCollectionName(this.index) .withDimensions(this.embeddingModel.dimensions()) diff --git a/vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/OracleVectorStore.java b/vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/OracleVectorStore.java index 53bb84737..0f20d9a98 100644 --- a/vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/OracleVectorStore.java +++ b/vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/OracleVectorStore.java @@ -549,7 +549,7 @@ public class OracleVectorStore extends AbstractObservationVectorStore implements } @Override - public Builder createObservationContextBuilder(String operationName) { + public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) { return VectorStoreObservationContext.builder(VectorStoreProvider.ORACLE.value(), operationName) .withDimensions(this.embeddingModel.dimensions()) .withCollectionName(this.getTableName())