Improve ChromaVectorStore with builder pattern and package organization

- Introduce builder pattern for ChromaVectorStore configuration
- Move Chroma-related classes to org.springframework.ai.chroma package
- Add proper null-safety annotations and validation
- Deprecate direct constructors in favor of builder
- Improve JavaDoc documentation
- Clean up code organization and imports

Author: sobychacko

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

@@ -48,7 +48,7 @@ public abstract class AbstractObservationVectorStore implements VectorStore {
 	 * @param customObservationConvention the custom observation convention to use
 	 */
 	public AbstractObservationVectorStore(ObservationRegistry observationRegistry,
-			VectorStoreObservationConvention customObservationConvention) {
+			@Nullable VectorStoreObservationConvention customObservationConvention) {
 		this.observationRegistry = observationRegistry;
 		this.customObservationConvention = customObservationConvention;
 	}

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-chroma-store/src/main/java/org/springframework/ai/chroma/ChromaApi.java renamed to 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 renamed to 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;

vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaVectorStore.java renamed to 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,23 +27,26 @@
 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.api.ChromaApi;
+import org.springframework.ai.chroma.api.ChromaApi.AddEmbeddingsRequest;
+import org.springframework.ai.chroma.api.ChromaApi.DeleteEmbeddingsRequest;
+import org.springframework.ai.chroma.api.ChromaApi.Embedding;
 import org.springframework.ai.document.Document;
 import org.springframework.ai.embedding.BatchingStrategy;
 import org.springframework.ai.embedding.EmbeddingModel;
 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.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;
 
@@ -73,6 +76,7 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements
 
 	private FilterExpressionConverter filterExpressionConverter;
 
+	@Nullable
 	private String collectionId;
 
 	private final boolean initializeSchema;
@@ -83,19 +87,31 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements
 
 	private boolean initialized = false;
 
+	/**
+	 * @deprecated in favor of the builder method
+	 */
+	@Deprecated(forRemoval = true)
 	public ChromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi, boolean initializeSchema) {
 		this(embeddingModel, chromaApi, DEFAULT_COLLECTION_NAME, initializeSchema);
 	}
 
+	/**
+	 * @deprecated in favor of the builder method
+	 */
+	@Deprecated(forRemoval = true)
 	public ChromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi, String collectionName,
 			boolean initializeSchema) {
 		this(embeddingModel, chromaApi, collectionName, initializeSchema, ObservationRegistry.NOOP, null,
 				new TokenCountBatchingStrategy());
 	}
 
+	/**
+	 * @deprecated in favor of the builder method
+	 */
+	@Deprecated(forRemoval = true)
 	public ChromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi, String collectionName,
 			boolean initializeSchema, ObservationRegistry observationRegistry,
-			VectorStoreObservationConvention customObservationConvention, BatchingStrategy batchingStrategy) {
+			@Nullable VectorStoreObservationConvention customObservationConvention, BatchingStrategy batchingStrategy) {
 
 		super(observationRegistry, customObservationConvention);
 
@@ -108,6 +124,9 @@ public ChromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi, Str
 		this.objectMapper = JsonMapper.builder().addModules(JacksonUtils.instantiateAvailableModules()).build();
 	}
 
+	/**
+	 * @param builder {@link Builder} for chroma vector store
+	 */
 	private ChromaVectorStore(Builder builder) {
 		super(builder.observationRegistry, builder.customObservationConvention);
 		this.embeddingModel = builder.embeddingModel;
@@ -128,9 +147,27 @@ private ChromaVectorStore(Builder builder) {
 		}
 	}
 
-	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 Builder builder(EmbeddingModel embeddingModel, ChromaApi chromaApi) {
+		return new Builder(embeddingModel, chromaApi);
 	}
 
 	@Override
@@ -212,36 +249,34 @@ private Map<String, Object> jsonToMap(String jsonText) {
 		}
 	}
 
+	/**
+	 * @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)
@@ -260,6 +295,7 @@ public static class Builder {
 
 		private ObservationRegistry observationRegistry = ObservationRegistry.NOOP;
 
+		@Nullable
 		private VectorStoreObservationConvention customObservationConvention = null;
 
 		private BatchingStrategy batchingStrategy = new TokenCountBatchingStrategy();
@@ -268,46 +304,103 @@ public static class Builder {
 
 		private boolean initializeImmediately = false;
 
+		/**
+		 * Constructor with required parameters.
+		 * @param embeddingModel The {@link EmbeddingModel} to use
+		 * @param chromaApi The {@link ChromaApi} instance
+		 * @throws IllegalArgumentException if any required parameter is null
+		 */
 		public Builder(EmbeddingModel embeddingModel, ChromaApi chromaApi) {
+			Assert.notNull(embeddingModel, "embeddingModel must not be null");
+			Assert.notNull(chromaApi, "chromaApi must not be null");
 			this.embeddingModel = embeddingModel;
 			this.chromaApi = chromaApi;
 		}
 
+		/**
+		 * Sets the collection name.
+		 * @param collectionName the name of the collection
+		 * @return the builder instance
+		 * @throws IllegalArgumentException if collectionName is null or empty
+		 */
 		public Builder collectionName(String collectionName) {
+			Assert.hasText(collectionName, "collectionName must not be null or empty");
 			this.collectionName = collectionName;
 			return this;
 		}
 
+		/**
+		 * Sets whether to initialize the schema.
+		 * @param initializeSchema true to initialize schema, false otherwise
+		 * @return the builder instance
+		 */
 		public Builder initializeSchema(boolean initializeSchema) {
 			this.initializeSchema = initializeSchema;
 			return this;
 		}
 
+		/**
+		 * Sets the {@link ObservationRegistry}
+		 * @param observationRegistry the observation registry to use
+		 * @return the builder instance
+		 * @throws IllegalArgumentException if observationRegistry is null
+		 */
 		public Builder observationRegistry(ObservationRegistry observationRegistry) {
+			Assert.notNull(observationRegistry, "observationRegistry must not be null");
 			this.observationRegistry = observationRegistry;
 			return this;
 		}
 
-		public Builder customObservationConvention(VectorStoreObservationConvention convention) {
+		/**
+		 * Sets the custom observation convention
+		 * {@link VectorStoreObservationConvention}.
+		 * @param convention the custom observation convention to use
+		 * @return the builder instance
+		 */
+		public Builder customObservationConvention(@Nullable VectorStoreObservationConvention convention) {
 			this.customObservationConvention = convention;
 			return this;
 		}
 
+		/**
+		 * Sets the batching strategy.
+		 * @param batchingStrategy the batching strategy to use
+		 * @return the builder instance
+		 * @throws IllegalArgumentException if batchingStrategy is null
+		 */
 		public Builder batchingStrategy(BatchingStrategy batchingStrategy) {
+			Assert.notNull(batchingStrategy, "batchingStrategy must not be null");
 			this.batchingStrategy = batchingStrategy;
 			return this;
 		}
 
+		/**
+		 * Sets the filter expression converter.
+		 * @param converter the filter expression converter to use
+		 * @return the builder instance
+		 * @throws IllegalArgumentException if converter is null
+		 */
 		public Builder filterExpressionConverter(FilterExpressionConverter converter) {
+			Assert.notNull(converter, "filterExpressionConverter must not be null");
 			this.filterExpressionConverter = converter;
 			return this;
 		}
 
+		/**
+		 * Sets whether to initialize immediately.
+		 * @param initialize true to initialize immediately, false otherwise
+		 * @return the builder instance
+		 */
 		public Builder 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() {
 			return new ChromaVectorStore(this);
 		}