Skip to content

Commit 589e56a

Browse files
OGuggenbuehlsjrl
authored and
kacperlukawski
committed
feat: MarkdownHeaderSplitter (#9660)
* implement md-header-splitter and add tests * rework md-header splitter to rewrite md-header levels * remove deprecated test * Update haystack/components/preprocessors/markdown_header_splitter.py use haystack logging Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> * use native types * move to haystack logging * docstrings improvements * Update haystack/components/preprocessors/markdown_header_splitter.py remove temp toc Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> * fix CustomDocumentSplitter arguments * remove header prefix from content * rework split_id assignment to avoid collisions * remove unneeded dese methods * cleanup * cleanup * add tests cleanup * move initialization of secondary-splitter out of run method * move _custom_document_splitter to class method * removed the _CustomDocumentSplitter class. splitting logic is now encapsulated within the MarkdownHeaderSplitter class as private methods. * return to standard feed-forward character and add tests for page break handling * quit exposing splitting_function param since it shouldn't be changed anyway * remove test section in module * add license header * add release note * minor refactor for type safety * Update haystack/components/preprocessors/markdown_header_splitter.py Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> * remove unneeded release notes entries * improved documentation for methods * improve method naming * improved page-number assignment & added return in docstring minor cleanup * unified page-counting * simplify conditional secondary-split initialization and usage * fix linting error * clearly specify the use of ATX-style headers (#) only * reference doc_id when logging no headers found * initialize md-header pattern as private variable once * add example to for inferring header levels to docstring * improve empty document handling add more logging for empty documents * more explicit testing for inferred headers * fix linting issue * improved empty content handling test cases * remove all functionality related to inferring md-header levels * compile regex-pattern in init for performance gains * Update haystack/components/preprocessors/markdown_header_splitter.py Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> * change all "none" to proper None values * fix minor * explicitly test doc content * rename parentheaders to parent_headers * test split_id, doc length * check meta content * remove unneeded test * make split_id testing more robust * more realistic overlap test sample * assign split_id globally to all output docs * taste page numbers explicitly * cleanup pagebreak test * minor * return doc unchunked if no headers have content * add doc-id to logging statement for unsplit documents * remove unneeded logs * minor cleanup * simplify page-number tracking method to not return count, just the updated page number * add dev comment to mypy check for doc.content is None * Update haystack/components/preprocessors/markdown_header_splitter.py Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> * remove split meta flattening * keep empty meta return consistent * remove unneeded content is none check * update tests to reflect empty meta dict for unsplit docs * clean up total_page counts * remove unneeded meta check * Update test/components/preprocessors/test_markdown_header_splitter.py Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> * implement keep_headers parameter * remove meta-dict flattening * add minor sanity checks * Update test/components/preprocessors/test_markdown_header_splitter.py Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> * add warmup * Update haystack/components/preprocessors/markdown_header_splitter.py Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> * fix splitting when keeping headers * test cleanup to cover keep_headers=True * add tests for keep_headers=False splitting * remove strip() * simplify doc handling * fix split id assignment * test cleanup * test splits more explicitly * cleanup tests minor commenting * Update haystack/components/preprocessors/markdown_header_splitter.py Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> * Update haystack/components/preprocessors/markdown_header_splitter.py Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> * Update haystack/components/preprocessors/markdown_header_splitter.py Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> * Update haystack/components/preprocessors/markdown_header_splitter.py Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> * Update test/components/preprocessors/test_markdown_header_splitter.py Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> * Update haystack/components/preprocessors/markdown_header_splitter.py Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> * Update releasenotes/notes/add-md-header-splitter-df5c024a6ddd2718.yaml Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> * Update haystack/components/preprocessors/markdown_header_splitter.py Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> * Update releasenotes/notes/add-md-header-splitter-df5c024a6ddd2718.yaml Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> * Update releasenotes/notes/add-md-header-splitter-df5c024a6ddd2718.yaml Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> * Update haystack/components/preprocessors/markdown_header_splitter.py Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> * fix test now that meta is always kept regardless of headers * update tests to consider headers always part of meta * remove trailing whitespace removal * remove redundant test * make test_split_multiple_documents more explicit * make tests more explicit * remove header level inference from release notes * improve splitting log message * add split ids and more explicit string assertions * add fixture & appropriate assertions for sample text with page breaks * add md-header-splitter to preprocessors api --------- Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com>
1 parent 554267b commit 589e56a

4 files changed

Lines changed: 915 additions & 0 deletions

File tree

‎haystack/components/preprocessors/markdown_header_splitter.py‎

Lines changed: 334 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,334 @@
1+
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2+
#
3+
# SPDX-License-Identifier: Apache-2.0
4+
5+
import re
6+
from typing import Literal
7+
8+
from haystack import Document, component, logging
9+
from haystack.components.preprocessors import DocumentSplitter
10+
11+
logger = logging.getLogger(__name__)
12+
13+
14+
@component
15+
class MarkdownHeaderSplitter:
16+
"""
17+
Split documents at ATX-style Markdown headers (#), with optional secondary splitting.
18+
19+
This component processes text documents by:
20+
- Splitting them into chunks at Markdown headers (e.g., '#', '##', etc.), preserving header hierarchy as metadata.
21+
- Optionally applying a secondary split (by word, passage, period, or line) to each chunk
22+
(using haystack's DocumentSplitter).
23+
- Preserving and propagating metadata such as parent headers, page numbers, and split IDs.
24+
"""
25+
26+
def __init__(
27+
self,
28+
*,
29+
page_break_character: str = "\f",
30+
keep_headers: bool = True,
31+
secondary_split: Literal["word", "passage", "period", "line"] | None = None,
32+
split_length: int = 200,
33+
split_overlap: int = 0,
34+
split_threshold: int = 0,
35+
skip_empty_documents: bool = True,
36+
):
37+
"""
38+
Initialize the MarkdownHeaderSplitter.
39+
40+
:param page_break_character: Character used to identify page breaks. Defaults to form feed ("\f").
41+
:param keep_headers: If True, headers are kept in the content. If False, headers are moved to metadata.
42+
Defaults to True.
43+
:param secondary_split: Optional secondary split condition after header splitting.
44+
Options are None, "word", "passage", "period", "line". Defaults to None.
45+
:param split_length: The maximum number of units in each split when using secondary splitting. Defaults to 200.
46+
:param split_overlap: The number of overlapping units for each split when using secondary splitting.
47+
Defaults to 0.
48+
:param split_threshold: The minimum number of units per split when using secondary splitting. Defaults to 0.
49+
:param skip_empty_documents: Choose whether to skip documents with empty content. Default is True.
50+
Set to False when downstream components in the Pipeline (like LLMDocumentContentExtractor) can extract text
51+
from non-textual documents.
52+
"""
53+
self.page_break_character = page_break_character
54+
self.secondary_split = secondary_split
55+
self.split_length = split_length
56+
self.split_overlap = split_overlap
57+
self.split_threshold = split_threshold
58+
self.skip_empty_documents = skip_empty_documents
59+
self.keep_headers = keep_headers
60+
self._header_pattern = re.compile(r"(?m)^(#{1,6}) (.+)$") # ATX-style .md-headers
61+
self._is_warmed_up = False
62+
63+
# initialize secondary_splitter only if needed
64+
if self.secondary_split:
65+
self.secondary_splitter = DocumentSplitter(
66+
split_by=self.secondary_split,
67+
split_length=self.split_length,
68+
split_overlap=self.split_overlap,
69+
split_threshold=self.split_threshold,
70+
)
71+
72+
def warm_up(self):
73+
"""
74+
Warm up the MarkdownHeaderSplitter.
75+
"""
76+
if self.secondary_split and not self._is_warmed_up:
77+
self.secondary_splitter.warm_up()
78+
self._is_warmed_up = True
79+
80+
def _split_text_by_markdown_headers(self, text: str, doc_id: str) -> list[dict]:
81+
"""Split text by ATX-style headers (#) and create chunks with appropriate metadata."""
82+
logger.debug("Splitting text by markdown headers")
83+
84+
# find headers
85+
matches = list(re.finditer(self._header_pattern, text))
86+
87+
# return unsplit if no headers found
88+
if not matches:
89+
logger.info(
90+
"No headers found in document {doc_id}; returning full document as single chunk.", doc_id=doc_id
91+
)
92+
return [{"content": text, "meta": {}}]
93+
94+
# process headers and build chunks
95+
chunks: list[dict] = []
96+
header_stack: list[str | None] = [None] * 6
97+
active_parents: list[str] = [] # track active parent headers
98+
pending_headers: list[str] = [] # store empty headers to prepend to next content
99+
has_content = False # flag to track if any header has content
100+
101+
for i, match in enumerate(matches):
102+
# extract header info
103+
header_prefix = match.group(1)
104+
header_text = match.group(2)
105+
level = len(header_prefix)
106+
107+
# get content
108+
start = match.end()
109+
end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
110+
content = text[start:end]
111+
112+
# update header stack to track nesting
113+
header_stack[level - 1] = header_text
114+
for j in range(level, 6):
115+
header_stack[j] = None
116+
117+
# skip splits w/o content
118+
if not content.strip(): # this strip is needed to avoid counting whitespace as content
119+
# add as parent for subsequent headers
120+
active_parents = [h for h in header_stack[: level - 1] if h is not None]
121+
active_parents.append(header_text)
122+
if self.keep_headers:
123+
header_line = f"{header_prefix} {header_text}"
124+
pending_headers.append(header_line)
125+
continue
126+
127+
has_content = True # at least one header has content
128+
parent_headers = list(active_parents)
129+
130+
logger.debug(
131+
"Creating chunk for header '{header_text}' at level {level}", header_text=header_text, level=level
132+
)
133+
134+
if self.keep_headers:
135+
header_line = f"{header_prefix} {header_text}"
136+
# add pending & current header to content
137+
chunk_content = ""
138+
if pending_headers:
139+
chunk_content += "\n".join(pending_headers) + "\n"
140+
chunk_content += f"{header_line}{content}"
141+
chunks.append(
142+
{"content": chunk_content, "meta": {"header": header_text, "parent_headers": parent_headers}}
143+
)
144+
pending_headers = [] # reset pending headers
145+
else:
146+
chunks.append({"content": content, "meta": {"header": header_text, "parent_headers": parent_headers}})
147+
148+
# reset active parents
149+
active_parents = [h for h in header_stack[: level - 1] if h is not None]
150+
151+
# return doc unchunked if no headers have content
152+
if not has_content:
153+
logger.info(
154+
"Document {doc_id} contains only headers with no content; returning original document.", doc_id=doc_id
155+
)
156+
return [{"content": text, "meta": {}}]
157+
158+
return chunks
159+
160+
def _apply_secondary_splitting(self, documents: list[Document]) -> list[Document]:
161+
"""
162+
Apply secondary splitting while preserving header metadata and structure.
163+
164+
Ensures page counting is maintained across splits.
165+
"""
166+
result_docs = []
167+
current_split_id = 0 # track split_id across all secondary splits from the same parent
168+
169+
for doc in documents:
170+
if doc.content is None:
171+
result_docs.append(doc)
172+
continue
173+
174+
content_for_splitting: str = doc.content
175+
176+
if not self.keep_headers: # skip header extraction if keep_headers
177+
# extract header information
178+
header_match = re.search(self._header_pattern, doc.content)
179+
if header_match:
180+
content_for_splitting = doc.content[header_match.end() :]
181+
182+
# track page from meta
183+
current_page = doc.meta.get("page_number", 1)
184+
185+
# create a clean meta dict without split_id for secondary splitting
186+
clean_meta = {k: v for k, v in doc.meta.items() if k != "split_id"}
187+
188+
secondary_splits = self.secondary_splitter.run(
189+
documents=[Document(content=content_for_splitting, meta=clean_meta)]
190+
)["documents"]
191+
192+
# split processing
193+
for i, split in enumerate(secondary_splits):
194+
# calculate page number for this split
195+
if i > 0 and secondary_splits[i - 1].content:
196+
current_page = self._update_page_number_with_breaks(secondary_splits[i - 1].content, current_page)
197+
198+
# set page number and split_id to meta
199+
split.meta["page_number"] = current_page
200+
split.meta["split_id"] = current_split_id
201+
# ensure source_id is preserved from the original document
202+
if "source_id" in doc.meta:
203+
split.meta["source_id"] = doc.meta["source_id"]
204+
current_split_id += 1
205+
206+
# preserve header metadata if we're not keeping headers in content
207+
if not self.keep_headers:
208+
for key in ["header", "parent_headers"]:
209+
if key in doc.meta:
210+
split.meta[key] = doc.meta[key]
211+
212+
result_docs.append(split)
213+
214+
logger.debug(
215+
"Secondary splitting complete. Final count: {final_count} documents.", final_count=len(result_docs)
216+
)
217+
return result_docs
218+
219+
def _update_page_number_with_breaks(self, content: str, current_page: int) -> int:
220+
"""
221+
Update page number based on page breaks in content.
222+
223+
:param content: Content to check for page breaks
224+
:param current_page: Current page number
225+
:return: New current page number
226+
"""
227+
if not isinstance(content, str):
228+
return current_page
229+
230+
page_breaks = content.count(self.page_break_character)
231+
new_page_number = current_page + page_breaks
232+
233+
if page_breaks > 0:
234+
logger.debug(
235+
"Found {page_breaks} page breaks, page number updated: {old} → {new}",
236+
page_breaks=page_breaks,
237+
old=current_page,
238+
new=new_page_number,
239+
)
240+
241+
return new_page_number
242+
243+
def _split_documents_by_markdown_headers(self, documents: list[Document]) -> list[Document]:
244+
"""Split a list of documents by markdown headers, preserving metadata."""
245+
246+
result_docs = []
247+
for doc in documents:
248+
logger.debug("Splitting document with id={doc_id}", doc_id=doc.id)
249+
# mypy: doc.content is Optional[str], so we must check for None before passing to splitting method
250+
if doc.content is None:
251+
continue
252+
splits = self._split_text_by_markdown_headers(doc.content, doc.id)
253+
docs = []
254+
255+
current_page = doc.meta.get("page_number", 1) if doc.meta else 1
256+
total_page_breaks = doc.content.count(self.page_break_character)
257+
logger.debug(
258+
"Processing document with id={doc_id}: starting at page {start_page}, "
259+
"contains {page_breaks} page breaks in total",
260+
doc_id=doc.id,
261+
start_page=current_page,
262+
page_breaks=total_page_breaks,
263+
)
264+
for split_idx, split in enumerate(splits):
265+
meta = doc.meta.copy() if doc.meta else {}
266+
meta.update({"source_id": doc.id, "page_number": current_page, "split_id": split_idx})
267+
if split.get("meta"):
268+
meta.update(split["meta"])
269+
current_page = self._update_page_number_with_breaks(split["content"], current_page)
270+
docs.append(Document(content=split["content"], meta=meta))
271+
logger.debug(
272+
"Split into {num_docs} documents for id={doc_id}, final page: {current_page}",
273+
num_docs=len(docs),
274+
doc_id=doc.id,
275+
current_page=current_page,
276+
)
277+
result_docs.extend(docs)
278+
return result_docs
279+
280+
@component.output_types(documents=list[Document])
281+
def run(self, documents: list[Document]) -> dict[str, list[Document]]:
282+
"""
283+
Run the markdown header splitter with optional secondary splitting.
284+
285+
:param documents: List of documents to split
286+
287+
:returns: A dictionary with the following key:
288+
- `documents`: List of documents with the split texts. Each document includes:
289+
- A metadata field `source_id` to track the original document.
290+
- A metadata field `page_number` to track the original page number.
291+
- A metadata field `split_id` to identify the split chunk index within its parent document.
292+
- All other metadata copied from the original document.
293+
"""
294+
if self.secondary_split and not self._is_warmed_up:
295+
self.warm_up()
296+
# validate input documents
297+
for doc in documents:
298+
if doc.content is None:
299+
raise ValueError(
300+
(
301+
"MarkdownHeaderSplitter only works with text documents but content for document ID"
302+
f" {doc.id} is None."
303+
)
304+
)
305+
if not isinstance(doc.content, str):
306+
raise ValueError("MarkdownHeaderSplitter only works with text documents (str content).")
307+
308+
final_docs = []
309+
for doc in documents:
310+
# handle empty documents
311+
if not doc.content or not doc.content.strip(): # avoid counting whitespace as content
312+
if self.skip_empty_documents:
313+
logger.warning("Document ID {doc_id} has an empty content. Skipping this document.", doc_id=doc.id)
314+
continue
315+
# keep empty documents
316+
final_docs.append(doc)
317+
logger.warning(
318+
"Document ID {doc_id} has an empty content. Keeping this document as per configuration.",
319+
doc_id=doc.id,
320+
)
321+
continue
322+
323+
# split this document by headers
324+
header_split_docs = self._split_documents_by_markdown_headers([doc])
325+
326+
# apply secondary splitting if configured
327+
if self.secondary_split:
328+
doc_splits = self._apply_secondary_splitting(header_split_docs)
329+
else:
330+
doc_splits = header_split_docs
331+
332+
final_docs.extend(doc_splits)
333+
334+
return {"documents": final_docs}

‎pydoc/preprocessors_api.yml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ loaders:
99
"document_splitter",
1010
"embedding_based_document_splitter",
1111
"hierarchical_document_splitter",
12+
"markdown_header_splitter",
1213
"recursive_splitter",
1314
"text_cleaner"]
1415
ignore_when_discovered: ["__init__"]

‎releasenotes/notes/add-md-header-splitter-df5c024a6ddd2718.yaml‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
---
2+
features:
3+
- |
4+
Introduced the ``MarkdownHeaderSplitter`` component:
5+
- Splits documents into chunks at Markdown headers (``#``, ``##``, etc.), preserving header hierarchy as metadata.
6+
- Supports secondary splitting (by word, passage, period, or line) for further chunking after header-based splitting using Haystack's ``DocumentSplitter``.
7+
- Preserves and propagates metadata such as parent headers and page numbers.
8+
- Handles edge cases such as documents with no headers, empty content, and non-text documents.

0 commit comments

Comments
 (0)