Coverage for haystack/components/preprocessors/recursive_splitter.py: 92%
242 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import re
6from copy import deepcopy
7from typing import Any, Literal
9from haystack import Document, component, logging
10from haystack.lazy_imports import LazyImport
12with LazyImport("Run 'pip install tiktoken'") as tiktoken_imports:
13 import tiktoken
15logger = logging.getLogger(__name__)
18@component
19class RecursiveDocumentSplitter:
20 """
21 Recursively chunk text into smaller chunks.
23 This component is used to split text into smaller chunks, it does so by recursively applying a list of separators
24 to the text.
26 The separators are applied in the order they are provided, typically this is a list of separators that are
27 applied in a specific order, being the last separator the most specific one.
29 Each separator is applied to the text, it then checks each of the resulting chunks, it keeps the chunks that
30 are within the split_length, for the ones that are larger than the split_length, it applies the next separator in the
31 list to the remaining text.
33 This is done until all chunks are smaller than the split_length parameter.
35 Example:
37 ```python
38 from haystack import Document
39 from haystack.components.preprocessors import RecursiveDocumentSplitter
41 chunker = RecursiveDocumentSplitter(split_length=15, split_overlap=0, separators=["\\n\\n", "\\n", ".", " "])
42 text = ('''Artificial intelligence (AI) - Introduction
44 AI, in its broadest sense, is intelligence exhibited by machines, particularly computer systems.
45 AI technology is widely used throughout industry, government, and science. Some high-profile applications include advanced web search engines; recommendation systems; interacting via human speech; autonomous vehicles; generative and creative tools; and superhuman play and analysis in strategy games.''')
46 doc = Document(content=text)
47 doc_chunks = chunker.run([doc])
48 print(doc_chunks["documents"])
49 # [
50 # Document(id=..., content: 'Artificial intelligence (AI) - Introduction\\n\\n', meta: {'source_id': '...', 'parent_id': '...', 'split_id': 0, 'split_idx_start': 0, '_split_overlap': None, 'page_number': 1})
51 # Document(id=..., content: 'AI, in its broadest sense, is intelligence exhibited by machines, particularly computer systems.\\n', meta: {'source_id': '...', 'parent_id': '...', 'split_id': 1, 'split_idx_start': 45, '_split_overlap': None, 'page_number': 1})
52 # Document(id=..., content: 'AI technology is widely used throughout industry, government, and science.', meta: {'source_id': '...', 'parent_id': '...', 'split_id': 2, 'split_idx_start': 142, '_split_overlap': None, 'page_number': 1})
53 # Document(id=..., content: ' Some high-profile applications include advanced web search engines; recommendation systems; interac...', meta: {'source_id': '...', 'parent_id': '...', 'split_id': 3, 'split_idx_start': 216, '_split_overlap': None, 'page_number': 1})
54 # Document(id=..., content: 'vehicles; generative and creative tools; and superhuman play and analysis in strategy games.', meta: {'source_id': '...', 'parent_id': '...', 'split_id': 4, 'split_idx_start': 350, '_split_overlap': None, 'page_number': 1})
55 # ]
56 ```
57 """ # noqa: E501
59 def __init__(
60 self,
61 *,
62 split_length: int = 200,
63 split_overlap: int = 0,
64 split_unit: Literal["word", "char", "token"] = "word",
65 separators: list[str] | None = None,
66 sentence_splitter_params: dict[str, Any] | None = None,
67 ) -> None:
68 """
69 Initializes a RecursiveDocumentSplitter.
71 :param split_length: The maximum length of each chunk by default in words, but can be in characters or tokens.
72 See the `split_units` parameter.
73 :param split_overlap: The number of characters to overlap between consecutive chunks.
74 :param split_unit: The unit of the split_length parameter. It can be either "word", "char", or "token".
75 If "token" is selected, the text will be split into tokens using the tiktoken tokenizer (o200k_base).
76 :param separators: An optional list of separator strings to use for splitting the text. The string
77 separators will be treated as regular expressions unless the separator is "sentence", in that case the
78 text will be split into sentences using a custom sentence tokenizer based on NLTK.
79 See: haystack.components.preprocessors.sentence_tokenizer.SentenceSplitter.
80 If no separators are provided, the default separators ["\\n\\n", "sentence", "\\n", " "] are used.
81 :param sentence_splitter_params: Optional parameters to pass to the sentence tokenizer.
82 See: haystack.components.preprocessors.sentence_tokenizer.SentenceSplitter for more information.
84 :raises ValueError: If the overlap is greater than or equal to the chunk size or if the overlap is negative, or
85 if any separator is not a string.
86 """
87 self.split_length = split_length
88 self.split_overlap = split_overlap
89 self.split_units = split_unit
90 self.separators = separators if separators else ["\n\n", "sentence", "\n", " "] # default separators
91 self._check_params()
92 self.nltk_tokenizer = None
93 self.sentence_splitter_params = (
94 {"keep_white_spaces": True} if sentence_splitter_params is None else sentence_splitter_params
95 )
96 self.tiktoken_tokenizer: "tiktoken.Encoding" | None = None
97 self._is_warmed_up = False
99 def warm_up(self) -> None:
100 """
101 Warm up the sentence tokenizer and tiktoken tokenizer if needed.
102 """
103 if self._is_warmed_up:
104 return
105 if "sentence" in self.separators:
106 self.nltk_tokenizer = self._get_custom_sentence_tokenizer(self.sentence_splitter_params)
107 if self.split_units == "token":
108 tiktoken_imports.check()
109 self.tiktoken_tokenizer = tiktoken.get_encoding("o200k_base")
110 self._is_warmed_up = True
112 def _check_params(self) -> None:
113 if self.split_length < 1:
114 raise ValueError("Split length must be at least 1 character.")
115 if self.split_overlap < 0:
116 raise ValueError("Overlap must be greater than zero.")
117 if self.split_overlap >= self.split_length:
118 raise ValueError("Overlap cannot be greater than or equal to the chunk size.")
119 if not all(isinstance(separator, str) for separator in self.separators):
120 raise ValueError("All separators must be strings.")
122 @staticmethod
123 def _get_custom_sentence_tokenizer(sentence_splitter_params: dict[str, Any]) -> Any:
124 from haystack.components.preprocessors.sentence_tokenizer import SentenceSplitter
126 return SentenceSplitter(**sentence_splitter_params)
128 def _split_chunk(self, current_chunk: str) -> tuple[str, str]:
129 """
130 Splits a chunk based on the split_length and split_units attribute.
132 :param current_chunk: The current chunk to be split.
133 :returns:
134 A tuple containing the current chunk and the remaining chunk.
135 """
136 if self.split_units == "word":
137 words = current_chunk.split()
138 current_chunk = " ".join(words[: self.split_length])
139 remaining_words = words[self.split_length :]
140 return current_chunk, " ".join(remaining_words)
141 if self.split_units == "char":
142 text = current_chunk
143 current_chunk = text[: self.split_length]
144 remaining_chars = text[self.split_length :]
145 return current_chunk, remaining_chars
147 # at this point we know that the tokenizer is already initialized
148 tokens = self.tiktoken_tokenizer.encode(current_chunk) # type: ignore
149 current_tokens = tokens[: self.split_length]
150 remaining_tokens = tokens[self.split_length :]
151 return self.tiktoken_tokenizer.decode(current_tokens), self.tiktoken_tokenizer.decode(remaining_tokens) # type: ignore
153 def _apply_overlap(self, chunks: list[str]) -> list[str]:
154 """
155 Applies an overlap between consecutive chunks if the chunk_overlap attribute is greater than zero.
157 Works for both word- and character-level splitting. It trims the last chunk if it exceeds the split_length and
158 adds the trimmed content to the next chunk. If the last chunk is still too long after trimming, it splits it
159 and adds the first chunk to the list. This process continues until the last chunk is within the split_length.
161 :param chunks: A list of text chunks.
162 :returns:
163 A list of text chunks with the overlap applied.
164 """
165 overlapped_chunks: list[str] = []
167 for idx, chunk in enumerate(chunks):
168 if idx == 0:
169 overlapped_chunks.append(chunk)
170 continue
172 # get the overlap between the current and previous chunk
173 overlap, prev_chunk = self._get_overlap(overlapped_chunks)
174 if overlap == prev_chunk:
175 logger.warning(
176 "Overlap is the same as the previous chunk. "
177 "Consider increasing the `split_length` parameter or decreasing the `split_overlap` parameter."
178 )
180 current_chunk = self._create_chunk_starting_with_overlap(chunk, overlap)
182 # if this new chunk exceeds 'split_length', trim it and move the remaining text to the next chunk
183 # if this is the last chunk, another new chunk will contain the trimmed text preceded by the overlap
184 # of the last chunk
185 if self._chunk_length(current_chunk) > self.split_length:
186 current_chunk, remaining_text = self._split_chunk(current_chunk)
187 if idx < len(chunks) - 1:
188 if self.split_units == "word":
189 chunks[idx + 1] = remaining_text + " " + chunks[idx + 1]
190 elif self.split_units == "token":
191 # For token-based splitting, combine at token level
192 # at this point we know that the tokenizer is already initialized
193 remaining_tokens = self.tiktoken_tokenizer.encode(remaining_text) # type: ignore
194 next_chunk_tokens = self.tiktoken_tokenizer.encode(chunks[idx + 1]) # type: ignore
195 chunks[idx + 1] = self.tiktoken_tokenizer.decode(remaining_tokens + next_chunk_tokens) # type: ignore
196 else: # char
197 chunks[idx + 1] = remaining_text + chunks[idx + 1]
198 elif remaining_text:
199 # create a new chunk with the trimmed text preceded by the overlap of the last chunk
200 overlapped_chunks.append(current_chunk)
201 chunk = remaining_text
202 overlap, _ = self._get_overlap(overlapped_chunks)
203 current_chunk = self._create_chunk_starting_with_overlap(chunk, overlap)
205 overlapped_chunks.append(current_chunk)
207 # it can still be that the new last chunk exceeds the 'split_length'
208 # continue splitting until the last chunk is within 'split_length'
209 if idx == len(chunks) - 1 and self._chunk_length(current_chunk) > self.split_length:
210 last_chunk = overlapped_chunks.pop()
211 first_chunk, remaining_chunk = self._split_chunk(last_chunk)
212 overlapped_chunks.append(first_chunk)
214 while remaining_chunk:
215 # combine overlap with remaining chunk
216 overlap, _ = self._get_overlap(overlapped_chunks)
217 current = self._create_chunk_starting_with_overlap(remaining_chunk, overlap)
219 # if it fits within split_length we are done
220 if self._chunk_length(current) <= self.split_length:
221 overlapped_chunks.append(current)
222 break
224 # otherwise split it again
225 first_chunk, remaining_chunk = self._split_chunk(current)
226 overlapped_chunks.append(first_chunk)
228 return overlapped_chunks
230 def _create_chunk_starting_with_overlap(self, chunk: str, overlap: str) -> str:
231 if self.split_units == "word":
232 current_chunk = overlap + " " + chunk
233 elif self.split_units == "token":
234 # For token-based splitting, combine at token level
235 # at this point we know that the tokenizer is already initialized
236 overlap_tokens = self.tiktoken_tokenizer.encode(overlap) # type: ignore
237 chunk_tokens = self.tiktoken_tokenizer.encode(chunk) # type: ignore
238 current_chunk = self.tiktoken_tokenizer.decode(overlap_tokens + chunk_tokens) # type: ignore
239 else: # char
240 current_chunk = overlap + chunk
241 return current_chunk
243 def _get_overlap(self, overlapped_chunks: list[str]) -> tuple[str, str]:
244 """Get the previous overlapped chunk instead of the original chunk."""
245 prev_chunk = overlapped_chunks[-1]
246 overlap_start = max(0, self._chunk_length(prev_chunk) - self.split_overlap)
248 if self.split_units == "word":
249 word_chunks = prev_chunk.split()
250 overlap = " ".join(word_chunks[overlap_start:])
251 elif self.split_units == "token":
252 # For token-based splitting, handle overlap at token level
253 # at this point we know that the tokenizer is already initialized
254 tokens = self.tiktoken_tokenizer.encode(prev_chunk) # type: ignore
255 overlap_tokens = tokens[overlap_start:]
256 overlap = self.tiktoken_tokenizer.decode(overlap_tokens) # type: ignore
257 else: # char
258 overlap = prev_chunk[overlap_start:]
260 return overlap, prev_chunk
262 def _chunk_length(self, text: str) -> int:
263 """
264 Get the length of the chunk in the specified units (words, characters, or tokens).
266 :param text: The text to measure.
267 :returns: The length of the text in the specified units.
268 """
269 if self.split_units == "word":
270 words = [word for word in text.split(" ") if word]
271 return len(words)
272 if self.split_units == "char":
273 return len(text)
274 # token
275 # at this point we know that the tokenizer is already initialized
276 return len(self.tiktoken_tokenizer.encode(text)) # type: ignore
278 def _chunk_text(self, text: str) -> list[str]:
279 """
280 Recursive chunking algorithm that divides text into smaller chunks based on a list of separator characters.
282 It starts with a list of separator characters (e.g., ["\n\n", "sentence", "\n", " "]) and attempts to divide
283 the text using the first separator. If the resulting chunks are still larger than the specified chunk size,
284 it moves to the next separator in the list. This process continues recursively, progressively applying each
285 specific separator until the chunks meet the desired size criteria.
287 :param text: The text to be split into chunks.
288 :returns:
289 A list of text chunks.
290 """
291 if self._chunk_length(text) <= self.split_length:
292 return [text]
294 for curr_separator in self.separators:
295 if curr_separator == "sentence":
296 # re. ignore: correct SentenceSplitter initialization is checked at the initialization of the component
297 sentence_with_spans = self.nltk_tokenizer.split_sentences(text) # type: ignore
298 splits = [sentence["sentence"] for sentence in sentence_with_spans]
299 else:
300 # add escape "\" to the separator and wrapped it in a group so that it's included in the splits as well
301 escaped_separator = re.escape(curr_separator)
302 escaped_separator = f"({escaped_separator})"
304 # split the text and merge every two consecutive splits, i.e.: the text and the separator after it
305 splits = re.split(escaped_separator, text)
306 splits = [
307 "".join([splits[i], splits[i + 1]]) if i < len(splits) - 1 else splits[i]
308 for i in range(0, len(splits), 2)
309 ]
311 # remove last split if it's empty
312 splits = splits[:-1] if splits[-1] == "" else splits
314 if len(splits) == 1: # go to next separator, if current separator not found in the text
315 continue
317 chunks = []
318 current_chunk: list[str] = []
319 current_length = 0
321 # check splits, if any is too long, recursively chunk it, otherwise add to current chunk
322 for split in splits:
323 split_text = split
325 # if adding this split exceeds chunk_size, process current_chunk
326 if current_length + self._chunk_length(split_text) > self.split_length:
327 # process current_chunk
328 if current_chunk: # keep the good splits
329 chunks.append("".join(current_chunk))
330 current_chunk = []
331 current_length = 0
333 # recursively handle splits that are too large
334 if self._chunk_length(split_text) > self.split_length:
335 if curr_separator == self.separators[-1]:
336 # tried last separator, can't split further, do a fixed-split based on word/character/token
337 fall_back_chunks = self._fall_back_to_fixed_chunking(split_text, self.split_units)
338 chunks.extend(fall_back_chunks)
339 else:
340 chunks.extend(self._chunk_text(split_text))
342 else:
343 current_chunk.append(split_text)
344 current_length += self._chunk_length(split_text)
345 else:
346 current_chunk.append(split_text)
347 current_length += self._chunk_length(split_text)
349 if current_chunk:
350 chunks.append("".join(current_chunk))
352 if self.split_overlap > 0:
353 chunks = self._apply_overlap(chunks)
355 if chunks:
356 return chunks
358 # if no separator worked, fall back to word- or character-level chunking
359 chunks = self._fall_back_to_fixed_chunking(text, self.split_units)
360 if self.split_overlap > 0:
361 chunks = self._apply_overlap(chunks)
362 return chunks
364 def _fall_back_to_fixed_chunking(self, text: str, split_units: Literal["word", "char", "token"]) -> list[str]:
365 """
366 Fall back to a fixed chunking approach if no separator works for the text.
368 Splits the text into smaller chunks based on the split_length and split_units attributes, either by words,
369 characters, or tokens.
371 :param text: The text to be split into chunks.
372 :param split_units: The unit of the split_length parameter. It can be either "word", "char", or "token".
373 :returns:
374 A list of text chunks.
375 """
376 chunks = []
378 if split_units == "word":
379 words = re.findall(r"\S+|\s+", text)
380 current_chunk = []
381 current_length = 0
383 for word in words:
384 # re.findall above also yields multi-character whitespace tokens (" ", "\t").
385 # Only count real words toward the length; any whitespace run is a separator.
386 if word.strip():
387 current_chunk.append(word)
388 current_length += 1
389 if current_length == self.split_length and current_chunk:
390 chunks.append("".join(current_chunk))
391 current_chunk = []
392 current_length = 0
393 else:
394 current_chunk.append(word)
396 if current_length > 0:
397 chunks.append("".join(current_chunk))
398 elif current_chunk:
399 # Only whitespace is left over (e.g. trailing whitespace): attach it to the
400 # previous chunk instead of emitting a whitespace-only chunk on its own.
401 if chunks:
402 chunks[-1] += "".join(current_chunk)
403 else:
404 chunks.append("".join(current_chunk))
405 elif split_units == "char":
406 for i in range(0, self._chunk_length(text), self.split_length):
407 chunks.append(text[i : i + self.split_length])
408 else: # token
409 # at this point we know that the tokenizer is already initialized
410 tokens = self.tiktoken_tokenizer.encode(text) # type: ignore
411 for i in range(0, len(tokens), self.split_length):
412 chunk_tokens = tokens[i : i + self.split_length]
413 chunks.append(self.tiktoken_tokenizer.decode(chunk_tokens)) # type: ignore
414 return chunks
416 def _add_overlap_info(self, curr_pos: int, new_doc: Document, new_docs: list[Document]) -> None:
417 prev_doc = new_docs[-1]
418 # curr_pos and split_idx_start are character offsets, so measure the
419 # overlap and range in characters too (not via _chunk_length, which returns a word/token count).
420 prev_doc_length = len(prev_doc.content) # type: ignore
421 overlap_length = prev_doc_length - (curr_pos - prev_doc.meta["split_idx_start"])
422 if overlap_length > 0:
423 prev_doc.meta["_split_overlap"].append({"doc_id": new_doc.id, "range": (0, overlap_length)})
424 new_doc.meta["_split_overlap"].append(
425 {"doc_id": prev_doc.id, "range": (prev_doc_length - overlap_length, prev_doc_length)}
426 )
428 def _run_one(self, doc: Document) -> list[Document]:
429 chunks = self._chunk_text(doc.content) # type: ignore # the caller already check for a non-empty doc.content
430 chunks = chunks[:-1] if len(chunks[-1]) == 0 else chunks # remove last empty chunk if it exists
431 current_position = 0
432 current_page = 1
434 new_docs: list[Document] = []
436 for split_nr, chunk in enumerate(chunks):
437 meta = deepcopy(doc.meta)
438 meta["source_id"] = doc.id
439 meta["parent_id"] = doc.id
440 meta["split_id"] = split_nr
441 meta["split_idx_start"] = current_position
442 meta["_split_overlap"] = [] if self.split_overlap > 0 else None
443 new_doc = Document(content=chunk, meta=meta)
445 # add overlap information to the previous and current doc
446 if split_nr > 0 and self.split_overlap > 0:
447 self._add_overlap_info(current_position, new_doc, new_docs)
449 # count page breaks in the chunk
450 current_page += chunk.count("\f")
452 # if there are consecutive page breaks at the end with no more text, adjust the page number
453 # e.g: "text\f\f\f" -> 3 page breaks, but current_page should be 1
454 consecutive_page_breaks = len(chunk) - len(chunk.rstrip("\f"))
456 if consecutive_page_breaks > 0:
457 new_doc.meta["page_number"] = current_page - consecutive_page_breaks
458 else:
459 new_doc.meta["page_number"] = current_page
461 # keep the new chunk doc and update the current position
462 new_docs.append(new_doc)
463 # Advance current_position by chunk length minus overlap.
464 # split_overlap is in split_units, not chars, so get the actual
465 # overlap string from _get_overlap() and use its char length.
466 if self.split_overlap > 0 and split_nr < len(chunks) - 1:
467 overlap_str, _ = self._get_overlap([doc.content for doc in new_docs]) # type: ignore[misc]
468 overlap_char_len = len(overlap_str)
469 else:
470 overlap_char_len = 0
471 current_position += len(chunk) - overlap_char_len
473 return new_docs
475 @component.output_types(documents=list[Document])
476 def run(self, documents: list[Document]) -> dict[str, list[Document]]:
477 """
478 Split a list of documents into documents with smaller chunks of text.
480 :param documents: List of Documents to split.
481 :returns:
482 A dictionary containing a key "documents" with a List of Documents with smaller chunks of text corresponding
483 to the input documents.
484 """
485 if not self._is_warmed_up and ("sentence" in self.separators or self.split_units == "token"):
486 self.warm_up()
488 docs = []
489 for doc in documents:
490 if not doc.content or doc.content == "":
491 logger.warning("Document ID {doc_id} has an empty content. Skipping this document.", doc_id=doc.id)
492 continue
493 docs.extend(self._run_one(doc))
495 return {"documents": docs}