Fix TokenTextSplitter to prevent mini chunks at end

The previous implementation advanced position based on optimized chunk size,
which could create a series of mini chunks when boundary optimization
aggressively shrunk chunks near the end of text.

Changes:
- Implement minimum advance guarantee to prevent consecutive mini chunks
- Add tests to verify no mini chunks are generated at end
- Clean up test code by removing redundant comments

Signed-off-by: logan-mac <[email protected]>

Author: SEUNGHWAN.JUNG

spring-ai-commons/src/main/java/org/springframework/ai/transformer/splitter/TokenTextSplitter.java

@@ -33,11 +33,14 @@
  * @author Raphael Yu
  * @author Christian Tzolov
  * @author Ricken Bazolo
+ * @author Seunghwan Jung
  */
 public class TokenTextSplitter extends TextSplitter {
 
 	private static final int DEFAULT_CHUNK_SIZE = 800;
 
+	private static final int DEFAULT_CHUNK_OVERLAP = 50;
+
 	private static final int MIN_CHUNK_SIZE_CHARS = 350;
 
 	private static final int MIN_CHUNK_LENGTH_TO_EMBED = 5;
@@ -53,6 +56,9 @@ public class TokenTextSplitter extends TextSplitter {
 	// The target size of each text chunk in tokens
 	private final int chunkSize;
 
+	// The overlap size of each text chunk in tokens
+	private final int chunkOverlap;
+
 	// The minimum size of each text chunk in characters
 	private final int minChunkSizeChars;
 
@@ -65,16 +71,20 @@ public class TokenTextSplitter extends TextSplitter {
 	private final boolean keepSeparator;
 
 	public TokenTextSplitter() {
-		this(DEFAULT_CHUNK_SIZE, MIN_CHUNK_SIZE_CHARS, MIN_CHUNK_LENGTH_TO_EMBED, MAX_NUM_CHUNKS, KEEP_SEPARATOR);
+		this(DEFAULT_CHUNK_SIZE, DEFAULT_CHUNK_OVERLAP, MIN_CHUNK_SIZE_CHARS, MIN_CHUNK_LENGTH_TO_EMBED, MAX_NUM_CHUNKS,
+				KEEP_SEPARATOR);
 	}
 
 	public TokenTextSplitter(boolean keepSeparator) {
-		this(DEFAULT_CHUNK_SIZE, MIN_CHUNK_SIZE_CHARS, MIN_CHUNK_LENGTH_TO_EMBED, MAX_NUM_CHUNKS, keepSeparator);
+		this(DEFAULT_CHUNK_SIZE, DEFAULT_CHUNK_OVERLAP, MIN_CHUNK_SIZE_CHARS, MIN_CHUNK_LENGTH_TO_EMBED, MAX_NUM_CHUNKS,
+				keepSeparator);
 	}
 
-	public TokenTextSplitter(int chunkSize, int minChunkSizeChars, int minChunkLengthToEmbed, int maxNumChunks,
-			boolean keepSeparator) {
+	public TokenTextSplitter(int chunkSize, int chunkOverlap, int minChunkSizeChars, int minChunkLengthToEmbed,
+			int maxNumChunks, boolean keepSeparator) {
+		Assert.isTrue(chunkOverlap < chunkSize, "chunk overlap must be less than chunk size");
 		this.chunkSize = chunkSize;
+		this.chunkOverlap = chunkOverlap;
 		this.minChunkSizeChars = minChunkSizeChars;
 		this.minChunkLengthToEmbed = minChunkLengthToEmbed;
 		this.maxNumChunks = maxNumChunks;
@@ -87,57 +97,89 @@ public static Builder builder() {
 
 	@Override
 	protected List<String> splitText(String text) {
-		return doSplit(text, this.chunkSize);
+		return doSplit(text, this.chunkSize, this.chunkOverlap);
 	}
 
-	protected List<String> doSplit(String text, int chunkSize) {
+	protected List<String> doSplit(String text, int chunkSize, int chunkOverlap) {
 		if (text == null || text.trim().isEmpty()) {
 			return new ArrayList<>();
 		}
 
 		List<Integer> tokens = getEncodedTokens(text);
-		List<String> chunks = new ArrayList<>();
-		int num_chunks = 0;
-		while (!tokens.isEmpty() && num_chunks < this.maxNumChunks) {
-			List<Integer> chunk = tokens.subList(0, Math.min(chunkSize, tokens.size()));
-			String chunkText = decodeTokens(chunk);
-
-			// Skip the chunk if it is empty or whitespace
-			if (chunkText.trim().isEmpty()) {
-				tokens = tokens.subList(chunk.size(), tokens.size());
-				continue;
-			}
-
-			// Find the last period or punctuation mark in the chunk
-			int lastPunctuation = Math.max(chunkText.lastIndexOf('.'), Math.max(chunkText.lastIndexOf('?'),
-					Math.max(chunkText.lastIndexOf('!'), chunkText.lastIndexOf('\n'))));
+		// If text is smaller than chunk size, return as a single chunk
+		if (tokens.size() <= chunkSize) {
+			String processedText = this.keepSeparator ? text.trim() : text.replace(System.lineSeparator(), " ").trim();
 
-			if (lastPunctuation != -1 && lastPunctuation > this.minChunkSizeChars) {
-				// Truncate the chunk text at the punctuation mark
-				chunkText = chunkText.substring(0, lastPunctuation + 1);
+			if (processedText.length() > this.minChunkLengthToEmbed) {
+				return List.of(processedText);
 			}
+			return new ArrayList<>();
+		}
+		List<String> chunks = new ArrayList<>();
 
-			String chunkTextToAppend = (this.keepSeparator) ? chunkText.trim()
-					: chunkText.replace(System.lineSeparator(), " ").trim();
-			if (chunkTextToAppend.length() > this.minChunkLengthToEmbed) {
-				chunks.add(chunkTextToAppend);
+		int position = 0;
+		int num_chunks = 0;
+		while (position < tokens.size() && num_chunks < this.maxNumChunks) {
+			int chunkEnd = Math.min(position + chunkSize, tokens.size());
+
+			// Extract tokens for this chunk
+			List<Integer> chunkTokens = tokens.subList(position, chunkEnd);
+			String chunkText = decodeTokens(chunkTokens);
+
+			// Apply sentence boundary optimization
+			String optimizedText = optimizeChunkBoundary(chunkText);
+			int optimizedTokenCount = getEncodedTokens(optimizedText).size();
+
+			// Use optimized chunk
+			String finalChunkText = optimizedText;
+			int finalChunkTokenCount = optimizedTokenCount;
+
+			// Advance position with minimum advance guarantee
+			// This prevents creating a series of mini chunks when boundary optimization
+			// aggressively shrinks chunks
+			int naturalAdvance = finalChunkTokenCount - chunkOverlap;
+			int minAdvance = Math.max(1, (chunkSize - chunkOverlap) / 2);
+			int advance = Math.max(naturalAdvance, minAdvance);
+			position += advance;
+
+			// Format according to keepSeparator setting
+			String formattedChunk = this.keepSeparator ? finalChunkText.trim()
+					: finalChunkText.replace(System.lineSeparator(), " ").trim();
+
+			// Add chunk if it meets minimum length
+			if (formattedChunk.length() > this.minChunkLengthToEmbed) {
+				chunks.add(formattedChunk);
+				num_chunks++;
 			}
+		}
 
-			// Remove the tokens corresponding to the chunk text from the remaining tokens
-			tokens = tokens.subList(getEncodedTokens(chunkText).size(), tokens.size());
+		return chunks;
+	}
 
-			num_chunks++;
+	private String optimizeChunkBoundary(String chunkText) {
+		if (chunkText.length() <= this.minChunkSizeChars) {
+			return chunkText;
 		}
 
-		// Handle the remaining tokens
-		if (!tokens.isEmpty()) {
-			String remaining_text = decodeTokens(tokens).replace(System.lineSeparator(), " ").trim();
-			if (remaining_text.length() > this.minChunkLengthToEmbed) {
-				chunks.add(remaining_text);
+		// Look for sentence endings: . ! ? \n
+		int bestCutPoint = -1;
+
+		// Check in reverse order to find the last sentence ending
+		for (int i = chunkText.length() - 1; i >= this.minChunkSizeChars; i--) {
+			char c = chunkText.charAt(i);
+			if (c == '.' || c == '!' || c == '?' || c == '\n') {
+				bestCutPoint = i + 1; // Include the punctuation
+				break;
 			}
 		}
 
-		return chunks;
+		// If we found a good cut point, use it
+		if (bestCutPoint > 0) {
+			return chunkText.substring(0, bestCutPoint);
+		}
+
+		// Otherwise return the original chunk
+		return chunkText;
 	}
 
 	private List<Integer> getEncodedTokens(String text) {
@@ -156,6 +198,8 @@ public static final class Builder {
 
 		private int chunkSize = DEFAULT_CHUNK_SIZE;
 
+		private int chunkOverlap = DEFAULT_CHUNK_OVERLAP;
+
 		private int minChunkSizeChars = MIN_CHUNK_SIZE_CHARS;
 
 		private int minChunkLengthToEmbed = MIN_CHUNK_LENGTH_TO_EMBED;
@@ -172,6 +216,11 @@ public Builder withChunkSize(int chunkSize) {
 			return this;
 		}
 
+		public Builder withChunkOverlap(int chunkOverlap) {
+			this.chunkOverlap = chunkOverlap;
+			return this;
+		}
+
 		public Builder withMinChunkSizeChars(int minChunkSizeChars) {
 			this.minChunkSizeChars = minChunkSizeChars;
 			return this;
@@ -193,8 +242,8 @@ public Builder withKeepSeparator(boolean keepSeparator) {
 		}
 
 		public TokenTextSplitter build() {
-			return new TokenTextSplitter(this.chunkSize, this.minChunkSizeChars, this.minChunkLengthToEmbed,
-					this.maxNumChunks, this.keepSeparator);
+			return new TokenTextSplitter(this.chunkSize, this.chunkOverlap, this.minChunkSizeChars,
+					this.minChunkLengthToEmbed, this.maxNumChunks, this.keepSeparator);
 		}
 
 	}