feat: make the tiktoken encoding configurable in RecursiveDocumentSpl…

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

@@ -31,6 +31,7 @@ The `RecursiveDocumentSplitter` expects a list of documents as input and returns
 - `split_unit`: The unit of the `split_length` parameter. Can be either `"word"`, `"char"`, or `"token"`.
 - `separators`: An optional list of separator strings to use for splitting the text. If you don’t provide any separators, the default ones are `["\n\n", "sentence", "\n", " "]`. The string separators will be treated as regular expressions. If the separator is `"sentence"`, the text will be split into sentences using a custom sentence tokenizer based on NLTK. See [SentenceSplitter](https://github.com/deepset-ai/haystack/blob/main/haystack/components/preprocessors/sentence_tokenizer.py#L116) code for more information.
 - `sentence_splitter_params`: Optional parameters to pass to the [SentenceSplitter](https://github.com/deepset-ai/haystack/blob/main/haystack/components/preprocessors/sentence_tokenizer.py#L116).
+- `tokenizer_encoding`: The tiktoken encoding used when `split_unit="token"`. 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, so that chunks fit the model's token limit.
 
 When `split_unit="token"`, strings such as `<|endoftext|>` in the source document are tokenized as ordinary text. They are preserved in the chunks and count toward the token limit.
 

haystack/components/preprocessors/recursive_splitter.py

@@ -65,6 +65,7 @@ def __init__(
         split_unit: Literal["word", "char", "token"] = "word",
         separators: list[str] | None = None,
         sentence_splitter_params: dict[str, Any] | None = None,
+        tokenizer_encoding: str = "o200k_base",
     ) -> None:
         """
         Initializes a RecursiveDocumentSplitter.
@@ -74,7 +75,7 @@ def __init__(
         :param split_overlap: The number of overlapping units (words, characters, or tokens, per
             `split_unit`) between consecutive chunks.
         :param split_unit: The unit of the split_length parameter. It can be either "word", "char", or "token".
-            If "token" is selected, the text will be split into tokens using the tiktoken tokenizer (o200k_base).
+            If "token" is selected, the text will be split into tokens using tiktoken (see `tokenizer_encoding`).
             Special-token strings in document content are encoded as ordinary text.
         :param separators: An optional list of separator strings to use for splitting the text. The string
             separators will be treated as regular expressions unless the separator is "sentence", in that case the
@@ -85,6 +86,8 @@ def __init__(
             See: haystack.components.preprocessors.sentence_tokenizer.SentenceSplitter for more information.
             The chunks keep the whitespace between sentences whether `keep_white_spaces` is True or False; it only
             changes the rules the tokenizer uses to decide where a sentence ends.
+        :param tokenizer_encoding: The tiktoken encoding to use when `split_unit="token"`. Defaults to `"o200k_base"`
+            (current OpenAI models). Only used when `split_unit="token"`.
 
         :raises ValueError: If the overlap is greater than or equal to the chunk size or if the overlap is negative, or
                             if any separator is not a string.
@@ -98,6 +101,7 @@ def __init__(
         self.sentence_splitter_params = (
             {"keep_white_spaces": True} if sentence_splitter_params is None else sentence_splitter_params
         )
+        self.tokenizer_encoding = tokenizer_encoding
         self.tiktoken_tokenizer: "tiktoken.Encoding" | None = None
         self._is_warmed_up = False
 
@@ -111,7 +115,7 @@ def warm_up(self) -> None:
             self.nltk_tokenizer = self._get_custom_sentence_tokenizer(self.sentence_splitter_params)
         if self.split_units == "token":
             tiktoken_imports.check()
-            self.tiktoken_tokenizer = tiktoken.get_encoding("o200k_base")
+            self.tiktoken_tokenizer = tiktoken.get_encoding(self.tokenizer_encoding)
         self._is_warmed_up = True
 
     def to_dict(self) -> dict[str, Any]:
@@ -128,6 +132,7 @@ def to_dict(self) -> dict[str, Any]:
             split_unit=self.split_units,
             separators=self.separators,
             sentence_splitter_params=self.sentence_splitter_params,
+            tokenizer_encoding=self.tokenizer_encoding,
         )
 
     def _check_params(self) -> None:

releasenotes/notes/recursive-splitter-tokenizer-encoding-5b7b308985596dcc.yaml

@@ -0,0 +1,7 @@
+---
+enhancements:
+  - |
+    Add a ``tokenizer_encoding`` parameter to ``RecursiveDocumentSplitter`` to choose the tiktoken encoding used
+    when ``split_unit="token"``, as in ``DocumentSplitter``. It defaults to ``"o200k_base"``, the encoding used
+    until now. Set it to match your model, for example ``"cl100k_base"`` for OpenAI's ``text-embedding-3`` models,
+    so that chunks fit the model's token limit.

test/components/preprocessors/test_recursive_splitter.py

@@ -1151,6 +1151,26 @@ def test_special_token_strings_are_split_as_literal_text(split_overlap):
     assert reconstructed == text
 
 
+@pytest.mark.integration
+@pytest.mark.parametrize("encoding", ["o200k_base", "cl100k_base"])
+def test_run_split_by_token_count_with_tokenizer_encoding(encoding):
+    splitter = RecursiveDocumentSplitter(
+        split_length=10, separators=["."], split_unit="token", tokenizer_encoding=encoding
+    )
+    text = (
+        "Haystack pipelines connect retrievers, rankers and generators into search systems. "
+        "Embedding models count tokens with their own encoding. Chunks must fit that limit."
+    )
+    chunks = splitter.run([Document(content=text)])["documents"]
+
+    assert splitter.tiktoken_tokenizer is not None
+    assert splitter.tiktoken_tokenizer.name == encoding
+    assert len(chunks) > 1
+    for chunk in chunks:
+        assert chunk.content is not None
+        assert len(splitter.tiktoken_tokenizer.encode_ordinary(chunk.content)) <= 10
+
+
 @pytest.mark.integration
 def test_run_split_by_token_count():
     splitter = RecursiveDocumentSplitter(split_length=5, separators=["."], split_unit="token")
@@ -1426,6 +1446,35 @@ def test_warm_up_is_idempotent_token(monkeypatch):
     assert splitter.tiktoken_tokenizer is sentinel
 
 
+def test_warm_up_uses_tokenizer_encoding(monkeypatch):
+    import haystack.components.preprocessors.recursive_splitter as mod
+
+    get_encoding = Mock()
+    monkeypatch.setattr(mod.tiktoken, "get_encoding", get_encoding)
+
+    splitter = RecursiveDocumentSplitter(split_unit="token", split_length=10, tokenizer_encoding="cl100k_base")
+    splitter.warm_up()
+
+    get_encoding.assert_called_once_with("cl100k_base")
+
+
+def test_serialization_keeps_tokenizer_encoding():
+    splitter = RecursiveDocumentSplitter(split_length=10, split_unit="token", tokenizer_encoding="cl100k_base")
+    pipeline = Pipeline()
+    pipeline.add_component("chunker", splitter)
+
+    restored = Pipeline.loads(pipeline.dumps()).get_component("chunker")
+    assert isinstance(restored, RecursiveDocumentSplitter)
+    assert restored.tokenizer_encoding == "cl100k_base"
+
+    # pipelines serialized before tokenizer_encoding existed get the default
+    data = pipeline.to_dict()
+    del data["components"]["chunker"]["init_parameters"]["tokenizer_encoding"]
+    restored = Pipeline.from_dict(data).get_component("chunker")
+    assert isinstance(restored, RecursiveDocumentSplitter)
+    assert restored.tokenizer_encoding == "o200k_base"
+
+
 def test_fallback_overlap_char_unit():
     """split_overlap must be applied even when no separator matches (char unit)."""
     splitter = RecursiveDocumentSplitter(split_length=5, split_overlap=2, separators=["\n\n"], split_unit="char")