Fix overlapped chunks exceeding chunk_size when separators count as tokens

src/semchunk/semchunk.py

@@ -487,13 +487,26 @@ def chunk(
                 unoverlapped_chunk_size / subchunk_size
             )  # NOTE `math.ceil` would cause overlaps to be missed.
 
-            offsets = [
-                (
-                    suboffsets[(start := i * subchunk_stride)][0],
-                    suboffsets[min(start + subchunks_per_chunk, num_subchunks) - 1][1],
-                )
-                for i in range(max(1, math.ceil((num_subchunks - subchunks_per_chunk) / subchunk_stride) + 1))
-            ]
+            offsets = []
+            i = 0
+
+            while True:
+                # Grow the window up to `subchunks_per_chunk` subchunks, then shrink it if the
+                # whitespace reintroduced between subchunks (stripped when the subchunks were
+                # formed but re-spanned here) would push the merged chunk over `chunk_size`.
+                end_i = min(i + subchunks_per_chunk, num_subchunks) - 1
+
+                while end_i > i and token_counter(text[suboffsets[i][0] : suboffsets[end_i][1]]) > chunk_size:
+                    end_i -= 1
+
+                offsets.append((suboffsets[i][0], suboffsets[end_i][1]))
+
+                if end_i >= num_subchunks - 1:
+                    break
+
+                # Advance by the stride but never past the last included subchunk so no
+                # subchunk is ever skipped (which would drop content).
+                i = min(i + subchunk_stride, end_i + 1)
 
         # Materialize chunks from offsets.
         chunks = [text[start:end] for start, end in offsets]
@@ -611,13 +624,26 @@ def chunk(
                 unoverlapped_chunk_size / subchunk_size
             )  # NOTE `math.ceil` would cause overlaps to be missed.
 
-            offsets = [
-                (
-                    suboffsets[(start := i * subchunk_stride)][0],
-                    suboffsets[min(start + subchunks_per_chunk, num_subchunks) - 1][1],
-                )
-                for i in range(max(1, math.ceil((num_subchunks - subchunks_per_chunk) / subchunk_stride) + 1))
-            ]
+            offsets = []
+            i = 0
+
+            while True:
+                # Grow the window up to `subchunks_per_chunk` subchunks, then shrink it if the
+                # whitespace reintroduced between subchunks (stripped when the subchunks were
+                # formed but re-spanned here) would push the merged chunk over `chunk_size`.
+                end_i = min(i + subchunks_per_chunk, num_subchunks) - 1
+
+                while end_i > i and token_counter(text[suboffsets[i][0] : suboffsets[end_i][1]]) > chunk_size:
+                    end_i -= 1
+
+                offsets.append((suboffsets[i][0], suboffsets[end_i][1]))
+
+                if end_i >= num_subchunks - 1:
+                    break
+
+                # Advance by the stride but never past the last included subchunk so no
+                # subchunk is ever skipped (which would drop content).
+                i = min(i + subchunk_stride, end_i + 1)
 
             chunks = [text[start:end] for start, end in offsets]
 

tests/test_semchunk.py

@@ -200,4 +200,19 @@ def test_semchunk() -> None:
     semchunk.chunk('\n\n', 512, lambda *args: 0)
 
 if __name__ == '__main__':
-    test_semchunk()
+    test_semchunk()
+
+def test_overlap_never_exceeds_chunk_size() -> None:
+    """Overlapped chunks must not exceed ``chunk_size`` when separators count as tokens.
+
+    Merging subchunks by offset re-includes the whitespace stripped between them; with a
+    whitespace-significant counter (``len``) that can push a chunk over ``chunk_size``.
+    """
+    text, chunk_size = "aa  bb  cc  dd", 4
+    chunks, offsets = semchunk.chunk(text, chunk_size, len, offsets=True, overlap=0.5, memoize=False)
+
+    # Every chunk stays within the documented maximum...
+    for chunk in chunks:
+        assert len(chunk) <= chunk_size, f"{chunk!r} has {len(chunk)} tokens > chunk_size {chunk_size}"
+    # ...and the offsets still reconstruct the chunks (no content dropped).
+    assert chunks == [text[start:end] for start, end in offsets]