Coverage for haystack/components/preprocessors/markdown_header_splitter.py: 97%
173 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 09:42 +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 Literal
9from haystack import Document, component, logging
10from haystack.components.preprocessors import DocumentSplitter
11from haystack.components.preprocessors._page_numbers import _leading_page_breaks
13logger = logging.getLogger(__name__)
16@component
17class MarkdownHeaderSplitter:
18 """
19 Split documents at ATX-style Markdown headers (#), with optional secondary splitting.
21 This component processes text documents by:
22 - Splitting them into chunks at Markdown headers (e.g., '#', '##', etc.), preserving header hierarchy as metadata.
23 - Optionally applying a secondary split (by word, passage, period, or line) to each chunk
24 (using haystack's DocumentSplitter).
25 - Preserving and propagating metadata such as parent headers, page numbers, and split IDs.
26 """
28 def __init__(
29 self,
30 *,
31 page_break_character: str = "\f",
32 keep_headers: bool = True,
33 header_split_levels: list[int] | None = None,
34 secondary_split: Literal["word", "passage", "period", "line"] | None = None,
35 split_length: int = 200,
36 split_overlap: int = 0,
37 split_threshold: int = 0,
38 skip_empty_documents: bool = True,
39 ) -> None:
40 """
41 Initialize the MarkdownHeaderSplitter.
43 :param page_break_character: Character used to identify page breaks. Defaults to form feed ("\f").
44 :param keep_headers: If True, headers are kept in the content. If False, headers are moved to metadata.
45 Defaults to True.
46 :param header_split_levels: List of header levels (1–6) to split on. For example, `[1, 2]` splits only
47 on `#` and `##` headers, merging content under deeper headers into the preceding chunk. Defaults to
48 all levels `[1, 2, 3, 4, 5, 6]`.
49 :param secondary_split: Optional secondary split condition after header splitting.
50 Options are None, "word", "passage", "period", "line". Defaults to None.
51 :param split_length: The maximum number of units in each split when using secondary splitting. Defaults to 200.
52 :param split_overlap: The number of overlapping units for each split when using secondary splitting.
53 Defaults to 0.
54 :param split_threshold: The minimum number of units per split when using secondary splitting. Defaults to 0.
55 :param skip_empty_documents: Choose whether to skip documents with empty content. Default is True.
56 Set to False when downstream components in the Pipeline (like LLMDocumentContentExtractor) can extract text
57 from non-textual documents.
58 """
59 if header_split_levels is None:
60 header_split_levels = [1, 2, 3, 4, 5, 6]
62 if not isinstance(header_split_levels, list) or len(header_split_levels) == 0:
63 raise ValueError("header_split_levels must be a non-empty list.")
64 invalid = [lvl for lvl in header_split_levels if not isinstance(lvl, int) or lvl < 1 or lvl > 6]
65 if invalid:
66 raise ValueError(
67 f"header_split_levels contains invalid values: {invalid}. All levels must be integers between 1 and 6."
68 )
69 if len(header_split_levels) != len(set(header_split_levels)):
70 raise ValueError("header_split_levels must not contain duplicate values.")
72 self.page_break_character = page_break_character
73 self.secondary_split = secondary_split
74 self.split_length = split_length
75 self.split_overlap = split_overlap
76 self.split_threshold = split_threshold
77 self.skip_empty_documents = skip_empty_documents
78 self.keep_headers = keep_headers
79 self.header_split_levels = header_split_levels
80 self._header_split_levels_set = set(header_split_levels)
81 self._header_pattern = re.compile(r"(?m)^(#{1,6}) (.+)$") # ATX-style .md-headers
83 # Matches fenced code blocks delimited by triple backticks (```) or triple tildes (~~~).
84 # Broken down:
85 # ^ - fence must start at the beginning of a line (MULTILINE)
86 # (?P<fence>`{3,}|~{3,})
87 # - named capture group "fence": three or more backticks OR three or
88 # more tildes. Capturing it allows the closing fence to be matched
89 # with a backreference, so ```-opened blocks must close with ```
90 # and ~~~-opened blocks must close with ~~~.
91 # [^\n]* - optional language identifier (e.g. "python") and any other text
92 # on the opening fence line, up to the newline
93 # \n - newline ending the opening fence line
94 # .*? - the code block body, matched lazily (DOTALL so . matches newlines)
95 # ^(?P=fence) - closing fence: must be identical to the opening fence (backreference),
96 # and must start at the beginning of a line
97 # \s*$ - optional trailing whitespace after the closing fence
98 self._code_block_pattern = re.compile(
99 r"^(?P<fence>`{3,}|~{3,})[^\n]*\n.*?^(?P=fence)\s*$", re.MULTILINE | re.DOTALL
100 )
102 self._is_warmed_up = False
104 # initialize secondary_splitter only if needed
105 if self.secondary_split:
106 self.secondary_splitter = DocumentSplitter(
107 split_by=self.secondary_split,
108 split_length=self.split_length,
109 split_overlap=self.split_overlap,
110 split_threshold=self.split_threshold,
111 )
113 def warm_up(self) -> None:
114 """
115 Warm up the MarkdownHeaderSplitter.
116 """
117 if self.secondary_split and not self._is_warmed_up:
118 self.secondary_splitter.warm_up()
119 self._is_warmed_up = True
121 def _code_block_spans(self, text: str) -> list[tuple[int, int]]:
122 """Return the (start, end) character spans of all fenced code blocks in text."""
123 return [(m.start(), m.end()) for m in self._code_block_pattern.finditer(text)]
125 def _split_text_by_markdown_headers(self, text: str, doc_id: str) -> list[dict]:
126 """
127 Split text by ATX-style headers (#) and create chunks with appropriate metadata.
129 The internal `source_start_idx` field tracks where each chunk begins in the original text for page counting.
130 The internal `from_header_split` field records whether a chunk starts with a header line, which decides
131 whether the secondary split may strip that line.
132 """
133 logger.debug("Splitting text by markdown headers")
135 # Pre-compute fenced code block spans so that # lines inside code blocks (e.g. Python comments) are not
136 # mistaken for Markdown headers.
137 code_spans = self._code_block_spans(text)
139 # find headers at the configured levels only, excluding any that fall inside a code block. Content between
140 # skipped headers is absorbed into the preceding chunk's span since end = next_match.start().
141 matches = [
142 m
143 for m in re.finditer(self._header_pattern, text)
144 if len(m.group(1)) in self._header_split_levels_set
145 and not any(start <= m.start() < end for start, end in code_spans)
146 ]
148 # return unsplit if no headers found
149 if not matches:
150 logger.info(
151 "No headers found in document {doc_id}; returning full document as single chunk.", doc_id=doc_id
152 )
153 return [{"content": text, "meta": {}, "source_start_idx": 0, "from_header_split": False}]
155 # process headers and build chunks
156 chunks: list[dict] = []
157 header_stack: list[str | None] = [None] * 6
158 pending_start: int | None = None # start offset in text of the first buffered empty header
159 pending_header_text: str | None = None # text of the last buffered empty header
160 pending_level = 0 # level of the last buffered empty header
161 has_content = False # flag to track if any header has content
163 # Text before the first header belongs to no section, but it is document content in its own
164 # right: a README's opening paragraph, a title block, or front matter. It becomes a chunk with
165 # empty header metadata, which counts as content for the headers-only check below. When it holds
166 # only whitespace it is buffered onto the first chunk that keeps its header line, so the chunks
167 # stay a byte-exact partition of the input without an empty chunk appearing in them.
168 preamble = text[: matches[0].start()]
169 if preamble.strip():
170 chunks.append(
171 {
172 "content": preamble,
173 "meta": {"header": "", "parent_headers": []},
174 "source_start_idx": 0,
175 "from_header_split": False,
176 }
177 )
178 has_content = True
179 elif preamble and self.keep_headers:
180 # Buffering is the mechanism that carries text into a later chunk, and only the
181 # `keep_headers` branch consumes and clears it. With headers in metadata the chunk
182 # content starts after the header line anyway, so there is nothing to carry into.
183 pending_start = 0
185 for i, match in enumerate(matches):
186 # extract header info
187 header_prefix = match.group(1)
188 header_text = match.group(2).strip() # keep surrounding whitespace out of the metadata
189 level = len(header_prefix)
191 # get content
192 start = match.end()
193 end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
194 content = text[start:end]
196 # update header stack to track nesting
197 header_stack[level - 1] = header_text
198 for j in range(level, 6):
199 header_stack[j] = None
201 # skip splits w/o content
202 if not content.strip(): # this strip is needed to avoid counting whitespace as content
203 if self.keep_headers:
204 # buffer this header so it gets prepended to the next contentful chunk; only the
205 # offset of the first header in the run is needed since the chunk is sliced from text
206 if pending_start is None:
207 pending_start = match.start()
208 pending_header_text = header_text
209 pending_level = level
210 continue
212 has_content = True # at least one header has content
213 # Build parent metadata from the current header stack so the first child of a
214 # contentful section still inherits its full ancestor chain.
215 parent_headers = [h for h in header_stack[: level - 1] if h is not None]
217 logger.debug(
218 "Creating chunk for header '{header_text}' at level {level}", header_text=header_text, level=level
219 )
221 if self.keep_headers:
222 # Slice from the first buffered empty header (if any) so the buffered header lines and
223 # the whitespace between them are preserved byte-exactly.
224 source_start_idx = pending_start if pending_start is not None else match.start()
225 chunk_content = text[source_start_idx:end]
226 chunks.append(
227 {
228 "content": chunk_content,
229 "meta": {"header": header_text, "parent_headers": parent_headers},
230 "source_start_idx": source_start_idx,
231 "from_header_split": True,
232 }
233 )
234 pending_start = None # reset buffered headers
235 else:
236 chunks.append(
237 {
238 "content": content,
239 "meta": {"header": header_text, "parent_headers": parent_headers},
240 "source_start_idx": start,
241 "from_header_split": True,
242 }
243 )
245 # return doc unchunked if no headers have content
246 if not has_content:
247 logger.info(
248 "Document {doc_id} contains only headers with no content; returning original document.", doc_id=doc_id
249 )
250 return [{"content": text, "meta": {}, "source_start_idx": 0, "from_header_split": False}]
252 # Flush any trailing headers that had no body text of their own. A header at the very end of the
253 # document (or a run of such headers) never gets a following content chunk to be prepended to, so
254 # without this it would be silently dropped. Slicing the original text keeps reconstruction exact.
255 if pending_start is not None:
256 parent_headers = [h for h in header_stack[: pending_level - 1] if h is not None]
257 chunks.append(
258 {
259 "content": text[pending_start:],
260 "meta": {"header": pending_header_text, "parent_headers": parent_headers},
261 "source_start_idx": pending_start,
262 "from_header_split": True,
263 }
264 )
266 return chunks
268 def _apply_secondary_splitting(self, documents: list[tuple[Document, bool]]) -> list[Document]:
269 """
270 Apply secondary splitting while preserving header metadata and structure.
272 Ensures page counting is maintained across splits.
273 """
274 result_docs = []
275 current_split_id = 0 # track split_id across all secondary splits from the same parent
277 for doc, from_header_split in documents:
278 if doc.content is None:
279 result_docs.append(doc)
280 continue
282 content_for_splitting: str = doc.content
283 secondary_content_start_idx = 0
285 # Only strip a leading header line from chunks that actually start with one.
286 # `from_header_split` is set by the chunk builder in _split_text_by_markdown_headers and carried
287 # through, rather than inferred from the presence of header metadata: the chunk holding the text
288 # before the first header carries header metadata too, and its first line can be a header at a
289 # level that is not being split on. Inferring the flag would strip that line and lose it. The
290 # fallback paths (no headers found / only headers with no content) return the whole document for
291 # the same reason.
292 if not self.keep_headers and from_header_split:
293 header_match = re.match(self._header_pattern, doc.content)
294 if header_match:
295 secondary_content_start_idx = header_match.end()
296 content_for_splitting = doc.content[secondary_content_start_idx:]
298 # The page this header chunk's first character is on; its splits are numbered relative to it.
299 # doc.meta holds the page the chunk's *text* starts on, which already counts the breaks the chunk
300 # opens with, and page_break_ends below counts those same breaks again - so take them back off.
301 chunk_start_page = doc.meta.get("page_number", 1) - _leading_page_breaks(
302 doc.content, self.page_break_character
303 )
305 clean_meta = {k: v for k, v in doc.meta.items() if k != "split_id"}
307 secondary_splits = self.secondary_splitter.run(
308 documents=[Document(content=content_for_splitting, meta=clean_meta)]
309 )["documents"]
311 # Store where each page break ends so splits after it get the next page number.
312 page_break_ends = [match.end() for match in re.finditer(re.escape(self.page_break_character), doc.content)]
313 page_break_count = 0
315 for split in secondary_splits:
316 split_start_idx = secondary_content_start_idx + split.meta["split_idx_start"]
317 while page_break_count < len(page_break_ends) and page_break_ends[page_break_count] <= split_start_idx:
318 page_break_count += 1
320 split.meta["page_number"] = (
321 chunk_start_page
322 + page_break_count
323 + _leading_page_breaks(split.content or "", self.page_break_character)
324 )
325 split.meta["split_id"] = current_split_id
326 if "source_id" in doc.meta:
327 split.meta["source_id"] = doc.meta["source_id"]
328 current_split_id += 1
330 if not self.keep_headers:
331 for key in ["header", "parent_headers"]:
332 if key in doc.meta:
333 split.meta[key] = deepcopy(doc.meta[key])
335 result_docs.append(split)
337 logger.debug(
338 "Secondary splitting complete. Final count: {final_count} documents.", final_count=len(result_docs)
339 )
340 return result_docs
342 def _split_documents_by_markdown_headers(self, documents: list[Document]) -> list[tuple[Document, bool]]:
343 """
344 Split a list of documents by markdown headers, preserving metadata.
346 Each returned Document is paired with whether it was produced by an actual header split (True) or is a
347 whole, unsplit document returned by a fallback path of _split_text_by_markdown_headers (False).
348 """
350 result_docs: list[tuple[Document, bool]] = []
351 for doc in documents:
352 logger.debug("Splitting document with id={doc_id}", doc_id=doc.id)
353 # mypy: doc.content is Optional[str], so we must check for None before passing to splitting method
354 if doc.content is None:
355 continue
356 splits = self._split_text_by_markdown_headers(doc.content, doc.id)
357 docs: list[tuple[Document, bool]] = []
359 document_start_page = doc.meta.get("page_number", 1) if doc.meta else 1
360 total_page_breaks = doc.content.count(self.page_break_character)
361 logger.debug(
362 "Processing document with id={doc_id}: starting at page {start_page}, "
363 "contains {page_breaks} page breaks in total",
364 doc_id=doc.id,
365 start_page=document_start_page,
366 page_breaks=total_page_breaks,
367 )
368 for split_idx, split in enumerate(splits):
369 meta = deepcopy(doc.meta) if doc.meta else {}
370 chunk_start_page = (
371 document_start_page
372 + doc.content.count(self.page_break_character, 0, split["source_start_idx"])
373 # breaks the chunk opens with end the previous page, so its text is on a later one
374 + _leading_page_breaks(split["content"], self.page_break_character)
375 )
376 meta.update({"source_id": doc.id, "page_number": chunk_start_page, "split_id": split_idx})
377 from_header_split = split["from_header_split"]
378 if split["meta"]:
379 meta.update(split["meta"])
380 docs.append((Document(content=split["content"], meta=meta), from_header_split))
381 final_page = document_start_page + total_page_breaks
382 logger.debug(
383 "Split into {num_docs} documents for id={doc_id}, final page: {final_page}",
384 num_docs=len(docs),
385 doc_id=doc.id,
386 final_page=final_page,
387 )
388 result_docs.extend(docs)
389 return result_docs
391 @component.output_types(documents=list[Document])
392 def run(self, documents: list[Document]) -> dict[str, list[Document]]:
393 """
394 Run the markdown header splitter with optional secondary splitting.
396 :param documents: List of documents to split
398 :returns: A dictionary with the following key:
399 - `documents`: List of documents with the split texts. Each document includes:
400 - A metadata field `source_id` to track the original document.
401 - A metadata field `page_number` with the page the chunk starts on, counting
402 `page_break_character` occurrences in the original document.
403 - A metadata field `split_id` to identify the split chunk index within its parent document.
404 - All other metadata copied from the original document.
405 :raises ValueError: If a document has `None` content.
406 :raises TypeError: If a document's content is not a string.
407 """
408 if self.secondary_split and not self._is_warmed_up:
409 self.warm_up()
410 # validate input documents
411 for doc in documents:
412 if doc.content is None:
413 raise ValueError(
414 "MarkdownHeaderSplitter only works with text documents but content for document ID"
415 f" {doc.id} is None."
416 )
417 if not isinstance(doc.content, str):
418 raise TypeError("MarkdownHeaderSplitter only works with text documents (str content).")
420 final_docs = []
421 for doc in documents:
422 # handle empty documents
423 if not doc.content or not doc.content.strip(): # avoid counting whitespace as content
424 if self.skip_empty_documents:
425 logger.warning("Document ID {doc_id} has an empty content. Skipping this document.", doc_id=doc.id)
426 continue
427 # keep empty documents
428 final_docs.append(doc)
429 logger.warning(
430 "Document ID {doc_id} has an empty content. Keeping this document as per configuration.",
431 doc_id=doc.id,
432 )
433 continue
435 # split this document by headers
436 header_split_docs = self._split_documents_by_markdown_headers([doc])
438 # apply secondary splitting if configured
439 if self.secondary_split:
440 doc_splits = self._apply_secondary_splitting(header_split_docs)
441 else:
442 doc_splits = [split_doc for split_doc, _ in header_split_docs]
444 final_docs.extend(doc_splits)
446 return {"documents": final_docs}