Chroma vector store builder refactoring

spring-ai-core/src/main/java/org/springframework/ai/vectorstore/AbstractVectorStoreBuilder.java

@@ -0,0 +1,103 @@
+/*
+ * 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 <T> the concrete builder type, enabling method chaining with the correct return
+ * type
+ */
+public abstract class AbstractVectorStoreBuilder<T extends AbstractVectorStoreBuilder<T>>
+		implements VectorStore.Builder<T> {
+
+	private final EmbeddingModel embeddingModel;
+
+	protected ObservationRegistry observationRegistry = ObservationRegistry.NOOP;
+
+	@Nullable
+	protected VectorStoreObservationConvention customObservationConvention;
+
+	/**
+	 * Creates a new builder instance with the required embedding model.
+	 * @param embeddingModel the embedding model to use for vector operations
+	 * @throws IllegalArgumentException if embeddingModel is null
+	 */
+	protected AbstractVectorStoreBuilder(EmbeddingModel embeddingModel) {
+		Assert.notNull(embeddingModel, "EmbeddingModel must not be null");
+		this.embeddingModel = embeddingModel;
+	}
+
+	/**
+	 * 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) {
+		this.observationRegistry = observationRegistry;
+		return self();
+	}
+
+	@Override
+	public T customObservationConvention(VectorStoreObservationConvention convention) {
+		this.customObservationConvention = convention;
+		return self();
+	}
+
+	/**
+	 * Returns the configured embedding model.
+	 * @return the embedding model
+	 */
+	@Override
+	public EmbeddingModel embeddingModel() {
+		return this.embeddingModel;
+	}
+
+	/**
+	 * Returns the configured observation registry.
+	 * @return the observation registry, never null
+	 */
+	@Override
+	public ObservationRegistry observationRegistry() {
+		return this.observationRegistry;
+	}
+
+	/**
+	 * Returns the configured custom observation convention.
+	 * @return the custom observation convention, may be null
+	 */
+	@Override
+	@Nullable
+	public VectorStoreObservationConvention customObservationConvention() {
+		return this.customObservationConvention;
+	}
+
+}

spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java

@@ -19,8 +19,14 @@
 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,56 @@ default List<Document> similaritySearch(String query) {
 		return this.similaritySearch(SearchRequest.query(query));
 	}
 
+	/**
+	 * Builder interface for creating VectorStore instances. Implements a fluent builder
+	 * pattern for configuring observation-related settings.
+	 *
+	 * @param <T> the concrete builder type, enabling method chaining with the correct
+	 * return type
+	 */
+	interface Builder<T extends Builder<T>> {
+
+		/**
+		 * 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);
+
+		/**
+		 * Returns the configured embedding model.
+		 * @return the embedding model, never null
+		 */
+		EmbeddingModel embeddingModel();
+
+		/**
+		 * Returns the configured observation registry.
+		 * @return the observation registry, never null
+		 */
+		ObservationRegistry observationRegistry();
+
+		/**
+		 * Returns the configured custom observation convention.
+		 * @return the custom observation convention, may be null
+		 */
+		@Nullable
+		VectorStoreObservationConvention customObservationConvention();
+
+		/**
+		 * Builds and returns a new VectorStore instance with the configured settings.
+		 * @return a new VectorStore instance
+		 */
+		VectorStore build();
+
+	}
+
 }

spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/AbstractObservationVectorStore.java

@@ -22,6 +22,8 @@
 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 @@
  * 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.embeddingModel(), builder.observationRegistry(), builder.customObservationConvention());
+	}
+
 	/**
 	 * Create a new {@link AbstractObservationVectorStore} instance.
 	 * @param documents the documents to add

spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationContext.java

@@ -194,7 +194,7 @@ public Builder withDimensions(Integer dimensions) {
 			return this;
 		}
 
-		public Builder withFieldName(String fieldName) {
+		public Builder withFieldName(@Nullable String fieldName) {
 			this.context.setFieldName(fieldName);
 			return this;
 		}

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

@@ -19,11 +19,11 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import io.micrometer.observation.ObservationRegistry;
 
-import org.springframework.ai.chroma.ChromaApi;
+import org.springframework.ai.chroma.api.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,13 @@ public ChromaVectorStore vectorStore(EmbeddingModel embeddingModel, ChromaApi ch
 			ChromaVectorStoreProperties storeProperties, ObjectProvider<ObservationRegistry> observationRegistry,
 			ObjectProvider<VectorStoreObservationConvention> customObservationConvention,
 			BatchingStrategy chromaBatchingStrategy) {
-		return new ChromaVectorStore(embeddingModel, chromaApi, storeProperties.getCollectionName(),
-				storeProperties.isInitializeSchema(), observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP),
-				customObservationConvention.getIfAvailable(() -> null), chromaBatchingStrategy);
+		return ChromaVectorStore.builder(embeddingModel, chromaApi)
+			.collectionName(storeProperties.getCollectionName())
+			.initializeSchema(storeProperties.isInitializeSchema())
+			.observationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP))
+			.customObservationConvention(customObservationConvention.getIfAvailable(() -> null))
+			.batchingStrategy(chromaBatchingStrategy)
+			.build();
 	}
 
 	static class PropertiesChromaConnectionDetails implements ChromaConnectionDetails {

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;
 
 /**

vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java

@@ -387,7 +387,7 @@ public void afterPropertiesSet() throws Exception {
 	}
 
 	@Override
-	public Builder createObservationContextBuilder(String operationName) {
+	public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
 
 		return VectorStoreObservationContext.builder(VectorStoreProvider.AZURE.value(), operationName)
 			.withCollectionName(this.indexName)

vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java

@@ -383,7 +383,7 @@ private String getDocumentId(Row row) {
 	}
 
 	@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())

vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/ChromaApi.java → vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/api/ChromaApi.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.springframework.ai.chroma;
+package org.springframework.ai.chroma.api;
 
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -28,7 +28,7 @@
 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.api.ChromaApi.QueryRequest.Include;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.MediaType;
 import org.springframework.http.client.SimpleClientHttpRequestFactory;
@@ -49,10 +49,10 @@
 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;
 

vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaFilterExpressionConverter.java → 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;