Fix merge_chunks pre-allocating a fixed chunk count that ignores overlap

crawl4ai/utils.py

@@ -198,26 +198,25 @@ def merge_chunks(
     if not total_tokens:
         return []
 
-    # Pre-allocate chunks
-    num_chunks = max(1, (total_tokens + target_size - 1) // target_size)
-    chunks: List[List[str]] = [[] for _ in range(num_chunks)]
-
-    curr_chunk = 0
+    # Chunks grow on demand instead of being pre-allocated from total_tokens // target_size:
+    # that estimate doesn't account for the tokens re-injected at every overlap boundary, so a
+    # fixed chunk count silently runs out and the last chunk absorbs all remaining tokens
+    # uncapped (growing without bound as the document gets longer).
+    chunks: List[List[str]] = [[]]
     curr_size = 0
-    
+
     # Distribute tokens
     for tokens in chain.from_iterable(all_tokens):
-        if curr_size >= target_size and curr_chunk < num_chunks - 1:
+        if curr_size >= target_size:
             if overlap > 0:
-                overlap_tokens = chunks[curr_chunk][-overlap:]
-                curr_chunk += 1
-                chunks[curr_chunk].extend(overlap_tokens)
+                overlap_tokens = chunks[-1][-overlap:]
+                chunks.append(list(overlap_tokens))
                 curr_size = len(overlap_tokens)
             else:
-                curr_chunk += 1
+                chunks.append([])
                 curr_size = 0
-                
-        chunks[curr_chunk].append(tokens)
+
+        chunks[-1].append(tokens)
         curr_size += 1
 
     # Return only non-empty chunks

tests/test_merge_chunks.py

@@ -0,0 +1,40 @@
+import unittest
+from crawl4ai.utils import merge_chunks
+
+class TestMergeChunks(unittest.TestCase):
+
+    def test_no_overlap_respects_target_size(self):
+        docs = [f"word{i} " * 50 for i in range(200)]  # 10000 tokens total
+        chunks = merge_chunks(docs, target_size=2048, overlap=0)
+        for chunk in chunks:
+            self.assertLessEqual(len(chunk.split()), 2048)
+
+    def test_overlap_does_not_overshoot_target_size(self):
+        # chunk_token_threshold=2048 / overlap_rate=0.1 are this repo's own defaults.
+        docs = [f"word{i} " * 50 for i in range(200)]  # 10000 tokens total
+        chunks = merge_chunks(docs, target_size=2048, overlap=205)
+        for chunk in chunks:
+            self.assertLessEqual(len(chunk.split()), 2048)
+
+    def test_overlap_does_not_overshoot_on_longer_documents(self):
+        # Regression guard: the pre-allocated-chunk-count bug made the last chunk grow
+        # without bound as more documents were merged in.
+        docs = [f"word{i} " * 50 for i in range(3200)]  # 160000 tokens total
+        chunks = merge_chunks(docs, target_size=2048, overlap=205)
+        for chunk in chunks:
+            self.assertLessEqual(len(chunk.split()), 2048)
+
+    def test_overlap_tokens_are_repeated_at_chunk_boundary(self):
+        docs = [f"word{i}" for i in range(10)]
+        chunks = merge_chunks(docs, target_size=6, overlap=2)
+        self.assertEqual(chunks[0].split()[-2:], chunks[1].split()[:2])
+
+    def test_all_tokens_are_preserved_without_overlap(self):
+        docs = [f"word{i}" for i in range(10)]
+        chunks = merge_chunks(docs, target_size=6, overlap=0)
+        merged = " ".join(chunks).split()
+        self.assertEqual(merged, [f"word{i}" for i in range(10)])
+
+    def test_empty_docs_returns_empty_list(self):
+        self.assertEqual(merge_chunks([], target_size=100), [])
+        self.assertEqual(merge_chunks(["", "  "], target_size=100), [])