feat: support split_by="token" in HierarchicalDocumentSplitter

docs-website/docs/pipeline-components/preprocessors/hierarchicaldocumentsplitter.mdx

@@ -40,8 +40,9 @@ This descending order of sizes builds the hierarchy.
 
 These additional parameters can be set when the component is initialized:
 
-- `split_by` can be `"word"` (default), `"sentence"`, `"passage"`, `"page"`.
-- `split_overlap` is an integer indicating the number of overlapping words, sentences, or passages between chunks, 0 being the default.
+- `split_by` can be `"word"` (default), `"sentence"`, `"passage"`, `"page"`, or `"token"`. With `"token"`, block sizes are counted in LLM tokens using tiktoken, so every block fits a known context budget. This requires `pip install tiktoken`.
+- `split_overlap` is an integer indicating the number of overlapping words, sentences, passages, or tokens between chunks, 0 being the default.
+- `tokenizer_encoding` is the tiktoken encoding used when `split_by="token"`. It defaults to `"o200k_base"`, used by current OpenAI models. Set it to match your model, for example `"cl100k_base"` for OpenAI's `text-embedding-3` embedding models.
 
 ## Usage
 

haystack/components/preprocessors/hierarchical_document_splitter.py

@@ -40,14 +40,20 @@ def __init__(
         self,
         block_sizes: set[int],
         split_overlap: int = 0,
-        split_by: Literal["word", "sentence", "page", "passage"] = "word",
+        split_by: Literal["word", "sentence", "page", "passage", "token"] = "word",
+        *,
+        tokenizer_encoding: str = "o200k_base",
     ) -> None:
         """
         Initialize HierarchicalDocumentSplitter.
 
         :param block_sizes: Set of block sizes to split the document into. The blocks are split in descending order.
         :param split_overlap: The number of overlapping units for each split.
-        :param split_by: The unit for splitting your documents.
+        :param split_by: The unit for splitting your documents: `word`, `sentence`, `page`, `passage`, or `token`.
+            `token` splits by token count using tiktoken (requires `pip install tiktoken`), so that block sizes
+            match an LLM's context budget.
+        :param tokenizer_encoding: The tiktoken encoding to use when `split_by="token"`. Defaults to `"o200k_base"`
+            (current OpenAI models). Only used when `split_by="token"`.
         :raises ValueError: If `block_sizes` is empty, if `split_overlap` is negative, or if `split_overlap` is
             greater than or equal to the smallest value in `block_sizes`.
         """
@@ -69,8 +75,16 @@ def __init__(
         self.splitters: dict[int, DocumentSplitter] = {}
         self.split_overlap = split_overlap
         self.split_by = split_by
+        self.tokenizer_encoding = tokenizer_encoding
         self._build_block_sizes()
 
+    def warm_up(self) -> None:
+        """
+        Warm up the splitters by loading the sentence tokenizer or the tiktoken encoding they need.
+        """
+        for splitter in self.splitters.values():
+            splitter.warm_up()
+
     @component.output_types(documents=list[Document])
     def run(self, documents: list[Document]) -> dict[str, list[Document]]:
         """
@@ -87,7 +101,10 @@ def run(self, documents: list[Document]) -> dict[str, list[Document]]:
     def _build_block_sizes(self) -> None:
         for block_size in self.block_sizes:
             self.splitters[block_size] = DocumentSplitter(
-                split_length=block_size, split_overlap=self.split_overlap, split_by=self.split_by
+                split_length=block_size,
+                split_overlap=self.split_overlap,
+                split_by=self.split_by,
+                tokenizer_encoding=self.tokenizer_encoding,
             )
 
     @staticmethod
@@ -141,7 +158,11 @@ def to_dict(self) -> dict[str, Any]:
                 Serialized dictionary representation of the component.
         """
         return default_to_dict(
-            self, block_sizes=self.block_sizes, split_overlap=self.split_overlap, split_by=self.split_by
+            self,
+            block_sizes=self.block_sizes,
+            split_overlap=self.split_overlap,
+            split_by=self.split_by,
+            tokenizer_encoding=self.tokenizer_encoding,
         )
 
     @classmethod

releasenotes/notes/hierarchical-splitter-token-split-62a86907ebb795f5.yaml

@@ -0,0 +1,8 @@
+---
+enhancements:
+  - |
+    Add ``split_by="token"`` to ``HierarchicalDocumentSplitter``, so block sizes are counted in LLM tokens using
+    tiktoken, as in ``DocumentSplitter``. The encoding defaults to ``"o200k_base"`` and can be changed with the new
+    ``tokenizer_encoding`` parameter. ``HierarchicalDocumentSplitter`` also gets a ``warm_up()`` method that loads the
+    sentence tokenizer or tiktoken encoding up front instead of on the first ``run()``. Requires ``pip install tiktoken``
+    for ``split_by="token"``.

test/components/preprocessors/test_hierarchical_doc_splitter.py

@@ -2,10 +2,13 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from unittest.mock import Mock
+
 import pytest
 
 from haystack import Document, Pipeline
 from haystack.components.preprocessors import HierarchicalDocumentSplitter
+from haystack.components.preprocessors import document_splitter as document_splitter_module
 from haystack.components.writers import DocumentWriter
 
 
@@ -43,7 +46,12 @@ def test_to_dict(self):
         expected = builder.to_dict()
         assert expected == {
             "type": "haystack.components.preprocessors.hierarchical_document_splitter.HierarchicalDocumentSplitter",
-            "init_parameters": {"block_sizes": [300, 200, 100], "split_overlap": 25, "split_by": "word"},
+            "init_parameters": {
+                "block_sizes": [300, 200, 100],
+                "split_overlap": 25,
+                "split_by": "word",
+                "tokenizer_encoding": "o200k_base",
+            },
         }
 
     def test_from_dict(self):
@@ -56,6 +64,8 @@ def test_from_dict(self):
         assert builder.block_sizes == [10, 5, 2]
         assert builder.split_overlap == 0
         assert builder.split_by == "word"
+        # dicts serialized before tokenizer_encoding existed get the default
+        assert builder.tokenizer_encoding == "o200k_base"
 
     def test_run(self):
         builder = HierarchicalDocumentSplitter(block_sizes={10, 5, 2}, split_overlap=0, split_by="word")
@@ -144,7 +154,12 @@ def test_to_dict_in_pipeline(self, in_memory_doc_store):
 
         assert expected["components"]["hierarchical_doc_splitter"] == {
             "type": "haystack.components.preprocessors.hierarchical_document_splitter.HierarchicalDocumentSplitter",
-            "init_parameters": {"block_sizes": [10, 5, 2], "split_overlap": 0, "split_by": "word"},
+            "init_parameters": {
+                "block_sizes": [10, 5, 2],
+                "split_overlap": 0,
+                "split_by": "word",
+                "tokenizer_encoding": "o200k_base",
+            },
         }
 
     def test_from_dict_in_pipeline(self):
@@ -276,3 +291,72 @@ def test_hierarchical_splitter_multiple_block_sizes(self):
                 for child in children:
                     assert child.meta["__parent_id"] == doc.id
                     assert child.meta["__level"] == doc.meta["__level"] + 1
+
+    def test_init_with_token_split(self):
+        splitter = HierarchicalDocumentSplitter(
+            block_sizes={100, 20}, split_by="token", tokenizer_encoding="cl100k_base"
+        )
+        assert splitter.split_by == "token"
+        assert splitter.tokenizer_encoding == "cl100k_base"
+        for block_size, inner_splitter in splitter.splitters.items():
+            assert inner_splitter.split_by == "token"
+            assert inner_splitter.split_length == block_size
+            assert inner_splitter.tokenizer_encoding == "cl100k_base"
+
+    def test_serialization_round_trip_with_token_split(self):
+        splitter = HierarchicalDocumentSplitter(
+            block_sizes={100, 20}, split_overlap=5, split_by="token", tokenizer_encoding="cl100k_base"
+        )
+        restored = HierarchicalDocumentSplitter.from_dict(splitter.to_dict())
+        assert restored.block_sizes == [100, 20]
+        assert restored.split_overlap == 5
+        assert restored.split_by == "token"
+        assert restored.tokenizer_encoding == "cl100k_base"
+        assert restored.splitters[20].tokenizer_encoding == "cl100k_base"
+
+    def test_warm_up_loads_the_encoding_for_every_block_size_once(self, monkeypatch):
+        encoding = Mock()
+        get_encoding = Mock(return_value=encoding)
+        monkeypatch.setattr(document_splitter_module.tiktoken, "get_encoding", get_encoding)
+
+        splitter = HierarchicalDocumentSplitter(
+            block_sizes={100, 20}, split_by="token", tokenizer_encoding="cl100k_base"
+        )
+        splitter.warm_up()
+        splitter.warm_up()
+
+        # one call per block size, the second warm_up() is a no-op
+        assert get_encoding.call_count == 2
+        get_encoding.assert_called_with("cl100k_base")
+        assert all(inner_splitter._tiktoken_tokenizer is encoding for inner_splitter in splitter.splitters.values())
+
+
+@pytest.mark.integration
+class TestHierarchicalDocumentSplitterByTokenIntegration:
+    """Integration tests for split_by="token" that load a real tiktoken encoding."""
+
+    @pytest.mark.parametrize("encoding", ["o200k_base", "cl100k_base"])
+    def test_run_split_by_token(self, encoding):
+        text = (
+            "Haystack pipelines connect retrievers, rankers and generators. Hierarchical chunks let an "
+            "AutoMergingRetriever return a parent block when enough of its children match.\f"
+            "Token-based block sizes keep every chunk inside the context budget of the model."
+        )
+        splitter = HierarchicalDocumentSplitter(block_sizes={20, 8}, split_by="token", tokenizer_encoding=encoding)
+        splitter.warm_up()
+
+        documents = splitter.run([Document(content=text)])["documents"]
+
+        tokenizer = splitter.splitters[20]._tiktoken_tokenizer
+        assert tokenizer is not None
+        docs_by_id = {doc.id: doc for doc in documents}
+        assert {doc.meta["__level"] for doc in documents} == {0, 1, 2}
+        for doc in documents[1:]:
+            assert doc.content is not None
+            assert len(tokenizer.encode_ordinary(doc.content)) <= doc.meta["__block_size"]
+            assert doc.id in docs_by_id[doc.meta["__parent_id"]].meta["__children_ids"]
+        # without overlap, the children of every block add up to exactly the block's text
+        for doc in documents:
+            if doc.meta["__children_ids"]:
+                children_text = "".join(str(docs_by_id[child_id].content) for child_id in doc.meta["__children_ids"])
+                assert children_text == doc.content