GH-2123: Add chunkOverlap support to TokenTextSplitter

Fixes GH-2123 (#2123)

  - Add chunkOverlap field and configuration to TokenTextSplitter class
  - Implement overlap functionality in doSplit method with boundary
  optimization
  - Add optimizeChunkBoundary method for sentence-aware chunk splitting
  - Add validation to ensure chunkOverlap < chunkSize
  - Update Builder pattern with withChunkOverlap method
  - Add comprehensive test coverage for overlap functionality
  - Improve existing tests to handle dynamic chunk counts

  Signed-off-by: Seunghwan Jung <[email protected]>

Author: ralla0405

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;
@@ -46,13 +49,17 @@ public class TokenTextSplitter extends TextSplitter {
 
 	private static final boolean KEEP_SEPARATOR = true;
 
+
 	private final EncodingRegistry registry = Encodings.newLazyEncodingRegistry();
 
 	private final Encoding encoding = this.registry.getEncoding(EncodingType.CL100K_BASE);
 
 	// 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 +72,18 @@ 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 +96,80 @@ 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;
-			}
+		// 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();
 
-			// 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 (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 finalChunkText = optimizeChunkBoundary(chunkText);
+			int finalChunkTokenCount = getEncodedTokens(finalChunkText).size();
+			int advance = Math.max(1, finalChunkTokenCount - chunkOverlap);
+			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 +188,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 +206,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,7 +232,7 @@ public Builder withKeepSeparator(boolean keepSeparator) {
 		}
 
 		public TokenTextSplitter build() {
-			return new TokenTextSplitter(this.chunkSize, this.minChunkSizeChars, this.minChunkLengthToEmbed,
+			return new TokenTextSplitter(this.chunkSize, this.chunkOverlap, this.minChunkSizeChars, this.minChunkLengthToEmbed,
 					this.maxNumChunks, this.keepSeparator);
 		}