feat: DocumentSplitter — add split_by="token" mode using tiktoken

haystack/components/preprocessors/document_splitter.py

@@ -11,8 +11,12 @@
 from haystack import Document, component, logging
 from haystack.components.preprocessors.sentence_tokenizer import Language, SentenceSplitter, nltk_imports
 from haystack.core.serialization import default_from_dict, default_to_dict
+from haystack.lazy_imports import LazyImport
 from haystack.utils import deserialize_callable, serialize_callable
 
+with LazyImport("Run 'pip install tiktoken'") as tiktoken_imports:
+    import tiktoken
+
 logger = logging.getLogger(__name__)
 
 # mapping of split by character, 'function' and 'sentence' don't split by character
@@ -54,7 +58,7 @@ class DocumentSplitter:
 
     def __init__(
         self,
-        split_by: Literal["function", "page", "passage", "period", "word", "line", "sentence"] = "word",
+        split_by: Literal["function", "page", "passage", "period", "word", "line", "sentence", "token"] = "word",
         split_length: int = 200,
         split_overlap: int = 0,
         split_threshold: int = 0,
@@ -65,6 +69,7 @@ def __init__(
         extend_abbreviations: bool = True,
         *,
         skip_empty_documents: bool = True,
+        tokenizer_encoding: str = "o200k_base",
     ) -> None:
         """
         Initialize DocumentSplitter.
@@ -76,6 +81,7 @@ def __init__(
             - `passage` for splitting by double line breaks ("\\n\\n")
             - `line` for splitting each line ("\\n")
             - `sentence` for splitting by NLTK sentence tokenizer
+            - `token` for splitting by token count using tiktoken (requires ``pip install tiktoken``)
 
         :param split_length: The maximum number of units in each split.
         :param split_overlap: The number of overlapping units for each split.
@@ -93,6 +99,8 @@ def __init__(
         :param skip_empty_documents: Choose whether to skip documents with empty content. Default is True.
             Set to False when downstream components in the Pipeline (like LLMDocumentContentExtractor) can extract text
             from non-textual documents.
+        :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"``.
         """
 
         self.split_by = split_by
@@ -105,6 +113,7 @@ def __init__(
         self.use_split_rules = use_split_rules
         self.extend_abbreviations = extend_abbreviations
         self.skip_empty_documents = skip_empty_documents
+        self.tokenizer_encoding = tokenizer_encoding
 
         self._init_checks(
             split_by=split_by,
@@ -117,6 +126,9 @@ def __init__(
         if self._use_sentence_splitter:
             nltk_imports.check()
             self.sentence_splitter: SentenceSplitter | None = None
+        if split_by == "token":
+            tiktoken_imports.check()
+            self._tiktoken_tokenizer: "tiktoken.Encoding | None" = None
 
     def _init_checks(
         self,
@@ -137,7 +149,7 @@ def _init_checks(
         :param respect_sentence_boundary: Whether to respect sentence boundaries when splitting
         :raises ValueError: If any parameter is invalid
         """
-        valid_split_by = ["function", "page", "passage", "period", "word", "line", "sentence"]
+        valid_split_by = ["function", "page", "passage", "period", "word", "line", "sentence", "token"]
         if split_by not in valid_split_by:
             raise ValueError(f"split_by must be one of {', '.join(valid_split_by)}.")
 
@@ -162,7 +174,7 @@ def _init_checks(
 
     def warm_up(self) -> None:
         """
-        Warm up the DocumentSplitter by loading the sentence tokenizer.
+        Warm up the DocumentSplitter by loading the sentence tokenizer or tiktoken encoder.
         """
         if self._use_sentence_splitter and self.sentence_splitter is None:
             self.sentence_splitter = SentenceSplitter(
@@ -171,6 +183,8 @@ def warm_up(self) -> None:
                 extend_abbreviations=self.extend_abbreviations,
                 keep_white_spaces=True,
             )
+        if self.split_by == "token" and self._tiktoken_tokenizer is None:
+            self._tiktoken_tokenizer = tiktoken.get_encoding(self.tokenizer_encoding)
 
     @component.output_types(documents=list[Document])
     def run(self, documents: list[Document]) -> dict[str, list[Document]]:
@@ -192,6 +206,8 @@ def run(self, documents: list[Document]) -> dict[str, list[Document]]:
         """
         if self._use_sentence_splitter and self.sentence_splitter is None:
             self.warm_up()
+        if self.split_by == "token" and self._tiktoken_tokenizer is None:
+            self.warm_up()
 
         if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
             raise TypeError("DocumentSplitter expects a List of Documents as input.")
@@ -216,6 +232,9 @@ def _split_document(self, doc: Document) -> list[Document]:
         if self.split_by == "function" and self.splitting_function is not None:
             return self._split_by_function(doc)
 
+        if self.split_by == "token":
+            return self._split_by_token(doc)
+
         return self._split_by_character(doc)
 
     def _split_by_nltk_sentence(self, doc: Document) -> list[Document]:
@@ -243,6 +262,40 @@ def _split_by_nltk_sentence(self, doc: Document) -> list[Document]:
 
         return split_docs
 
+    def _split_by_token(self, doc: Document) -> list[Document]:
+        """
+        Split a document by token count using tiktoken.
+
+        Encodes the full document text to tokens, slices into chunks of ``split_length`` tokens
+        with ``split_overlap`` overlap, then decodes each chunk back to a string.
+        """
+        tokens = self._tiktoken_tokenizer.encode(doc.content)  # type: ignore[union-attr, arg-type]
+        step = self.split_length - self.split_overlap
+
+        text_splits: list[str] = []
+        splits_pages: list[int] = []
+        splits_start_idxs: list[int] = []
+        cur_page = 1
+        cur_start_idx = 0
+
+        for i in range(0, len(tokens), step):
+            chunk_tokens = tokens[i : i + self.split_length]
+            chunk_text = self._tiktoken_tokenizer.decode(chunk_tokens)  # type: ignore[union-attr]
+            text_splits.append(chunk_text)
+            splits_pages.append(cur_page)
+            splits_start_idxs.append(cur_start_idx)
+
+            # Advance by the non-overlapping prefix only
+            non_overlap_text = self._tiktoken_tokenizer.decode(tokens[i : i + step])  # type: ignore[union-attr]
+            cur_page += non_overlap_text.count("\f")
+            cur_start_idx += len(non_overlap_text)
+
+        metadata = deepcopy(doc.meta)
+        metadata["source_id"] = doc.id
+        return self._create_docs_from_splits(
+            text_splits=text_splits, splits_pages=splits_pages, splits_start_idxs=splits_start_idxs, meta=metadata
+        )
+
     def _split_by_character(self, doc: Document) -> list[Document]:
         split_at = _CHARACTER_SPLIT_BY_MAPPING[self.split_by]
         units = doc.content.split(split_at)  # type: ignore[union-attr]
@@ -387,6 +440,7 @@ def to_dict(self) -> dict[str, Any]:
             use_split_rules=self.use_split_rules,
             extend_abbreviations=self.extend_abbreviations,
             skip_empty_documents=self.skip_empty_documents,
+            tokenizer_encoding=self.tokenizer_encoding,
         )
         if self.splitting_function:
             serialized["init_parameters"]["splitting_function"] = serialize_callable(self.splitting_function)

releasenotes/notes/document-splitter-split-by-token-8a3f69718d05dfad.yaml

@@ -0,0 +1,8 @@
+---
+enhancements:
+  - |
+    Add ``split_by="token"`` to ``DocumentSplitter``. Splits text by LLM token
+    count using tiktoken. The encoding defaults to ``"o200k_base"`` (current
+    OpenAI models) and can be changed via the new ``tokenizer_encoding``
+    parameter. ``split_length``, ``split_overlap``, and ``split_threshold`` all
+    work as usual. Requires ``pip install tiktoken``.

test/components/preprocessors/test_document_splitter.py

@@ -846,3 +846,73 @@ def test_duplicate_pages_get_different_doc_id(self):
         result = splitter.run(documents=[doc1])
 
         assert len({doc.id for doc in result["documents"]}) == 4
+
+
+@pytest.mark.integration
+class TestSplittingByToken:
+    """Tests for split_by="token" mode. Require tiktoken to be installed."""
+
+    def test_basic_chunking(self):
+        splitter = DocumentSplitter(split_by="token", split_length=5, split_overlap=0)
+        doc = Document(content="one two three four five six seven eight nine ten")
+        result = splitter.run(documents=[doc])["documents"]
+        assert len(result) > 1
+        for chunk in result:
+            # tiktoken: each simple word is 1 token, so no chunk should exceed 5 tokens
+            tokens = splitter._tiktoken_tokenizer.encode(chunk.content)
+            assert len(tokens) <= 5
+
+    def test_metadata_set(self):
+        splitter = DocumentSplitter(split_by="token", split_length=5, split_overlap=0)
+        doc = Document(content="one two three four five six seven eight nine ten")
+        result = splitter.run(documents=[doc])["documents"]
+        for i, chunk in enumerate(result):
+            assert chunk.meta["source_id"] == doc.id
+            assert chunk.meta["split_id"] == i
+            assert "split_idx_start" in chunk.meta
+            assert chunk.meta["page_number"] == 1
+
+    def test_overlap_produces_shared_text(self):
+        splitter = DocumentSplitter(split_by="token", split_length=6, split_overlap=2)
+        # 12 simple tokens → expect > 2 chunks, each pair sharing 2 tokens of text
+        doc = Document(content="a b c d e f g h i j k l")
+        result = splitter.run(documents=[doc])["documents"]
+        assert len(result) > 1
+        for i in range(len(result) - 1):
+            # The start index of chunk[i+1] must be less than the end index of chunk[i] due to overlap
+            assert result[i + 1].meta["split_idx_start"] < result[i].meta["split_idx_start"] + len(result[i].content)
+
+    def test_page_tracking(self):
+        splitter = DocumentSplitter(split_by="token", split_length=5, split_overlap=0)
+        # 5 tokens on page 1, form-feed, then more tokens on page 2
+        doc = Document(content="a b c d e\fg h i j k")
+        result = splitter.run(documents=[doc])["documents"]
+        assert result[0].meta["page_number"] == 1
+        # The second chunk (after the \f) should be on page 2
+        assert result[-1].meta["page_number"] == 2
+
+    def test_empty_document_skipped(self, caplog):
+        splitter = DocumentSplitter(split_by="token", split_length=5)
+        result = splitter.run(documents=[Document(content="")])["documents"]
+        assert result == []
+        assert "has an empty content. Skipping this document." in caplog.text
+
+    def test_custom_encoding(self):
+        splitter = DocumentSplitter(split_by="token", split_length=5, tokenizer_encoding="cl100k_base")
+        doc = Document(content="one two three four five six seven eight")
+        result = splitter.run(documents=[doc])["documents"]
+        assert len(result) > 0
+        assert splitter.tokenizer_encoding == "cl100k_base"
+
+    def test_serialization_roundtrip(self):
+        splitter = DocumentSplitter(
+            split_by="token", split_length=10, split_overlap=2, tokenizer_encoding="cl100k_base"
+        )
+        serialized = splitter.to_dict()
+        assert serialized["init_parameters"]["split_by"] == "token"
+        assert serialized["init_parameters"]["tokenizer_encoding"] == "cl100k_base"
+        restored = DocumentSplitter.from_dict(serialized)
+        assert restored.split_by == "token"
+        assert restored.split_length == 10
+        assert restored.split_overlap == 2
+        assert restored.tokenizer_encoding == "cl100k_base"