Move batching strategy to base vector store builder

Moving BatchingStrategy configuration from individual vector store implementations
to the base AbstractVectorStoreBuilder to reduce code duplication and provide consistent
batching behavior across all vector stores.

The default TokenCountBatchingStrategy is now set in the base builder class.

Author: sobychacko

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

@@ -18,7 +18,9 @@
 
 import io.micrometer.observation.ObservationRegistry;
 
+import org.springframework.ai.embedding.BatchingStrategy;
 import org.springframework.ai.embedding.EmbeddingModel;
+import org.springframework.ai.embedding.TokenCountBatchingStrategy;
 import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
@@ -40,6 +42,8 @@ public abstract class AbstractVectorStoreBuilder<T extends AbstractVectorStoreBu
 	@Nullable
 	protected VectorStoreObservationConvention customObservationConvention;
 
+	protected BatchingStrategy batchingStrategy = new TokenCountBatchingStrategy();
+
 	public AbstractVectorStoreBuilder(EmbeddingModel embeddingModel) {
 		Assert.notNull(embeddingModel, "EmbeddingModel must be configured");
 		this.embeddingModel = embeddingModel;
@@ -49,6 +53,10 @@ public EmbeddingModel getEmbeddingModel() {
 		return this.embeddingModel;
 	}
 
+	public BatchingStrategy getBatchingStrategy() {
+		return this.batchingStrategy;
+	}
+
 	public ObservationRegistry getObservationRegistry() {
 		return this.observationRegistry;
 	}
@@ -81,4 +89,15 @@ public T customObservationConvention(@Nullable VectorStoreObservationConvention
 		return self();
 	}
 
+	/**
+	 * Sets the batching strategy.
+	 * @param batchingStrategy the strategy to use
+	 * @return the builder instance
+	 */
+	public T batchingStrategy(BatchingStrategy batchingStrategy) {
+		Assert.notNull(batchingStrategy, "BatchingStrategy must not be null");
+		this.batchingStrategy = batchingStrategy;
+		return self();
+	}
+
 }

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

@@ -23,6 +23,7 @@
 
 import org.springframework.ai.document.Document;
 import org.springframework.ai.document.DocumentWriter;
+import org.springframework.ai.embedding.BatchingStrategy;
 import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention;
 import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
 import org.springframework.lang.Nullable;
@@ -108,6 +109,13 @@ interface Builder<T extends Builder<T>> {
 		 */
 		T customObservationConvention(VectorStoreObservationConvention convention);
 
+		/**
+		 * Sets the batching strategy.
+		 * @param batchingStrategy the strategy to use
+		 * @return the builder instance for method chaining
+		 */
+		T batchingStrategy(BatchingStrategy batchingStrategy);
+
 		/**
 		 * Builds and returns a new VectorStore instance with the configured settings.
 		 * @return a new VectorStore instance

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

@@ -22,6 +22,7 @@
 import io.micrometer.observation.ObservationRegistry;
 
 import org.springframework.ai.document.Document;
+import org.springframework.ai.embedding.BatchingStrategy;
 import org.springframework.ai.embedding.EmbeddingModel;
 import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder;
 import org.springframework.ai.vectorstore.SearchRequest;
@@ -47,11 +48,14 @@ public abstract class AbstractObservationVectorStore implements VectorStore {
 
 	protected final EmbeddingModel embeddingModel;
 
+	protected final BatchingStrategy batchingStrategy;
+
 	private AbstractObservationVectorStore(EmbeddingModel embeddingModel, ObservationRegistry observationRegistry,
-			@Nullable VectorStoreObservationConvention customObservationConvention) {
+			@Nullable VectorStoreObservationConvention customObservationConvention, BatchingStrategy batchingStrategy) {
 		this.embeddingModel = embeddingModel;
 		this.observationRegistry = observationRegistry;
 		this.customObservationConvention = customObservationConvention;
+		this.batchingStrategy = batchingStrategy;
 	}
 
 	/**
@@ -60,7 +64,8 @@ private AbstractObservationVectorStore(EmbeddingModel embeddingModel, Observatio
 	 * @param builder the builder containing configuration settings
 	 */
 	public AbstractObservationVectorStore(AbstractVectorStoreBuilder<?> builder) {
-		this(builder.getEmbeddingModel(), builder.getObservationRegistry(), builder.getCustomObservationConvention());
+		this(builder.getEmbeddingModel(), builder.getObservationRegistry(), builder.getCustomObservationConvention(),
+				builder.getBatchingStrategy());
 	}
 
 	/**

vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/cosmosdb/CosmosDBVectorStore.java

@@ -51,24 +51,20 @@
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.ObjectNode;
-import io.micrometer.observation.ObservationRegistry;
 import org.apache.commons.lang3.tuple.ImmutablePair;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Flux;
 
 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.vectorstore.AbstractVectorStoreBuilder;
 import org.springframework.ai.vectorstore.SearchRequest;
 import org.springframework.ai.vectorstore.filter.Filter;
 import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
 import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
-import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 
@@ -98,8 +94,6 @@ public class CosmosDBVectorStore extends AbstractObservationVectorStore implemen
 
 	private final List<String> metadataFieldsList;
 
-	private final BatchingStrategy batchingStrategy;
-
 	private CosmosAsyncContainer container;
 
 	/**
@@ -122,7 +116,6 @@ protected CosmosDBVectorStore(Builder builder) {
 		this.vectorStoreThroughput = builder.vectorStoreThroughput;
 		this.vectorDimensions = builder.vectorDimensions;
 		this.metadataFieldsList = builder.metadataFieldsList;
-		this.batchingStrategy = builder.batchingStrategy;
 
 		cosmosClient.createDatabaseIfNotExists(databaseName).block();
 		initializeContainer(containerName, databaseName, vectorStoreThroughput, vectorDimensions, partitionKeyPath);
@@ -404,8 +397,6 @@ public static class Builder extends AbstractVectorStoreBuilder<Builder> {
 
 		private List<String> metadataFieldsList = new ArrayList<>();
 
-		private BatchingStrategy batchingStrategy = new TokenCountBatchingStrategy();
-
 		private Builder(CosmosAsyncClient cosmosClient, EmbeddingModel embeddingModel) {
 			super(embeddingModel);
 			Assert.notNull(cosmosClient, "CosmosClient must not be null");
@@ -483,18 +474,6 @@ public Builder metadataFields(List<String> metadataFieldsList) {
 			return this;
 		}
 
-		/**
-		 * Sets the batching strategy.
-		 * @param batchingStrategy the 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;
-		}
-
 		@Override
 		public CosmosDBVectorStore build() {
 			return new CosmosDBVectorStore(this);

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

@@ -42,16 +42,13 @@
 import com.azure.search.documents.models.SearchOptions;
 import com.azure.search.documents.models.VectorSearchOptions;
 import com.azure.search.documents.models.VectorizedQuery;
-import io.micrometer.observation.ObservationRegistry;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import org.springframework.ai.document.Document;
 import org.springframework.ai.document.DocumentMetadata;
-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.model.EmbeddingUtils;
 import org.springframework.ai.observation.conventions.VectorStoreProvider;
 import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
@@ -60,7 +57,6 @@
 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.Nullable;
 import org.springframework.util.Assert;
@@ -110,8 +106,6 @@ public class AzureVectorStore extends AbstractObservationVectorStore implements
 
 	private final boolean initializeSchema;
 
-	private final BatchingStrategy batchingStrategy;
-
 	/**
 	 * List of metadata fields (as field name and type) that can be used in similarity
 	 * search query filter expressions. The {@link Document#getMetadata()} can contain
@@ -146,7 +140,6 @@ protected AzureVectorStore(Builder builder) {
 		this.searchIndexClient = builder.searchIndexClient;
 		this.initializeSchema = builder.initializeSchema;
 		this.filterMetadataFields = builder.filterMetadataFields;
-		this.batchingStrategy = builder.batchingStrategy;
 		this.defaultTopK = builder.defaultTopK;
 		this.defaultSimilarityThreshold = builder.defaultSimilarityThreshold;
 		this.indexName = builder.indexName;
@@ -389,8 +382,6 @@ public static class Builder extends AbstractVectorStoreBuilder<Builder> {
 
 		private List<MetadataField> filterMetadataFields = List.of();
 
-		private BatchingStrategy batchingStrategy = new TokenCountBatchingStrategy();
-
 		private int defaultTopK = DEFAULT_TOP_K;
 
 		private Double defaultSimilarityThreshold = DEFAULT_SIMILARITY_THRESHOLD;
@@ -423,17 +414,6 @@ public Builder filterMetadataFields(List<MetadataField> filterMetadataFields) {
 			return this;
 		}
 
-		/**
-		 * Sets the batching strategy.
-		 * @param batchingStrategy the strategy to use
-		 * @return the builder instance
-		 */
-		public Builder batchingStrategy(BatchingStrategy batchingStrategy) {
-			Assert.notNull(batchingStrategy, "BatchingStrategy must not be null");
-			this.batchingStrategy = batchingStrategy;
-			return this;
-		}
-
 		/**
 		 * Sets the index name for the Azure Vector Store.
 		 * @param indexName the name of the index to use

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

@@ -217,8 +217,6 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme
 
 	private final boolean closeSessionOnClose;
 
-	private final BatchingStrategy batchingStrategy;
-
 	private final ConcurrentMap<Set<String>, PreparedStatement> addStmts = new ConcurrentHashMap<>();
 
 	private final PreparedStatement deleteStmt;
@@ -239,7 +237,6 @@ protected CassandraVectorStore(Builder builder) {
 		this.primaryKeyTranslator = builder.primaryKeyTranslator;
 		this.executor = Executors.newFixedThreadPool(builder.fixedThreadPoolExecutorSize);
 		this.closeSessionOnClose = builder.closeSessionOnClose;
-		this.batchingStrategy = builder.batchingStrategy;
 
 		ensureSchemaExists(embeddingModel.dimensions());
 		prepareAddStatement(Set.of());
@@ -777,8 +774,6 @@ public static class Builder extends AbstractVectorStoreBuilder<Builder> {
 
 		private int fixedThreadPoolExecutorSize = DEFAULT_ADD_CONCURRENCY;
 
-		private BatchingStrategy batchingStrategy = new TokenCountBatchingStrategy();
-
 		private FilterExpressionConverter filterExpressionConverter;
 
 		private DocumentIdTranslator documentIdTranslator = (String id) -> List.of(id);
@@ -917,18 +912,6 @@ public Builder disallowSchemaChanges(boolean disallowSchemaChanges) {
 			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

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

@@ -79,8 +79,6 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements
 
 	private final boolean initializeSchema;
 
-	private final BatchingStrategy batchingStrategy;
-
 	private final ObjectMapper objectMapper;
 
 	private boolean initialized = false;
@@ -95,7 +93,6 @@ protected ChromaVectorStore(Builder builder) {
 		this.collectionName = builder.collectionName;
 		this.initializeSchema = builder.initializeSchema;
 		this.filterExpressionConverter = builder.filterExpressionConverter;
-		this.batchingStrategy = builder.batchingStrategy;
 		this.objectMapper = JsonMapper.builder().addModules(JacksonUtils.instantiateAvailableModules()).build();
 
 		if (builder.initializeImmediately) {
@@ -232,8 +229,6 @@ public static class Builder extends AbstractVectorStoreBuilder<Builder> {
 
 		private boolean initializeSchema = false;
 
-		private BatchingStrategy batchingStrategy = new TokenCountBatchingStrategy();
-
 		private FilterExpressionConverter filterExpressionConverter = new ChromaFilterExpressionConverter();
 
 		private boolean initializeImmediately = false;
@@ -266,18 +261,6 @@ public Builder initializeSchema(boolean initializeSchema) {
 			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

vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStore.java

@@ -164,8 +164,6 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp
 
 	private final boolean initializeSchema;
 
-	private final BatchingStrategy batchingStrategy;
-
 	protected ElasticsearchVectorStore(Builder builder) {
 		super(builder);
 
@@ -174,7 +172,6 @@ protected ElasticsearchVectorStore(Builder builder) {
 		this.initializeSchema = builder.initializeSchema;
 		this.options = builder.options;
 		this.filterExpressionConverter = builder.filterExpressionConverter;
-		this.batchingStrategy = builder.batchingStrategy;
 
 		String version = Version.VERSION == null ? "Unknown" : Version.VERSION.toString();
 		this.elasticsearchClient = new ElasticsearchClient(new RestClientTransport(builder.restClient,
@@ -371,8 +368,6 @@ public static class Builder extends AbstractVectorStoreBuilder<Builder> {
 
 		private boolean initializeSchema = false;
 
-		private BatchingStrategy batchingStrategy = new TokenCountBatchingStrategy();
-
 		private FilterExpressionConverter filterExpressionConverter = new ElasticsearchAiSearchFilterExpressionConverter();
 
 		/**
@@ -408,18 +403,6 @@ public Builder initializeSchema(boolean initializeSchema) {
 			return this;
 		}
 
-		/**
-		 * Sets the batching strategy for vector operations.
-		 * @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