fix: preserve whitespace during split in _split_text_fallback

src/chonkie/chunker/slumber.py

@@ -1,5 +1,6 @@
 """Module containing the SlumberChunker."""
 
+import re
 from bisect import bisect_left
 from itertools import accumulate
 from typing import Optional, Union
@@ -103,20 +104,15 @@ class Split(BaseModel):  # type: ignore
 
     def _split_text(self, text: str, recursive_level: RecursiveLevel) -> list[str]:
         """Split the text into chunks using the delimiters."""
+        # Use Python fallback for whitespace mode because it correctly handles
+        # all three include_delim modes ("prev", "next", None).
+        # The Cython extension ignores include_delim for whitespace splitting.
+        if recursive_level.whitespace:
+            return self._split_text_fallback(text, recursive_level)
+
         if _CYTHON_AVAILABLE:
-            # Use the optimized Cython split function
-            if recursive_level.whitespace:
-                return list(
-                    split_text(
-                        text,
-                        delim=None,
-                        include_delim=recursive_level.include_delim,
-                        min_characters_per_segment=self.min_characters_per_chunk,
-                        whitespace_mode=True,
-                        character_fallback=False,
-                    ),
-                )
-            elif recursive_level.delimiters:
+            # Use the optimized Cython split function for non-whitespace modes
+            if recursive_level.delimiters:
                 return list(
                     split_text(
                         text,
@@ -148,15 +144,15 @@ def _split_text_fallback(self, text: str, recursive_level: RecursiveLevel) -> li
             # Add whitespace back; assumes that the whitespace is uniform across the text
             # if the whitespace is not uniform, the split will not be reconstructable.
             if recursive_level.include_delim == "prev":
-                splits = [" " + split for (i, split) in enumerate(candidate_splits) if i > 0]
+                # Include first word without space, prepend space to subsequent words
+                splits = [candidate_splits[0]] + [" " + s for s in candidate_splits[1:]]
             elif recursive_level.include_delim == "next":
-                splits = [
-                    split + " "
-                    for (i, split) in enumerate(candidate_splits)
-                    if i < len(candidate_splits) - 1
-                ]
+                # Append space to all words except the last one
+                splits = [s + " " for s in candidate_splits[:-1]] + [candidate_splits[-1]]
             else:
-                splits = candidate_splits
+                # Preserve spaces by keeping them as separate elements
+                splits = re.split(r"( )", text)
+                splits = [s for s in splits if s]
 
         elif recursive_level.delimiters:
             if recursive_level.include_delim == "prev":

tests/chunkers/test_slumber_chunker.py

@@ -174,6 +174,139 @@ def test_split_text_whitespace(self, mock_genie: MockGenie) -> None:
         assert isinstance(splits, list)
         assert len(splits) > 0
 
+    def test_split_text_whitespace_reconstruction(self, mock_genie: MockGenie) -> None:
+        """Test that whitespace splits can be reconstructed to original text."""
+        chunker = SlumberChunker(genie=mock_genie, min_characters_per_chunk=1)
+        text = "hello world test string"
+
+        # Test with include_delim="prev"
+        level = RecursiveLevel(whitespace=True, include_delim="prev")
+        splits = chunker._split_text(text, level)
+        reconstructed = "".join(splits)
+        assert reconstructed == text, f"prev: '{reconstructed}' != '{text}'"
+
+        # Test with include_delim="next"
+        level = RecursiveLevel(whitespace=True, include_delim="next")
+        splits = chunker._split_text(text, level)
+        reconstructed = "".join(splits)
+        assert reconstructed == text, f"next: '{reconstructed}' != '{text}'"
+
+        # Test with no include_delim
+        level = RecursiveLevel(whitespace=True, include_delim=None)
+        splits = chunker._split_text(text, level)
+        reconstructed = "".join(splits)
+        assert reconstructed == text, f"none: '{reconstructed}' != '{text}'"
+
+    def test_split_text_whitespace_index_alignment(self, mock_genie: MockGenie) -> None:
+        """Test that sum of split lengths equals original text length."""
+        chunker = SlumberChunker(genie=mock_genie, min_characters_per_chunk=1)
+        text = "The quick brown fox"
+
+        for include_delim in ["prev", "next", None]:
+            level = RecursiveLevel(whitespace=True, include_delim=include_delim)
+            splits = chunker._split_text(text, level)
+            total_length = sum(len(s) for s in splits)
+            assert total_length == len(text), (
+                f"include_delim={include_delim}: total_length={total_length} != {len(text)}"
+            )
+
+    def test_split_text_whitespace_multiple_spaces(self, mock_genie: MockGenie) -> None:
+        """Test whitespace splitting with multiple consecutive spaces."""
+        chunker = SlumberChunker(genie=mock_genie, min_characters_per_chunk=1)
+        text = "hello  world"  # Two consecutive spaces
+
+        for include_delim in ["prev", "next", None]:
+            level = RecursiveLevel(whitespace=True, include_delim=include_delim)
+            splits = chunker._split_text(text, level)
+            reconstructed = "".join(splits)
+            assert reconstructed == text, (
+                f"include_delim={include_delim}: '{reconstructed}' != '{text}'"
+            )
+
+    def test_split_text_whitespace_leading_trailing(self, mock_genie: MockGenie) -> None:
+        """Test whitespace splitting with leading and trailing spaces."""
+        chunker = SlumberChunker(genie=mock_genie, min_characters_per_chunk=1)
+        text = " hello world "  # Leading and trailing spaces
+
+        for include_delim in ["prev", "next", None]:
+            level = RecursiveLevel(whitespace=True, include_delim=include_delim)
+            splits = chunker._split_text(text, level)
+            reconstructed = "".join(splits)
+            assert reconstructed == text, (
+                f"include_delim={include_delim}: '{reconstructed}' != '{text}'"
+            )
+
+    def test_split_text_whitespace_empty_string(self, mock_genie: MockGenie) -> None:
+        """Test whitespace splitting with empty string."""
+        chunker = SlumberChunker(genie=mock_genie, min_characters_per_chunk=1)
+        for include_delim in ["prev", "next", None]:
+            level = RecursiveLevel(whitespace=True, include_delim=include_delim)
+            splits = chunker._split_text("", level)
+            assert "".join(splits) == ""
+
+    def test_split_text_whitespace_only_spaces(self, mock_genie: MockGenie) -> None:
+        """Test whitespace splitting with only spaces."""
+        chunker = SlumberChunker(genie=mock_genie, min_characters_per_chunk=1)
+        for include_delim in ["prev", "next", None]:
+            level = RecursiveLevel(whitespace=True, include_delim=include_delim)
+            splits = chunker._split_text("   ", level)
+            assert "".join(splits) == "   "
+
+    def test_split_text_whitespace_single_word(self, mock_genie: MockGenie) -> None:
+        """Test whitespace splitting with single word (no spaces)."""
+        chunker = SlumberChunker(genie=mock_genie, min_characters_per_chunk=1)
+        for include_delim in ["prev", "next", None]:
+            level = RecursiveLevel(whitespace=True, include_delim=include_delim)
+            splits = chunker._split_text("hello", level)
+            assert "".join(splits) == "hello"
+
+    def test_whitespace_regression_no_concatenation(self, mock_genie: MockGenie) -> None:
+        """Regression test for bug where whitespace splitting caused concatenation.
+
+        Original bug: "The quick brown fox" -> "Thequickbrownfox"
+        Root cause: Whitespace lost during split/merge causing index drift
+        """
+        chunker = SlumberChunker(genie=mock_genie, min_characters_per_chunk=1)
+        text = "The quick brown fox"
+
+        level = RecursiveLevel(whitespace=True, include_delim="prev")
+        splits = chunker._split_text(text, level)
+        reconstructed = "".join(splits)
+
+        assert "Thequickbrown" not in reconstructed  # Verify no concatenation
+        assert reconstructed == text
+
+    def test_whitespace_mode_specific_structure(self, mock_genie: MockGenie) -> None:
+        """Test that each include_delim mode produces correct structure."""
+        chunker = SlumberChunker(genie=mock_genie, min_characters_per_chunk=1)
+        text = "hello world"
+
+        # "prev" mode: space goes with following word
+        level = RecursiveLevel(whitespace=True, include_delim="prev")
+        splits = chunker._split_text(text, level)
+        assert splits == ["hello", " world"]
+
+        # "next" mode: space goes with preceding word
+        level = RecursiveLevel(whitespace=True, include_delim="next")
+        splits = chunker._split_text(text, level)
+        assert splits == ["hello ", "world"]
+
+        # None mode: space is separate element
+        level = RecursiveLevel(whitespace=True, include_delim=None)
+        splits = chunker._split_text(text, level)
+        assert splits == ["hello", " ", "world"]
+
+    def test_whitespace_with_merging_enabled(self, mock_genie: MockGenie) -> None:
+        """Test whitespace preservation with default merging logic."""
+        chunker = SlumberChunker(genie=mock_genie)  # Default min_characters_per_chunk=24
+        text = "a b c d e f g h i j k l m n o p q r s t u v w x y z"
+
+        level = RecursiveLevel(whitespace=True, include_delim="prev")
+        splits = chunker._split_text(text, level)
+
+        # Verify reconstruction even after merging
+        assert "".join(splits) == text
+
     def test_split_text_delimiters(self, mock_genie: MockGenie) -> None:
         """Test text splitting with custom delimiters."""
         chunker = SlumberChunker(genie=mock_genie)