From c658c7b52ebd46052831be61ab6f584ee8fd0131 Mon Sep 17 00:00:00 2001
From: Laura Trotta <153528055+l-trotta@users.noreply.github.com>
Date: Tue, 10 Sep 2024 17:35:24 +0200
Subject: [PATCH] GH-1341: Fixing bug in TokenCountBatchingStrategy

Fixes: https://github.com/spring-projects/spring-ai/issues/1341

* Instead of clearing the current batch, create a new instance each time it needs to be reset
* added author
---
 .../ai/embedding/TokenCountBatchingStrategy.java               | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/spring-ai-core/src/main/java/org/springframework/ai/embedding/TokenCountBatchingStrategy.java b/spring-ai-core/src/main/java/org/springframework/ai/embedding/TokenCountBatchingStrategy.java
index 2a9aefa23..d790a9de4 100644
--- a/spring-ai-core/src/main/java/org/springframework/ai/embedding/TokenCountBatchingStrategy.java
+++ b/spring-ai-core/src/main/java/org/springframework/ai/embedding/TokenCountBatchingStrategy.java
@@ -46,6 +46,7 @@ import com.knuddels.jtokkit.api.EncodingType;
  *
  * @author Soby Chacko
  * @author Mark Pollack
+ * @author Laura Trotta
  * @since 1.0.0
  */
 public class TokenCountBatchingStrategy implements BatchingStrategy {
@@ -122,7 +123,7 @@ public class TokenCountBatchingStrategy implements BatchingStrategy {
 			Integer tokenCount = documentTokens.get(document);
 			if (currentSize + tokenCount > maxInputTokenCount) {
 				batches.add(currentBatch);
-				currentBatch.clear();
+				currentBatch = new ArrayList<>();
 				currentSize = 0;
 			}
 			currentBatch.add(document);