Add Batching strategy for embedding documents

- When embedding documents, allow batching the documents using some criteria.
 - `BatchingStrategy` interface with a `TokenCountBatchingStrategy` implementation that uses
   the openai max input token size of 8191 as the default.
 - Add a default method in EmbeddingModel to embed document using this new batching strategy.
 - Change `MilvusVectorStore` to make use of this new batching API.
 - Adding unit tests for `TokenCountBatchingStrategy`.
 - Adding openai integration test to call the embed API that uses batching.

Resolves https://github.com/spring-projects/spring-ai/issues/1214

Other vector stores will be updated seperately

models/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 - 2024 the original author or authors.
+ * 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.
@@ -18,23 +18,33 @@ package org.springframework.ai.openai.embedding;
 import org.junit.jupiter.api.Test;
 
 import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
+
+import org.springframework.ai.document.Document;
 import org.springframework.ai.embedding.EmbeddingRequest;
 import org.springframework.ai.embedding.EmbeddingResponse;
+import org.springframework.ai.embedding.TokenCountBatchingStrategy;
 import org.springframework.ai.openai.OpenAiEmbeddingModel;
 import org.springframework.ai.openai.OpenAiEmbeddingOptions;
 import org.springframework.ai.openai.OpenAiTestConfiguration;
+import org.springframework.ai.openai.api.OpenAiApi;
 import org.springframework.ai.openai.testutils.AbstractIT;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.core.io.DefaultResourceLoader;
+import org.springframework.core.io.Resource;
 
+import java.nio.charset.StandardCharsets;
 import java.util.List;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 @SpringBootTest(classes = OpenAiTestConfiguration.class)
 @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
 class EmbeddingIT extends AbstractIT {
 
+	private Resource resource = new DefaultResourceLoader().getResource("classpath:text_source.txt");
+
 	@Autowired
 	private OpenAiEmbeddingModel embeddingModel;
 
@@ -53,6 +63,28 @@ class EmbeddingIT extends AbstractIT {
 		assertThat(embeddingModel.dimensions()).isEqualTo(1536);
 	}
 
+	@Test
+	void embeddingBatchDocuments() throws Exception {
+		assertThat(embeddingModel).isNotNull();
+		List<float[]> embedded = this.embeddingModel.embed(
+				List.of(new Document("Hello world"), new Document("Hello Spring"), new Document("Hello Spring AI!")),
+				OpenAiEmbeddingOptions.builder().withModel(OpenAiApi.DEFAULT_EMBEDDING_MODEL).build(),
+				new TokenCountBatchingStrategy());
+		assertThat(embedded.size()).isEqualTo(3);
+		embedded.forEach(embedding -> assertThat(embedding.length).isEqualTo(this.embeddingModel.dimensions()));
+	}
+
+	@Test
+	void embeddingBatchDocumentsThatExceedTheLimit() throws Exception {
+		assertThat(embeddingModel).isNotNull();
+		String contentAsString = resource.getContentAsString(StandardCharsets.UTF_8);
+		assertThatThrownBy(() -> {
+			embeddingModel.embed(List.of(new Document("Hello World"), new Document(contentAsString)),
+					OpenAiEmbeddingOptions.builder().withModel(OpenAiApi.DEFAULT_EMBEDDING_MODEL).build(),
+					new TokenCountBatchingStrategy());
+		}).isInstanceOf(IllegalArgumentException.class);
+	}
+
 	@Test
 	void embedding3Large() {
 

spring-ai-core/src/main/java/org/springframework/ai/embedding/BatchingStrategy.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 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.embedding;
+
+import java.util.List;
+
+import org.springframework.ai.document.Document;
+
+/**
+ * Contract for batching {@link Document} objects so that the call to embed them could be
+ * optimized.
+ *
+ * @author Soby Chacko
+ * @since 1.0.0
+ */
+public interface BatchingStrategy {
+
+	/**
+	 * {@link EmbeddingModel} implementations can call this method to optimize embedding
+	 * tokens. The incoming collection of {@link Document}s are split into su-batches.
+	 * @param documents to batch
+	 * @return a list of sub-batches that contain {@link Document}s.
+	 */
+	List<List<Document>> batch(List<Document> documents);
+
+}

spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingModel.java

@@ -19,10 +19,18 @@ import org.springframework.ai.document.Document;
 import org.springframework.ai.model.Model;
 import org.springframework.util.Assert;
 
+import java.util.ArrayList;
 import java.util.List;
 
 /**
  * EmbeddingModel is a generic interface for embedding models.
+ *
+ * @author Mark Pollack
+ * @author Christian Tzolov
+ * @author Josh Long
+ * @author Soby Chacko
+ * @since 1.0.0
+ *
  */
 public interface EmbeddingModel extends Model<EmbeddingRequest, EmbeddingResponse> {
 
@@ -61,6 +69,35 @@ public interface EmbeddingModel extends Model<EmbeddingRequest, EmbeddingRespons
 			.toList();
 	}
 
+	/**
+	 * Embeds a batch of {@link Document}s into vectors based on a
+	 * {@link BatchingStrategy}.
+	 * @param documents list of {@link Document}s.
+	 * @param options {@link EmbeddingOptions}.
+	 * @param batchingStrategy {@link BatchingStrategy}.
+	 * @return a list of float[] that represents the vectors for the incoming
+	 * {@link Document}s.
+	 */
+	default List<float[]> embed(List<Document> documents, EmbeddingOptions options, BatchingStrategy batchingStrategy) {
+		Assert.notNull(documents, "Documents must not be null");
+		List<float[]> embeddings = new ArrayList<>();
+
+		List<List<Document>> batch = batchingStrategy.batch(documents);
+
+		for (List<Document> subBatch : batch) {
+			List<String> texts = subBatch.stream().map(Document::getContent).toList();
+			EmbeddingRequest request = new EmbeddingRequest(texts, options);
+			EmbeddingResponse response = this.call(request);
+			for (int i = 0; i < subBatch.size(); i++) {
+				Document document = subBatch.get(i);
+				float[] output = response.getResults().get(i).getOutput();
+				embeddings.add(output);
+				document.setEmbedding(output);
+			}
+		}
+		return embeddings;
+	}
+
 	/**
 	 * Embeds a batch of texts into vectors and returns the {@link EmbeddingResponse}.
 	 * @param texts list of texts to embed.

spring-ai-core/src/main/java/org/springframework/ai/embedding/TokenCountBatchingStrategy.java

@@ -0,0 +1,105 @@
+/*
+ * Copyright 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.embedding;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.springframework.ai.document.ContentFormatter;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.document.MetadataMode;
+import org.springframework.ai.tokenizer.JTokkitTokenCountEstimator;
+import org.springframework.ai.tokenizer.TokenCountEstimator;
+
+import com.knuddels.jtokkit.api.EncodingType;
+
+/**
+ * Token count based strategy implementation for {@link BatchingStrategy}. Using openai
+ * max input token as the default:
+ * https://platform.openai.com/docs/guides/embeddings/embedding-models.
+ *
+ * @author Soby Chacko
+ * @since 1.0.0
+ */
+public class TokenCountBatchingStrategy implements BatchingStrategy {
+
+	/**
+	 * Using openai upper limit of input token count as the default.
+	 */
+	private static final int MAX_INPUT_TOKEN_COUNT = 8191;
+
+	private final TokenCountEstimator tokenCountEstimator;
+
+	private final int maxInputTokenCount;
+
+	private final ContentFormatter contentFormater;
+
+	private final MetadataMode metadataMode;
+
+	public TokenCountBatchingStrategy() {
+		this(EncodingType.CL100K_BASE, MAX_INPUT_TOKEN_COUNT);
+	}
+
+	/**
+	 * @param encodingType {@link EncodingType}
+	 * @param maxInputTokenCount upper limit for input tokens
+	 */
+	public TokenCountBatchingStrategy(EncodingType encodingType, int maxInputTokenCount) {
+		this(encodingType, maxInputTokenCount, Document.DEFAULT_CONTENT_FORMATTER, MetadataMode.NONE);
+	}
+
+	/**
+	 * @param encodingType {@link EncodingType}
+	 * @param maxInputTokenCount upper limit for input tokens
+	 * @param contentFormatter {@link ContentFormatter}
+	 * @param metadataMode {@link MetadataMode}
+	 */
+	public TokenCountBatchingStrategy(EncodingType encodingType, int maxInputTokenCount,
+			ContentFormatter contentFormatter, MetadataMode metadataMode) {
+		this.tokenCountEstimator = new JTokkitTokenCountEstimator(encodingType);
+		this.maxInputTokenCount = (int) Math.round(maxInputTokenCount - (maxInputTokenCount * .1));
+		this.contentFormater = contentFormatter;
+		this.metadataMode = metadataMode;
+	}
+
+	@Override
+	public List<List<Document>> batch(List<Document> documents) {
+		List<List<Document>> batches = new ArrayList<>();
+		int currentSize = 0;
+		List<Document> currentBatch = new ArrayList<>();
+
+		for (Document document : documents) {
+			int tokenCount = this.tokenCountEstimator
+				.estimate(document.getFormattedContent(this.contentFormater, this.metadataMode));
+			if (tokenCount > this.maxInputTokenCount) {
+				throw new IllegalArgumentException(
+						"Tokens in a single document exceeds the maximum number of allowed input tokens");
+			}
+			if (currentSize + tokenCount > maxInputTokenCount) {
+				batches.add(currentBatch);
+				currentBatch.clear();
+				currentSize = 0;
+			}
+			currentBatch.add(document);
+			currentSize += tokenCount;
+		}
+		if (!currentBatch.isEmpty()) {
+			batches.add(currentBatch);
+		}
+		return batches;
+	}
+
+}

spring-ai-core/src/main/java/org/springframework/ai/tokenizer/JTokkitTokenCountEstimator.java

@@ -21,23 +21,27 @@ import com.knuddels.jtokkit.api.Encoding;
 import com.knuddels.jtokkit.api.EncodingType;
 
 import org.springframework.ai.model.Media;
-import org.springframework.ai.model.Content;
 import org.springframework.ai.model.MediaContent;
 import org.springframework.util.CollectionUtils;
 
 /**
+ * Estimates the number of tokens in a given text or message using the JTokkit encoding
+ * library.
+ *
  * @author Christian Tzolov
+ * @author Soby Chacko
+ * @since 1.0.0
  */
 public class JTokkitTokenCountEstimator implements TokenCountEstimator {
 
 	private final Encoding estimator;
 
 	public JTokkitTokenCountEstimator() {
-		this.estimator = Encodings.newLazyEncodingRegistry().getEncoding(EncodingType.CL100K_BASE);
+		this(EncodingType.CL100K_BASE);
 	}
 
-	public JTokkitTokenCountEstimator(Encoding tokenEncoding) {
-		this.estimator = tokenEncoding;
+	public JTokkitTokenCountEstimator(EncodingType tokenEncodingType) {
+		this.estimator = Encodings.newLazyEncodingRegistry().getEncoding(tokenEncodingType);
 	}
 
 	@Override

spring-ai-core/src/test/java/org/springframework/ai/embedding/TokenCountBatchingStrategyTests.java

@@ -0,0 +1,57 @@
+/*
+ * Copyright 2024-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.embedding;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+
+import org.junit.jupiter.api.Test;
+
+import org.springframework.ai.document.Document;
+import org.springframework.core.io.DefaultResourceLoader;
+import org.springframework.core.io.Resource;
+
+/**
+ * Basic unit test for {@link TokenCountBatchingStrategy}.
+ *
+ * @author Soby Chacko
+ */
+public class TokenCountBatchingStrategyTests {
+
+	@Test
+	void batchEmbeddingHappyPath() {
+		TokenCountBatchingStrategy tokenCountBatchingStrategy = new TokenCountBatchingStrategy();
+		List<List<Document>> batch = tokenCountBatchingStrategy.batch(
+				List.of(new Document("Hello world"), new Document("Hello Spring"), new Document("Hello Spring AI!")));
+		assertThat(batch.size()).isEqualTo(1);
+		assertThat(batch.get(0).size()).isEqualTo(3);
+	}
+
+	@Test
+	void batchEmbeddingWithLargeDocumentExceedsMaxTokenSize() throws IOException {
+		Resource resource = new DefaultResourceLoader().getResource("classpath:text_source.txt");
+		String contentAsString = resource.getContentAsString(StandardCharsets.UTF_8);
+		TokenCountBatchingStrategy tokenCountBatchingStrategy = new TokenCountBatchingStrategy();
+		assertThatThrownBy(() -> {
+			tokenCountBatchingStrategy.batch(List.of(new Document(contentAsString)));
+		}).isInstanceOf(IllegalArgumentException.class);
+	}
+
+}

spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreAutoConfiguration.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 - 2024 the original author or authors.
+ * 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.
@@ -15,15 +15,14 @@
  */
 package org.springframework.ai.autoconfigure.vectorstore.milvus;
 
-import java.util.concurrent.TimeUnit;
-
 import io.micrometer.observation.ObservationRegistry;
 import io.milvus.client.MilvusServiceClient;
 import io.milvus.param.ConnectParam;
 import io.milvus.param.IndexType;
 import io.milvus.param.MetricType;
-
+import org.springframework.ai.embedding.BatchingStrategy;
 import org.springframework.ai.embedding.EmbeddingModel;
+import org.springframework.ai.embedding.TokenCountBatchingStrategy;
 import org.springframework.ai.vectorstore.MilvusVectorStore;
 import org.springframework.ai.vectorstore.MilvusVectorStore.MilvusVectorStoreConfig;
 import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
@@ -35,9 +34,12 @@ import org.springframework.boot.context.properties.EnableConfigurationProperties
 import org.springframework.context.annotation.Bean;
 import org.springframework.util.StringUtils;
 
+import java.util.concurrent.TimeUnit;
+
 /**
  * @author Christian Tzolov
  * @author Eddú Meléndez
+ * @author Soby Chacko
  */
 @AutoConfiguration
 @ConditionalOnClass({ MilvusVectorStore.class, EmbeddingModel.class })
@@ -51,10 +53,17 @@ public class MilvusVectorStoreAutoConfiguration {
 		return new PropertiesMilvusServiceClientConnectionDetails(properties);
 	}
 
+	@Bean
+	@ConditionalOnMissingBean(BatchingStrategy.class)
+	BatchingStrategy milvusBatchingStrategy() {
+		return new TokenCountBatchingStrategy();
+	}
+
 	@Bean
 	@ConditionalOnMissingBean
 	public MilvusVectorStore vectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel,
-			MilvusVectorStoreProperties properties, ObjectProvider<ObservationRegistry> observationRegistry,
+			MilvusVectorStoreProperties properties, BatchingStrategy batchingStrategy,
+			ObjectProvider<ObservationRegistry> observationRegistry,
 			ObjectProvider<VectorStoreObservationConvention> customObservationConvention) {
 
 		MilvusVectorStoreConfig config = MilvusVectorStoreConfig.builder()
@@ -67,7 +76,7 @@ public class MilvusVectorStoreAutoConfiguration {
 			.build();
 
 		return new MilvusVectorStore(milvusClient, embeddingModel, config, properties.isInitializeSchema(),
-				observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP),
+				batchingStrategy, observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP),
 				customObservationConvention.getIfAvailable(() -> null));
 	}
 

vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 - 2024 the original author or authors.
+ * 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.
@@ -15,29 +15,7 @@
  */
 package org.springframework.ai.vectorstore;
 
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
-import java.util.stream.Collectors;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-import org.springframework.ai.document.Document;
-import org.springframework.ai.embedding.EmbeddingModel;
-import org.springframework.ai.model.EmbeddingUtils;
-import org.springframework.ai.observation.conventions.VectorStoreProvider;
-import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
-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.util.Assert;
-import org.springframework.util.StringUtils;
-
 import com.alibaba.fastjson.JSONObject;
-
 import io.micrometer.observation.ObservationRegistry;
 import io.milvus.client.MilvusServiceClient;
 import io.milvus.common.clientenum.ConsistencyLevelEnum;
@@ -64,9 +42,33 @@ import io.milvus.param.index.DescribeIndexParam;
 import io.milvus.param.index.DropIndexParam;
 import io.milvus.response.QueryResultsWrapper.RowRecord;
 import io.milvus.response.SearchResultsWrapper;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+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.model.EmbeddingUtils;
+import org.springframework.ai.observation.conventions.VectorStoreProvider;
+import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
+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.util.Assert;
+import org.springframework.util.StringUtils;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.stream.Collectors;
 
 /**
  * @author Christian Tzolov
+ * @author Soby Chacko
  */
 public class MilvusVectorStore extends AbstractObservationVectorStore implements InitializingBean {
 
@@ -104,6 +106,8 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements
 
 	private final boolean initializeSchema;
 
+	private final BatchingStrategy batchingStrategy;
+
 	/**
 	 * Configuration for the Milvus vector store.
 	 */
@@ -134,7 +138,6 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements
 		 * {@return the default config}
 		 */
 		public static MilvusVectorStoreConfig defaultConfig() {
-
 			return builder().build();
 		}
 
@@ -252,20 +255,25 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements
 
 	public MilvusVectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel,
 			boolean initializeSchema) {
-		this(milvusClient, embeddingModel, MilvusVectorStoreConfig.defaultConfig(), initializeSchema);
+		this(milvusClient, embeddingModel, MilvusVectorStoreConfig.defaultConfig(), initializeSchema,
+				new TokenCountBatchingStrategy());
+	}
+
+	public MilvusVectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel, boolean initializeSchema,
+			BatchingStrategy batchingStrategy) {
+		this(milvusClient, embeddingModel, MilvusVectorStoreConfig.defaultConfig(), initializeSchema, batchingStrategy);
 	}
 
 	public MilvusVectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel,
-			MilvusVectorStoreConfig config, boolean initializeSchema) {
-		this(milvusClient, embeddingModel, config, initializeSchema, ObservationRegistry.NOOP, null);
+			MilvusVectorStoreConfig config, boolean initializeSchema, BatchingStrategy batchingStrategy) {
+		this(milvusClient, embeddingModel, config, initializeSchema, batchingStrategy, ObservationRegistry.NOOP, null);
 	}
 
 	public MilvusVectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel,
-			MilvusVectorStoreConfig config, boolean initializeSchema, ObservationRegistry observationRegistry,
-			VectorStoreObservationConvention customObservationConvention) {
+			MilvusVectorStoreConfig config, boolean initializeSchema, BatchingStrategy batchingStrategy,
+			ObservationRegistry observationRegistry, VectorStoreObservationConvention customObservationConvention) {
 
 		super(observationRegistry, customObservationConvention);
-
 		this.initializeSchema = initializeSchema;
 
 		Assert.notNull(milvusClient, "MilvusServiceClient must not be null");
@@ -274,6 +282,7 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements
 		this.milvusClient = milvusClient;
 		this.embeddingModel = embeddingModel;
 		this.config = config;
+		this.batchingStrategy = batchingStrategy;
 	}
 
 	@Override
@@ -286,15 +295,16 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements
 		List<JSONObject> metadataArray = new ArrayList<>();
 		List<List<Float>> embeddingArray = new ArrayList<>();
 
+		// TODO: Need to customize how we pass the embedding options
+		this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
+
 		for (Document document : documents) {
-			float[] embedding = this.embeddingModel.embed(document);
-			document.setEmbedding(embedding);
 			docIdArray.add(document.getId());
 			// Use a (future) DocumentTextLayoutFormatter instance to extract
 			// the content used to compute the embeddings
 			contentArray.add(document.getContent());
 			metadataArray.add(new JSONObject(document.getMetadata()));
-			embeddingArray.add(EmbeddingUtils.toList(embedding));
+			embeddingArray.add(EmbeddingUtils.toList(document.getEmbedding()));
 		}
 
 		List<InsertParam.Field> fields = new ArrayList<>();

vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusEmbeddingDimensionsTests.java

@@ -25,6 +25,7 @@ import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
 
 import org.springframework.ai.embedding.EmbeddingModel;
+import org.springframework.ai.embedding.TokenCountBatchingStrategy;
 import org.springframework.ai.vectorstore.MilvusVectorStore.MilvusVectorStoreConfig;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -56,7 +57,8 @@ public class MilvusEmbeddingDimensionsTests {
 			.withEmbeddingDimension(explicitDimensions)
 			.build();
 
-		var dim = new MilvusVectorStore(milvusClient, embeddingModel, config, true).embeddingDimensions();
+		var dim = new MilvusVectorStore(milvusClient, embeddingModel, config, true, new TokenCountBatchingStrategy())
+			.embeddingDimensions();
 
 		assertThat(dim).isEqualTo(explicitDimensions);
 		verify(embeddingModel, never()).dimensions();
@@ -68,7 +70,7 @@ public class MilvusEmbeddingDimensionsTests {
 
 		MilvusVectorStoreConfig config = MilvusVectorStoreConfig.builder().build();
 
-		var dim = new MilvusVectorStore(milvusClient, embeddingModel, config ,true)
+		var dim = new MilvusVectorStore(milvusClient, embeddingModel, config ,true, new TokenCountBatchingStrategy())
 				.embeddingDimensions();
 
 		assertThat(dim).isEqualTo(969);
@@ -82,7 +84,7 @@ public class MilvusEmbeddingDimensionsTests {
 		when(embeddingModel.dimensions()).thenThrow(new RuntimeException());
 
 		var dim = new MilvusVectorStore(milvusClient, embeddingModel,
-				MilvusVectorStoreConfig.builder().build() ,true)
+				MilvusVectorStoreConfig.builder().build() ,true, new TokenCountBatchingStrategy())
 						.embeddingDimensions();
 
 		assertThat(dim).isEqualTo(MilvusVectorStore.OPENAI_EMBEDDING_DIMENSION_SIZE);

vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java

@@ -34,6 +34,7 @@ import org.testcontainers.junit.jupiter.Testcontainers;
 
 import org.springframework.ai.document.Document;
 import org.springframework.ai.embedding.EmbeddingModel;
+import org.springframework.ai.embedding.TokenCountBatchingStrategy;
 import org.springframework.ai.openai.OpenAiEmbeddingModel;
 import org.springframework.ai.openai.api.OpenAiApi;
 import org.springframework.ai.vectorstore.MilvusVectorStore.MilvusVectorStoreConfig;
@@ -265,7 +266,7 @@ public class MilvusVectorStoreIT {
 				.withIndexType(IndexType.IVF_FLAT)
 				.withMetricType(metricType)
 				.build();
-			return new MilvusVectorStore(milvusClient, embeddingModel, config, true);
+			return new MilvusVectorStore(milvusClient, embeddingModel, config, true, new TokenCountBatchingStrategy());
 		}
 
 		@Bean

vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreObservationIT.java

@@ -25,6 +25,7 @@ import java.util.Map;
 import org.junit.jupiter.api.Test;
 import org.springframework.ai.document.Document;
 import org.springframework.ai.embedding.EmbeddingModel;
+import org.springframework.ai.embedding.TokenCountBatchingStrategy;
 import org.springframework.ai.observation.conventions.VectorStoreProvider;
 import org.springframework.ai.openai.OpenAiEmbeddingModel;
 import org.springframework.ai.openai.api.OpenAiApi;
@@ -158,7 +159,8 @@ public class MilvusVectorStoreObservationIT {
 				.withIndexType(IndexType.IVF_FLAT)
 				.withMetricType(MetricType.COSINE)
 				.build();
-			return new MilvusVectorStore(milvusClient, embeddingModel, config, true, observationRegistry, null);
+			return new MilvusVectorStore(milvusClient, embeddingModel, config, true, new TokenCountBatchingStrategy(),
+					observationRegistry, null);
 		}
 
 		@Bean