Coverage for haystack/components/preprocessors/markdown_header_splitter.py: 96%
166 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 Literal
9from haystack import Document, component, logging
10from haystack.components.preprocessors import DocumentSplitter
12logger = logging.getLogger(__name__)
15@component
16class MarkdownHeaderSplitter:
17 """
18 Split documents at ATX-style Markdown headers (#), with optional secondary splitting.
20 This component processes text documents by:
21 - Splitting them into chunks at Markdown headers (e.g., '#', '##', etc.), preserving header hierarchy as metadata.
22 - Optionally applying a secondary split (by word, passage, period, or line) to each chunk
23 (using haystack's DocumentSplitter).
24 - Preserving and propagating metadata such as parent headers, page numbers, and split IDs.
25 """
27 def __init__(
28 self,
29 *,
30 page_break_character: str = "\f",
31 keep_headers: bool = True,
32 header_split_levels: list[int] | None = None,
33 secondary_split: Literal["word", "passage", "period", "line"] | None = None,
34 split_length: int = 200,
35 split_overlap: int = 0,
36 split_threshold: int = 0,
37 skip_empty_documents: bool = True,
38 ) -> None:
39 """
40 Initialize the MarkdownHeaderSplitter.
42 :param page_break_character: Character used to identify page breaks. Defaults to form feed ("\f").
43 :param keep_headers: If True, headers are kept in the content. If False, headers are moved to metadata.
44 Defaults to True.
45 :param header_split_levels: List of header levels (1–6) to split on. For example, `[1, 2]` splits only
46 on `#` and `##` headers, merging content under deeper headers into the preceding chunk. Defaults to
47 all levels `[1, 2, 3, 4, 5, 6]`.
48 :param secondary_split: Optional secondary split condition after header splitting.
49 Options are None, "word", "passage", "period", "line". Defaults to None.
50 :param split_length: The maximum number of units in each split when using secondary splitting. Defaults to 200.
51 :param split_overlap: The number of overlapping units for each split when using secondary splitting.
52 Defaults to 0.
53 :param split_threshold: The minimum number of units per split when using secondary splitting. Defaults to 0.
54 :param skip_empty_documents: Choose whether to skip documents with empty content. Default is True.
55 Set to False when downstream components in the Pipeline (like LLMDocumentContentExtractor) can extract text
56 from non-textual documents.
57 """
58 if header_split_levels is None:
59 header_split_levels = [1, 2, 3, 4, 5, 6]
61 if not isinstance(header_split_levels, list) or len(header_split_levels) == 0:
62 raise ValueError("header_split_levels must be a non-empty list.")
63 invalid = [lvl for lvl in header_split_levels if not isinstance(lvl, int) or lvl < 1 or lvl > 6]
64 if invalid:
65 raise ValueError(
66 f"header_split_levels contains invalid values: {invalid}. All levels must be integers between 1 and 6."
67 )
68 if len(header_split_levels) != len(set(header_split_levels)):
69 raise ValueError("header_split_levels must not contain duplicate values.")
71 self.page_break_character = page_break_character
72 self.secondary_split = secondary_split
73 self.split_length = split_length
74 self.split_overlap = split_overlap
75 self.split_threshold = split_threshold
76 self.skip_empty_documents = skip_empty_documents
77 self.keep_headers = keep_headers
78 self.header_split_levels = header_split_levels
79 self._header_split_levels_set = set(header_split_levels)
80 self._header_pattern = re.compile(r"(?m)^(#{1,6}) (.+)$") # ATX-style .md-headers
82 # Matches fenced code blocks delimited by triple backticks (```) or triple tildes (~~~).
83 # Broken down:
84 # ^ - fence must start at the beginning of a line (MULTILINE)
85 # (?P<fence>`{3,}|~{3,})
86 # - named capture group "fence": three or more backticks OR three or
87 # more tildes. Capturing it allows the closing fence to be matched
88 # with a backreference, so ```-opened blocks must close with ```
89 # and ~~~-opened blocks must close with ~~~.
90 # [^\n]* - optional language identifier (e.g. "python") and any other text
91 # on the opening fence line, up to the newline
92 # \n - newline ending the opening fence line
93 # .*? - the code block body, matched lazily (DOTALL so . matches newlines)
94 # ^(?P=fence) - closing fence: must be identical to the opening fence (backreference),
95 # and must start at the beginning of a line
96 # \s*$ - optional trailing whitespace after the closing fence
97 self._code_block_pattern = re.compile(
98 r"^(?P<fence>`{3,}|~{3,})[^\n]*\n.*?^(?P=fence)\s*$", re.MULTILINE | re.DOTALL
99 )
101 self._is_warmed_up = False
103 # initialize secondary_splitter only if needed
104 if self.secondary_split:
105 self.secondary_splitter = DocumentSplitter(
106 split_by=self.secondary_split,
107 split_length=self.split_length,
108 split_overlap=self.split_overlap,
109 split_threshold=self.split_threshold,
110 )
112 def warm_up(self) -> None:
113 """
114 Warm up the MarkdownHeaderSplitter.
115 """
116 if self.secondary_split and not self._is_warmed_up:
117 self.secondary_splitter.warm_up()
118 self._is_warmed_up = True
120 def _code_block_spans(self, text: str) -> list[tuple[int, int]]:
121 """Return the (start, end) character spans of all fenced code blocks in text."""
122 return [(m.start(), m.end()) for m in self._code_block_pattern.finditer(text)]
124 def _split_text_by_markdown_headers(self, text: str, doc_id: str) -> list[dict]:
125 """Split text by ATX-style headers (#) and create chunks with appropriate metadata."""
126 logger.debug("Splitting text by markdown headers")
128 # Pre-compute fenced code block spans so that # lines inside code blocks (e.g. Python comments) are not
129 # mistaken for Markdown headers.
130 code_spans = self._code_block_spans(text)
132 # find headers at the configured levels only, excluding any that fall inside a code block. Content between
133 # skipped headers is absorbed into the preceding chunk's span since end = next_match.start().
134 matches = [
135 m
136 for m in re.finditer(self._header_pattern, text)
137 if len(m.group(1)) in self._header_split_levels_set
138 and not any(start <= m.start() < end for start, end in code_spans)
139 ]
141 # return unsplit if no headers found
142 if not matches:
143 logger.info(
144 "No headers found in document {doc_id}; returning full document as single chunk.", doc_id=doc_id
145 )
146 return [{"content": text, "meta": {}}]
148 # process headers and build chunks
149 chunks: list[dict] = []
150 header_stack: list[str | None] = [None] * 6
151 pending_start: int | None = None # start offset in text of the first buffered empty header
152 pending_header_text: str | None = None # text of the last buffered empty header
153 pending_level = 0 # level of the last buffered empty header
154 has_content = False # flag to track if any header has content
156 for i, match in enumerate(matches):
157 # extract header info
158 header_prefix = match.group(1)
159 header_text = match.group(2).strip() # keep surrounding whitespace out of the metadata
160 level = len(header_prefix)
162 # get content
163 start = match.end()
164 end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
165 content = text[start:end]
167 # update header stack to track nesting
168 header_stack[level - 1] = header_text
169 for j in range(level, 6):
170 header_stack[j] = None
172 # skip splits w/o content
173 if not content.strip(): # this strip is needed to avoid counting whitespace as content
174 if self.keep_headers:
175 # buffer this header so it gets prepended to the next contentful chunk; only the
176 # offset of the first header in the run is needed since the chunk is sliced from text
177 if pending_start is None:
178 pending_start = match.start()
179 pending_header_text = header_text
180 pending_level = level
181 continue
183 has_content = True # at least one header has content
184 # Build parent metadata from the current header stack so the first child of a
185 # contentful section still inherits its full ancestor chain.
186 parent_headers = [h for h in header_stack[: level - 1] if h is not None]
188 logger.debug(
189 "Creating chunk for header '{header_text}' at level {level}", header_text=header_text, level=level
190 )
192 if self.keep_headers:
193 # Slice from the first buffered empty header (if any) so the buffered header lines and
194 # the whitespace between them are preserved byte-exactly.
195 chunk_content = text[(pending_start if pending_start is not None else match.start()) : end]
196 chunks.append(
197 {"content": chunk_content, "meta": {"header": header_text, "parent_headers": parent_headers}}
198 )
199 pending_start = None # reset buffered headers
200 else:
201 chunks.append({"content": content, "meta": {"header": header_text, "parent_headers": parent_headers}})
203 # return doc unchunked if no headers have content
204 if not has_content:
205 logger.info(
206 "Document {doc_id} contains only headers with no content; returning original document.", doc_id=doc_id
207 )
208 return [{"content": text, "meta": {}}]
210 # Flush any trailing headers that had no body text of their own. A header at the very end of the
211 # document (or a run of such headers) never gets a following content chunk to be prepended to, so
212 # without this it would be silently dropped. Slicing the original text keeps reconstruction exact.
213 if pending_start is not None:
214 parent_headers = [h for h in header_stack[: pending_level - 1] if h is not None]
215 chunks.append(
216 {
217 "content": text[pending_start:],
218 "meta": {"header": pending_header_text, "parent_headers": parent_headers},
219 }
220 )
222 return chunks
224 def _apply_secondary_splitting(self, documents: list[Document]) -> list[Document]:
225 """
226 Apply secondary splitting while preserving header metadata and structure.
228 Ensures page counting is maintained across splits.
229 """
230 result_docs = []
231 current_split_id = 0 # track split_id across all secondary splits from the same parent
233 for doc in documents:
234 if doc.content is None:
235 result_docs.append(doc)
236 continue
238 content_for_splitting: str = doc.content
240 if not self.keep_headers: # skip header extraction if keep_headers
241 # extract header information
242 header_match = re.match(self._header_pattern, doc.content)
243 if header_match:
244 content_for_splitting = doc.content[header_match.end() :]
246 # track page from meta
247 current_page = doc.meta.get("page_number", 1)
249 # create a clean meta dict without split_id for secondary splitting
250 clean_meta = {k: v for k, v in doc.meta.items() if k != "split_id"}
252 secondary_splits = self.secondary_splitter.run(
253 documents=[Document(content=content_for_splitting, meta=clean_meta)]
254 )["documents"]
256 # split processing
257 for i, split in enumerate(secondary_splits):
258 # calculate page number for this split
259 if i > 0 and secondary_splits[i - 1].content:
260 current_page = self._update_page_number_with_breaks(secondary_splits[i - 1].content, current_page)
262 # set page number and split_id to meta
263 split.meta["page_number"] = current_page
264 split.meta["split_id"] = current_split_id
265 # ensure source_id is preserved from the original document
266 if "source_id" in doc.meta:
267 split.meta["source_id"] = doc.meta["source_id"]
268 current_split_id += 1
270 # preserve header metadata if we're not keeping headers in content
271 if not self.keep_headers:
272 for key in ["header", "parent_headers"]:
273 if key in doc.meta:
274 split.meta[key] = doc.meta[key]
276 result_docs.append(split)
278 logger.debug(
279 "Secondary splitting complete. Final count: {final_count} documents.", final_count=len(result_docs)
280 )
281 return result_docs
283 def _update_page_number_with_breaks(self, content: str | None, current_page: int) -> int:
284 """
285 Update page number based on page breaks in content.
287 :param content: Content to check for page breaks
288 :param current_page: Current page number
289 :return: New current page number
290 """
291 if not isinstance(content, str):
292 return current_page
294 page_breaks = content.count(self.page_break_character)
295 new_page_number = current_page + page_breaks
297 if page_breaks > 0:
298 logger.debug(
299 "Found {page_breaks} page breaks, page number updated: {old} → {new}",
300 page_breaks=page_breaks,
301 old=current_page,
302 new=new_page_number,
303 )
305 return new_page_number
307 def _split_documents_by_markdown_headers(self, documents: list[Document]) -> list[Document]:
308 """Split a list of documents by markdown headers, preserving metadata."""
310 result_docs = []
311 for doc in documents:
312 logger.debug("Splitting document with id={doc_id}", doc_id=doc.id)
313 # mypy: doc.content is Optional[str], so we must check for None before passing to splitting method
314 if doc.content is None:
315 continue
316 splits = self._split_text_by_markdown_headers(doc.content, doc.id)
317 docs = []
319 current_page = doc.meta.get("page_number", 1) if doc.meta else 1
320 total_page_breaks = doc.content.count(self.page_break_character)
321 logger.debug(
322 "Processing document with id={doc_id}: starting at page {start_page}, "
323 "contains {page_breaks} page breaks in total",
324 doc_id=doc.id,
325 start_page=current_page,
326 page_breaks=total_page_breaks,
327 )
328 for split_idx, split in enumerate(splits):
329 meta = deepcopy(doc.meta) if doc.meta else {}
330 meta.update({"source_id": doc.id, "page_number": current_page, "split_id": split_idx})
331 if split.get("meta"):
332 meta.update(split["meta"])
333 current_page = self._update_page_number_with_breaks(split["content"], current_page)
334 docs.append(Document(content=split["content"], meta=meta))
335 logger.debug(
336 "Split into {num_docs} documents for id={doc_id}, final page: {current_page}",
337 num_docs=len(docs),
338 doc_id=doc.id,
339 current_page=current_page,
340 )
341 result_docs.extend(docs)
342 return result_docs
344 @component.output_types(documents=list[Document])
345 def run(self, documents: list[Document]) -> dict[str, list[Document]]:
346 """
347 Run the markdown header splitter with optional secondary splitting.
349 :param documents: List of documents to split
351 :returns: A dictionary with the following key:
352 - `documents`: List of documents with the split texts. Each document includes:
353 - A metadata field `source_id` to track the original document.
354 - A metadata field `page_number` to track the original page number.
355 - A metadata field `split_id` to identify the split chunk index within its parent document.
356 - All other metadata copied from the original document.
357 :raises ValueError: If a document has `None` content.
358 :raises TypeError: If a document's content is not a string.
359 """
360 if self.secondary_split and not self._is_warmed_up:
361 self.warm_up()
362 # validate input documents
363 for doc in documents:
364 if doc.content is None:
365 raise ValueError(
366 "MarkdownHeaderSplitter only works with text documents but content for document ID"
367 f" {doc.id} is None."
368 )
369 if not isinstance(doc.content, str):
370 raise TypeError("MarkdownHeaderSplitter only works with text documents (str content).")
372 final_docs = []
373 for doc in documents:
374 # handle empty documents
375 if not doc.content or not doc.content.strip(): # avoid counting whitespace as content
376 if self.skip_empty_documents:
377 logger.warning("Document ID {doc_id} has an empty content. Skipping this document.", doc_id=doc.id)
378 continue
379 # keep empty documents
380 final_docs.append(doc)
381 logger.warning(
382 "Document ID {doc_id} has an empty content. Keeping this document as per configuration.",
383 doc_id=doc.id,
384 )
385 continue
387 # split this document by headers
388 header_split_docs = self._split_documents_by_markdown_headers([doc])
390 # apply secondary splitting if configured
391 if self.secondary_split:
392 doc_splits = self._apply_secondary_splitting(header_split_docs)
393 else:
394 doc_splits = header_split_docs
396 final_docs.extend(doc_splits)
398 return {"documents": final_docs}