Repository: griptape-ai/griptape

Issue:
Chunker aggressively chunking text when using `" "` as a separator.
When using `" "` as a `ChunkSeparator`, the `TextChunker` appears to be overly aggressive with chunking the text. In this example it's splitting the first "chunk" into two, even though their total chunks would not be more than `max_tokens`.

```python
from griptape.chunkers import ChunkSeparator, TextChunker
from griptape.tokenizers import OpenAiTokenizer
from rich import print

text = "This is a paragraph of text. I'll count to three: one, two three."
tokenizer = OpenAiTokenizer(model="gpt-4o-mini")
max_tokens = 10
chunker = TextChunker(
    max_tokens=max_tokens, separators=[ChunkSeparator(" ")], tokenizer=tokenizer
).chunk(text)

print(f"Max Tokens: {max_tokens}")
for i, chunk in enumerate(chunker):
    chunk_size = tokenizer.count_tokens(chunk.value)
    print(f"{chunk_size}: {chunk.value}")
```

and the result:

```bash
Max Tokens: 10
4: This is a paragraph
3: of text.
10: I'll count to three: one, two three.
```

The first two chunks should be a single chunk.

Discussion:

Solution patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index acdde9dd..832708a1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -18,6 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 
 - Occasional crash during `FuturesExecutorMixin` cleanup.
+- `BaseChunker` chunking more than necessary.
 
 ### Deprecated
 
diff --git a/griptape/chunkers/base_chunker.py b/griptape/chunkers/base_chunker.py
index 9b6ef64b..c93ef5d2 100644
--- a/griptape/chunkers/base_chunker.py
+++ b/griptape/chunkers/base_chunker.py
@@ -40,15 +40,11 @@ class BaseChunker(ABC):
 
     def _chunk_recursively(self, chunk: str, current_separator: Optional[ChunkSeparator] = None) -> list[str]:
         token_count = self.tokenizer.count_tokens(chunk)
+        half_token_count = token_count // 2
 
         if token_count <= self.max_tokens:
             return [chunk]
         else:
-            balance_index = -1
-            balance_diff = float("inf")
-            tokens_count = 0
-            half_token_count = token_count // 2
-
             # If a separator is provided, only use separators after it.
             separators = (
                 self.separators[self.separators.index(current_separator) :] if current_separator else self.separators
@@ -59,22 +55,13 @@ class BaseChunker(ABC):
                 # Split the chunk into subchunks using the current separator.
                 subchunks = list(filter(None, chunk.split(separator.value)))
 
-                # Check if the split resulted in more than one subchunk.
+                # We can only recurse if there are multiple subchunks.
                 if len(subchunks) > 1:
-                    # Iterate through the subchunks and calculate token counts.
-                    for index, subchunk in enumerate(subchunks):
-                        if index < len(subchunks):
-                            subchunk = separator.value + subchunk if separator.is_prefix else subchunk + separator.value
-
-                        tokens_count += self.tokenizer.count_tokens(subchunk)
-
-                        # Update the best split if the current one is more balanced.
-                        if abs(tokens_count - half_token_count) < balance_diff:
-                            balance_index = index
-                            balance_diff = abs(tokens_count - half_token_count)
+                    # Find what combination of subchunks results in the most balanced split of the chunk.
+                    midpoint_index = self.__find_midpoint_index(subchunks, half_token_count)
 
                     # Create the two subchunks based on the best separator.
-                    first_subchunk, second_subchunk = self.__get_subchunks(separator, subchunks, balance_index)
+                    first_subchunk, second_subchunk = self.__get_subchunks(separator, subchunks, midpoint_index)
 
                     # Continue recursively chunking the subchunks.
                     first_subchunk_rec = self._chunk_recursively(first_subchunk.strip(), separator)
@@ -106,3 +93,17 @@ class BaseChunker(ABC):
             second_subchunk = separator.value.join(subchunks[balance_index + 1 :])
 
         return first_subchunk, second_subchunk
+
+    def __find_midpoint_index(self, subchunks: list[str], half_token_count: int) -> int:
+        midpoint_index = -1
+        best_midpoint_distance = float("inf")
+
+        for index, _ in enumerate(subchunks):
+            subchunk_tokens_count = self.tokenizer.count_tokens("".join(subchunks[: index + 1]))
+
+            midpoint_distance = abs(subchunk_tokens_count - half_token_count)
+            if midpoint_distance < best_midpoint_distance:
+                midpoint_index = index
+                best_midpoint_distance = midpoint_distance
+
+        return midpoint_index


Test patch:
diff --git a/tests/unit/chunkers/test_markdown_chunker.py b/tests/unit/chunkers/test_markdown_chunker.py
index 30db6461..436e09b1 100644
--- a/tests/unit/chunkers/test_markdown_chunker.py
+++ b/tests/unit/chunkers/test_markdown_chunker.py
@@ -33,12 +33,12 @@ class TestTextChunker:
         assert chunks[1].value.startswith("## Header 2\nfoo-0")
         assert chunks[2].value.startswith("foo-0.")
         assert chunks[3].value.startswith("## Header 3\nfoo-0")
-        assert chunks[4].value.startswith("foo-9.")
-        assert chunks[5].value.startswith("foo-15.")
+        assert chunks[4].value.startswith("foo-10.")
+        assert chunks[5].value.startswith("foo-16.")
 
         assert chunks[0].value.endswith(". foo-5.")
         assert chunks[1].value.endswith(". foo-5.")
         assert chunks[2].value.endswith(". foo-5.")
-        assert chunks[3].value.endswith(". foo-8.")
-        assert chunks[4].value.endswith(". foo-14.")
+        assert chunks[3].value.endswith(". foo-9.")
+        assert chunks[4].value.endswith(". foo-15.")
         assert chunks[5].value.endswith(". foo-24.")
diff --git a/tests/unit/chunkers/test_text_chunker.py b/tests/unit/chunkers/test_text_chunker.py
index d9d01f16..99f83d94 100644
--- a/tests/unit/chunkers/test_text_chunker.py
+++ b/tests/unit/chunkers/test_text_chunker.py
@@ -2,6 +2,7 @@ import pytest
 
 from griptape.artifacts import TextArtifact
 from griptape.chunkers import TextChunker
+from griptape.chunkers.chunk_separator import ChunkSeparator
 from tests.unit.chunkers.utils import gen_paragraph
 
 MAX_TOKENS = 50
@@ -55,12 +56,12 @@ class TestTextChunker:
             assert chunker.tokenizer.count_tokens(chunk.value) <= MAX_TOKENS
 
         assert chunks[0].value.startswith("foo-0!")
-        assert chunks[1].value.startswith("foo-10!")
-        assert chunks[2].value.startswith("foo-16!")
+        assert chunks[1].value.startswith("foo-11!")
+        assert chunks[2].value.startswith("foo-17!")
         assert chunks[3].value.startswith("foo-0.")
 
-        assert chunks[0].value.endswith("! foo-9!")
-        assert chunks[1].value.endswith("! foo-15!")
+        assert chunks[0].value.endswith("! foo-10!")
+        assert chunks[1].value.endswith("! foo-16!")
         assert chunks[2].value.endswith("! foo-24!")
         assert chunks[3].value.endswith(". foo-11.")
 
@@ -91,23 +92,34 @@ class TestTextChunker:
             assert chunker.tokenizer.count_tokens(chunk.value) <= MAX_TOKENS
 
         assert chunks[0].value.startswith("foo-0!")
-        assert chunks[1].value.startswith("foo-10!")
-        assert chunks[2].value.startswith("foo-16!")
+        assert chunks[1].value.startswith("foo-11!")
+        assert chunks[2].value.startswith("foo-17!")
         assert chunks[3].value.startswith("foo-0.")
         assert chunks[4].value.startswith("foo-0?")
-        assert chunks[5].value.startswith("foo-5?")
+        assert chunks[5].value.startswith("foo-9?")
         assert chunks[6].value.startswith("foo-0")
-        assert chunks[7].value.startswith("foo-6")
+        assert chunks[7].value.startswith("foo-8")
 
-        assert chunks[0].value.endswith("! foo-9!")
-        assert chunks[1].value.endswith("! foo-15!")
+        assert chunks[0].value.endswith("! foo-10!")
+        assert chunks[1].value.endswith("! foo-16!")
         assert chunks[2].value.endswith("! foo-24!")
         assert chunks[3].value.endswith(". foo-11.")
-        assert chunks[4].value.endswith("? foo-4?")
+        assert chunks[4].value.endswith("? foo-8?")
         assert chunks[5].value.endswith("? foo-12?")
-        assert chunks[6].value.endswith(" foo-5")
+        assert chunks[6].value.endswith(" foo-7")
         assert chunks[7].value.endswith(" foo-16")
 
     def test_chunk_with_max_tokens(self, chunker):
         with pytest.raises(ValueError):
             TextChunker(max_tokens=-1)
+
+    def test_tiny_max_tokens(self, chunker):
+        max_tokens = 10
+        chunker.separators = [ChunkSeparator(" ")]
+        chunker.max_tokens = max_tokens
+        text = "This is a paragraph of text. I'll count to three: one, two three."
+        chunks = chunker.chunk(text)
+
+        assert len(chunks) == 2
+        for chunk in chunks:
+            assert chunker.tokenizer.count_tokens(chunk.value) <= max_tokens
diff --git a/tests/unit/engines/extraction/test_json_extraction_engine.py b/tests/unit/engines/extraction/test_json_extraction_engine.py
index b4aa58b7..e8a98589 100644
--- a/tests/unit/engines/extraction/test_json_extraction_engine.py
+++ b/tests/unit/engines/extraction/test_json_extraction_engine.py
@@ -29,7 +29,7 @@ class TestJsonExtractionEngine:
         large_text = Path(normpath(join(dirname(__file__), "../../../resources", "test.txt"))).read_text()
 
         extracted = engine.extract_text(large_text * 50)
-        assert len(extracted) == 354
+        assert len(extracted) == 322
         assert extracted[0].value == {"test_key_1": "test_value_1"}
 
     def test_extract_error(self, engine):