Fix TokenCountBatchingStrategy ignoring metadata when sizing batches

spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc

@@ -236,7 +236,10 @@ In this configuration:
 2. `8000`: Sets the maximum input token count. This value should be less than or equal to the maximum context window size of your embedding model.
 3. `0.1`: Sets the reserve percentage. The percentage of tokens to reserve from the max input token count. This creates a buffer for potential token count increases during processing.
 
-By default, this constructor uses `Document.DEFAULT_CONTENT_FORMATTER` for content formatting and `MetadataMode.NONE` for metadata handling. If you need to customize these parameters, you can use the full constructor with additional parameters.
+By default, this constructor uses `Document.DEFAULT_CONTENT_FORMATTER` for content formatting and `MetadataMode.EMBED` for metadata handling.
+For a metadata-aware embedding model that includes document metadata in the text it submits, estimating with `MetadataMode.NONE` undercounts that text and can produce batches larger than the configured limit.
+Applications whose embedding model embeds the document text only can pass `MetadataMode.NONE` to the full constructor.
+If you need to customize these parameters, you can use the full constructor with additional parameters.
 
 Once defined, this custom `TokenCountBatchingStrategy` bean will be automatically used by the `EmbeddingModel` implementations in your application, replacing the default strategy.
 
@@ -253,7 +256,7 @@ TokenCountBatchingStrategy strategy = new TokenCountBatchingStrategy(
     8000,  // maxInputTokenCount
     0.1,   // reservePercentage
     Document.DEFAULT_CONTENT_FORMATTER,
-    MetadataMode.NONE
+    MetadataMode.EMBED
 );
 ----
 

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

@@ -44,11 +44,19 @@
  * The strategy batches documents based on their token counts, ensuring that each batch
  * does not exceed the calculated max input token count.
  *
+ * Token counts are estimated with {@link MetadataMode#EMBED} by default. For a
+ * metadata-aware {@link EmbeddingModel} that includes document metadata in the text it
+ * submits, estimating with {@link MetadataMode#NONE} undercounts that text and can
+ * produce batches larger than the configured limit. The constructor taking an explicit
+ * {@link MetadataMode} still accepts {@link MetadataMode#NONE} for models that embed the
+ * document text only.
+ *
  * @author Soby Chacko
  * @author Mark Pollack
  * @author Laura Trotta
  * @author Jihoon Kim
  * @author Yanming Zhou
+ * @author Subhash Polisetti
  * @since 1.0.0
  */
 public class TokenCountBatchingStrategy implements BatchingStrategy {
@@ -84,7 +92,7 @@ public TokenCountBatchingStrategy() {
 	 */
 	public TokenCountBatchingStrategy(EncodingType encodingType, int maxInputTokenCount, double reservePercentage) {
 		this(encodingType, maxInputTokenCount, reservePercentage, Document.DEFAULT_CONTENT_FORMATTER,
-				MetadataMode.NONE);
+				MetadataMode.EMBED);
 	}
 
 	/**

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

@@ -25,6 +25,7 @@
 import org.junit.jupiter.api.Test;
 
 import org.springframework.ai.document.Document;
+import org.springframework.ai.document.MetadataMode;
 import org.springframework.core.io.DefaultResourceLoader;
 import org.springframework.core.io.Resource;
 
@@ -38,6 +39,9 @@
  */
 public class TokenCountBatchingStrategyTests {
 
+	private static final String SUMMARY_METADATA_VALUE = "This chunk was extracted from the quarterly financial "
+			+ "report and covers revenue, margin and outlook for the reporting period.";
+
 	@Test
 	void batchEmbeddingHappyPath() {
 		TokenCountBatchingStrategy tokenCountBatchingStrategy = new TokenCountBatchingStrategy();
@@ -93,6 +97,39 @@ void batchKeepsEqualDocuments() {
 		assertThat(batches.get(0)).containsExactly(first, second);
 	}
 
+	@Test
+	void batchRejectsDocumentWhoseMetadataPushesItOverTheLimit() {
+		TokenCountBatchingStrategy strategy = new TokenCountBatchingStrategy(EncodingType.CL100K_BASE, 10, 0.0);
+		Document document = new Document("Hello world", Map.of("summary", SUMMARY_METADATA_VALUE));
+
+		assertThatThrownBy(() -> strategy.batch(List.of(document))).isInstanceOf(IllegalArgumentException.class);
+	}
+
+	@Test
+	void batchWithExplicitMetadataModeNoneIgnoresMetadata() {
+		TokenCountBatchingStrategy strategy = new TokenCountBatchingStrategy(EncodingType.CL100K_BASE, 10, 0.0,
+				Document.DEFAULT_CONTENT_FORMATTER, MetadataMode.NONE);
+		Document document = new Document("Hello world", Map.of("summary", SUMMARY_METADATA_VALUE));
+
+		List<List<Document>> batches = strategy.batch(List.of(document));
+
+		assertThat(batches).hasSize(1);
+		assertThat(batches.get(0)).containsExactly(document);
+	}
+
+	@Test
+	void batchSplitsOnTheTokenCountIncludingMetadata() {
+		TokenCountBatchingStrategy strategy = new TokenCountBatchingStrategy(EncodingType.CL100K_BASE, 50, 0.0);
+		List<Document> documents = List.of(new Document("first chunk", Map.of("summary", SUMMARY_METADATA_VALUE)),
+				new Document("second chunk", Map.of("summary", SUMMARY_METADATA_VALUE)),
+				new Document("third chunk", Map.of("summary", SUMMARY_METADATA_VALUE)));
+
+		List<List<Document>> batches = strategy.batch(documents);
+
+		assertThat(batches.stream().mapToInt(List::size).sum()).isEqualTo(3);
+		assertThat(batches).hasSizeGreaterThan(1);
+	}
+
 	@Test
 	void batchKeepsRepeatedDocumentInstance() {
 		TokenCountBatchingStrategy strategy = new TokenCountBatchingStrategy();