From ab907926a12313ab61480162884df9beef5adbc4 Mon Sep 17 00:00:00 2001 From: Andrew Chen <48723787+chuenchen309@users.noreply.github.com> Date: Wed, 15 Jul 2026 13:47:26 +0800 Subject: [PATCH] Fix merge_chunks pre-allocating a fixed chunk count that ignores overlap, letting the last chunk grow unbounded merge_chunks() pre-computes num_chunks = ceil(total_tokens / target_size), then guards the rollover with `curr_chunk < num_chunks - 1`. That estimate never accounts for the `overlap` tokens re-injected at every chunk boundary -- with the repo's own defaults (chunk_token_threshold=2048, overlap_rate=0.1 -> overlap=205), each boundary adds ~205 tokens beyond the naive total_tokens/target_size budget. Once curr_chunk reaches the last pre-allocated slot, the rollover condition can never fire again (the second half of the `and` is always False), so every remaining token gets dumped into that one last chunk with no cap. Verified independently with real (non-mocked) calls to merge_chunks: at 10,000 tokens the last chunk overshoots target_size by 28%; at 40,000 tokens by 143%; at 160,000 tokens by 693% -- the overshoot is unbounded and grows with document length, not a fixed off-by-one. merge_chunks feeds both LLMContentFilter (content_filter_strategy.py:913) and LLMExtractionStrategy (extraction_strategy.py:778), so a long enough document silently produces a "chunk" far larger than the token budget the caller configured -- risking truncation or blown context windows in the downstream LLM call, with no error or warning. Fix: stop pre-allocating a fixed chunk count. Grow the chunk list on demand (`chunks.append(...)` whenever curr_size reaches target_size, unconditionally) instead of capping rollover at a chunk count computed without knowledge of the overlap that will be injected. Added tests/test_merge_chunks.py: no-overlap stays within target_size, overlap stays within target_size at both 10k and 160k tokens (regression guard for the unbounded-growth behavior), overlap tokens are correctly repeated at chunk boundaries, full token content is preserved without overlap, and empty input returns an empty list. Co-Authored-By: Claude Sonnet 5 --- crawl4ai/utils.py | 25 ++++++++++++------------ tests/test_merge_chunks.py | 40 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 13 deletions(-) create mode 100644 tests/test_merge_chunks.py diff --git a/crawl4ai/utils.py b/crawl4ai/utils.py index 279c27708..8e3556319 100644 --- a/crawl4ai/utils.py +++ b/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 diff --git a/tests/test_merge_chunks.py b/tests/test_merge_chunks.py new file mode 100644 index 000000000..847ea53b4 --- /dev/null +++ b/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), [])