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
This commit is contained in:
Laura Trotta
2024-09-10 17:35:24 +02:00
committed by GitHub
parent 9f33e326b1
commit c658c7b52e

View File

@@ -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);