chore: import upstream snapshot with attribution
CodeQL / Analyze (python) (push) Has been cancelled
Update Platform Components Table / update (push) Has been cancelled
Docker image release / Build base image (push) Has been cancelled
Sync docs with Docusaurus / sync (push) Has been cancelled
Tests / Check if changed (push) Has been cancelled
Tests / format (push) Has been cancelled
Tests / check-imports (push) Has been cancelled
Tests / Unit / macos-latest (push) Has been cancelled
Tests / Unit / ubuntu-latest (push) Has been cancelled
Tests / Unit / windows-latest (push) Has been cancelled
Tests / mypy (push) Has been cancelled
Tests / Integration / ubuntu-latest (push) Has been cancelled
Tests / Integration / macos-latest (push) Has been cancelled
Tests / Integration / windows-latest (push) Has been cancelled
Tests / notify-slack-on-failure (push) Has been cancelled
Tests / Mark tests as completed (push) Has been cancelled

This commit is contained in:
wehub-resource-sync
2026-07-13 13:22:28 +08:00
commit c56bef871b
9296 changed files with 1854228 additions and 0 deletions
@@ -0,0 +1,3 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
@@ -0,0 +1,220 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
from haystack import Document
from haystack.components.preprocessors.csv_document_cleaner import CSVDocumentCleaner
def test_empty_column() -> None:
csv_content = """,A,B,C
,1,2,3
,4,5,6
"""
csv_document = Document(content=csv_content)
csv_document_cleaner = CSVDocumentCleaner()
result = csv_document_cleaner.run([csv_document])
cleaned_document = result["documents"][0]
assert cleaned_document.content == "A,B,C\n1,2,3\n4,5,6\n"
def test_empty_row() -> None:
csv_content = """A,B,C
1,2,3
,,
4,5,6
"""
csv_document = Document(content=csv_content)
csv_document_cleaner = CSVDocumentCleaner()
result = csv_document_cleaner.run([csv_document])
cleaned_document = result["documents"][0]
assert cleaned_document.content == "A,B,C\n1,2,3\n4,5,6\n"
def test_empty_column_and_row() -> None:
csv_content = """,A,B,C
,1,2,3
,,,
,4,5,6
"""
csv_document = Document(content=csv_content)
csv_document_cleaner = CSVDocumentCleaner()
result = csv_document_cleaner.run([csv_document])
cleaned_document = result["documents"][0]
assert cleaned_document.content == "A,B,C\n1,2,3\n4,5,6\n"
def test_ignore_rows() -> None:
csv_content = """,,
A,B,C
4,5,6
7,8,9
"""
csv_document = Document(content=csv_content, meta={"name": "test.csv"})
csv_document_cleaner = CSVDocumentCleaner(ignore_rows=1)
result = csv_document_cleaner.run([csv_document])
cleaned_document = result["documents"][0]
assert cleaned_document.content == ",,\nA,B,C\n4,5,6\n7,8,9\n"
assert cleaned_document.meta == {"name": "test.csv"}
def test_ignore_rows_2() -> None:
csv_content = """A,B,C
,,
4,5,6
7,8,9
"""
csv_document = Document(content=csv_content, meta={"name": "test.csv"})
csv_document_cleaner = CSVDocumentCleaner(ignore_rows=1)
result = csv_document_cleaner.run([csv_document])
cleaned_document = result["documents"][0]
assert cleaned_document.content == "A,B,C\n4,5,6\n7,8,9\n"
assert cleaned_document.meta == {"name": "test.csv"}
def test_ignore_rows_3() -> None:
csv_content = """A,B,C
4,,6
7,,9
"""
csv_document = Document(content=csv_content, meta={"name": "test.csv"})
csv_document_cleaner = CSVDocumentCleaner(ignore_rows=1)
result = csv_document_cleaner.run([csv_document])
cleaned_document = result["documents"][0]
assert cleaned_document.content == "A,C\n4,6\n7,9\n"
assert cleaned_document.meta == {"name": "test.csv"}
def test_ignore_columns() -> None:
csv_content = """,,A,B
,2,3,4
,7,8,9
"""
csv_document = Document(content=csv_content)
csv_document_cleaner = CSVDocumentCleaner(ignore_columns=1)
result = csv_document_cleaner.run([csv_document])
cleaned_document = result["documents"][0]
assert cleaned_document.content == ",,A,B\n,2,3,4\n,7,8,9\n"
def test_too_many_ignore_rows() -> None:
csv_content = """,,
A,B,C
4,5,6
"""
csv_document = Document(content=csv_content)
csv_document_cleaner = CSVDocumentCleaner(ignore_rows=4)
result = csv_document_cleaner.run([csv_document])
cleaned_document = result["documents"][0]
assert cleaned_document.content == ",,\nA,B,C\n4,5,6\n"
def test_too_many_ignore_columns() -> None:
csv_content = """,,
A,B,C
4,5,6
"""
csv_document = Document(content=csv_content)
csv_document_cleaner = CSVDocumentCleaner(ignore_columns=4)
result = csv_document_cleaner.run([csv_document])
cleaned_document = result["documents"][0]
assert cleaned_document.content == ",,\nA,B,C\n4,5,6\n"
def test_ignore_rows_and_columns() -> None:
csv_content = """,A,B,C
1,item,s,
2,item2,fd,
"""
csv_document = Document(content=csv_content)
csv_document_cleaner = CSVDocumentCleaner(ignore_columns=1, ignore_rows=1)
result = csv_document_cleaner.run([csv_document])
cleaned_document = result["documents"][0]
assert cleaned_document.content == ",A,B\n1,item,s\n2,item2,fd\n"
def test_zero_ignore_rows_and_columns() -> None:
csv_content = """,A,B,C
1,item,s,
2,item2,fd,
"""
csv_document = Document(content=csv_content)
csv_document_cleaner = CSVDocumentCleaner(ignore_columns=0, ignore_rows=0)
result = csv_document_cleaner.run([csv_document])
cleaned_document = result["documents"][0]
assert cleaned_document.content == ",A,B,C\n1,item,s,\n2,item2,fd,\n"
def test_empty_document() -> None:
csv_document = Document(content="")
csv_document_cleaner = CSVDocumentCleaner()
result = csv_document_cleaner.run([csv_document])
cleaned_document = result["documents"][0]
assert cleaned_document.content == ""
assert cleaned_document.meta == {}
def test_empty_documents() -> None:
csv_document_cleaner = CSVDocumentCleaner()
result = csv_document_cleaner.run([])
assert result["documents"] == []
def test_keep_id() -> None:
csv_content = """,A,B,C
1,item,s,
"""
csv_document = Document(id="123", content=csv_content)
csv_document_cleaner = CSVDocumentCleaner(keep_id=True)
result = csv_document_cleaner.run([csv_document])
cleaned_document = result["documents"][0]
assert cleaned_document.id == "123"
assert cleaned_document.content == ",A,B,C\n1,item,s,\n"
def test_id_not_none() -> None:
csv_content = """,A,B,C
1,item,s,
"""
csv_document = Document(content=csv_content)
csv_document_cleaner = CSVDocumentCleaner()
result = csv_document_cleaner.run([csv_document])
cleaned_document = result["documents"][0]
assert cleaned_document.id != ""
assert cleaned_document.content == ",A,B,C\n1,item,s,\n"
def test_remove_empty_rows_false() -> None:
csv_content = """,B,C
,,
,5,6
"""
csv_document = Document(content=csv_content)
csv_document_cleaner = CSVDocumentCleaner(remove_empty_rows=False)
result = csv_document_cleaner.run([csv_document])
cleaned_document = result["documents"][0]
assert cleaned_document.content == "B,C\n,\n5,6\n"
def test_remove_empty_columns_false() -> None:
csv_content = """,B,C
,,
,,4
"""
csv_document = Document(content=csv_content)
csv_document_cleaner = CSVDocumentCleaner(remove_empty_columns=False)
result = csv_document_cleaner.run([csv_document])
cleaned_document = result["documents"][0]
assert cleaned_document.content == ",B,C\n,,4\n"
def test_remove_empty_rows_and_columns_false() -> None:
csv_content = """,B,C
,,4
,,
"""
csv_document = Document(content=csv_content)
csv_document_cleaner = CSVDocumentCleaner(remove_empty_rows=False, remove_empty_columns=False)
result = csv_document_cleaner.run([csv_document])
cleaned_document = result["documents"][0]
assert cleaned_document.content == ",B,C\n,,4\n,,\n"
@@ -0,0 +1,345 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
import logging
from io import StringIO
import pytest
from pandas import read_csv
from haystack import Document
from haystack.components.preprocessors.csv_document_splitter import CSVDocumentSplitter
from haystack.core.serialization import component_from_dict, component_to_dict
@pytest.fixture
def splitter() -> CSVDocumentSplitter:
return CSVDocumentSplitter()
@pytest.fixture
def csv_with_four_rows() -> str:
return """A,B,C
1,2,3
X,Y,Z
7,8,9
"""
@pytest.fixture
def two_tables_sep_by_two_empty_rows() -> str:
return """A,B,C
1,2,3
,,
,,
X,Y,Z
7,8,9
"""
@pytest.fixture
def three_tables_sep_by_empty_rows() -> str:
return """A,B,C
,,
1,2,3
,,
,,
X,Y,Z
7,8,9
"""
@pytest.fixture
def two_tables_sep_by_two_empty_columns() -> str:
return """A,B,,,X,Y
1,2,,,7,8
3,4,,,9,10
"""
class TestFindSplitIndices:
def test_find_split_indices_row_two_tables(
self, splitter: CSVDocumentSplitter, two_tables_sep_by_two_empty_rows: str
) -> None:
df = read_csv(StringIO(two_tables_sep_by_two_empty_rows), header=None, dtype=object) # type: ignore
result = splitter._find_split_indices(df, split_threshold=2, axis="row")
assert result == [(2, 3)]
def test_find_split_indices_row_two_tables_with_empty_row(
self, splitter: CSVDocumentSplitter, three_tables_sep_by_empty_rows: str
) -> None:
df = read_csv(StringIO(three_tables_sep_by_empty_rows), header=None, dtype=object) # type: ignore
result = splitter._find_split_indices(df, split_threshold=2, axis="row")
assert result == [(3, 4)]
def test_find_split_indices_row_three_tables(self, splitter: CSVDocumentSplitter) -> None:
csv_content = """A,B,C
1,2,3
,,
,,
X,Y,Z
7,8,9
,,
,,
P,Q,R
"""
df = read_csv(StringIO(csv_content), header=None, dtype=object) # type: ignore
result = splitter._find_split_indices(df, split_threshold=2, axis="row")
assert result == [(2, 3), (6, 7)]
def test_find_split_indices_column_two_tables(
self, splitter: CSVDocumentSplitter, two_tables_sep_by_two_empty_columns: str
) -> None:
df = read_csv(StringIO(two_tables_sep_by_two_empty_columns), header=None, dtype=object) # type: ignore
result = splitter._find_split_indices(df, split_threshold=1, axis="column")
assert result == [(2, 3)]
def test_find_split_indices_column_two_tables_with_empty_column(self, splitter: CSVDocumentSplitter) -> None:
csv_content = """A,,B,,,X,Y
1,,2,,,7,8
3,,4,,,9,10
"""
df = read_csv(StringIO(csv_content), header=None, dtype=object) # type: ignore
result = splitter._find_split_indices(df, split_threshold=2, axis="column")
assert result == [(3, 4)]
def test_find_split_indices_column_three_tables(self, splitter: CSVDocumentSplitter) -> None:
csv_content = """A,B,,,X,Y,,,P,Q
1,2,,,7,8,,,11,12
3,4,,,9,10,,,13,14
"""
df = read_csv(StringIO(csv_content), header=None, dtype=object) # type: ignore
result = splitter._find_split_indices(df, split_threshold=2, axis="column")
assert result == [(2, 3), (6, 7)]
class TestInit:
def test_row_split_threshold_raises_error(self) -> None:
with pytest.raises(ValueError, match="row_split_threshold must be greater than 0"):
CSVDocumentSplitter(row_split_threshold=-1)
def test_column_split_threshold_raises_error(self) -> None:
with pytest.raises(ValueError, match="column_split_threshold must be greater than 0"):
CSVDocumentSplitter(column_split_threshold=-1)
def test_row_split_threshold_and_row_column_threshold_none(self) -> None:
with pytest.raises(
ValueError, match="At least one of row_split_threshold or column_split_threshold must be specified."
):
CSVDocumentSplitter(row_split_threshold=None, column_split_threshold=None)
class TestCSVDocumentSplitter:
def test_single_table_no_split(self, splitter: CSVDocumentSplitter) -> None:
csv_content = """A,B,C
1,2,3
4,5,6
"""
doc = Document(content=csv_content, id="test_id")
result = splitter.run([doc])["documents"]
assert len(result) == 1
assert result[0].content == csv_content
assert result[0].meta == {"source_id": "test_id", "row_idx_start": 0, "col_idx_start": 0, "split_id": 0}
def test_row_split(self, splitter: CSVDocumentSplitter, two_tables_sep_by_two_empty_rows: str) -> None:
doc = Document(content=two_tables_sep_by_two_empty_rows, id="test_id")
result = splitter.run([doc])["documents"]
assert len(result) == 2
expected_tables = ["A,B,C\n1,2,3\n", "X,Y,Z\n7,8,9\n"]
expected_meta = [
{"source_id": "test_id", "row_idx_start": 0, "col_idx_start": 0, "split_id": 0},
{"source_id": "test_id", "row_idx_start": 4, "col_idx_start": 0, "split_id": 1},
]
for i, table in enumerate(result):
assert table.content == expected_tables[i]
assert table.meta == expected_meta[i]
def test_column_split(self, splitter: CSVDocumentSplitter, two_tables_sep_by_two_empty_columns: str) -> None:
doc = Document(content=two_tables_sep_by_two_empty_columns, id="test_id")
result = splitter.run([doc])["documents"]
assert len(result) == 2
expected_tables = ["A,B\n1,2\n3,4\n", "X,Y\n7,8\n9,10\n"]
expected_meta = [
{"source_id": "test_id", "row_idx_start": 0, "col_idx_start": 0, "split_id": 0},
{"source_id": "test_id", "row_idx_start": 0, "col_idx_start": 4, "split_id": 1},
]
for i, table in enumerate(result):
assert table.content == expected_tables[i]
assert table.meta == expected_meta[i]
def test_recursive_split_one_level(self, splitter: CSVDocumentSplitter) -> None:
csv_content = """A,B,,,X,Y
1,2,,,7,8
,,,,,
,,,,,
P,Q,,,M,N
3,4,,,9,10
"""
doc = Document(content=csv_content, id="test_id")
result = splitter.run([doc])["documents"]
assert len(result) == 4
expected_tables = ["A,B\n1,2\n", "X,Y\n7,8\n", "P,Q\n3,4\n", "M,N\n9,10\n"]
expected_meta = [
{"source_id": "test_id", "row_idx_start": 0, "col_idx_start": 0, "split_id": 0},
{"source_id": "test_id", "row_idx_start": 0, "col_idx_start": 4, "split_id": 1},
{"source_id": "test_id", "row_idx_start": 4, "col_idx_start": 0, "split_id": 2},
{"source_id": "test_id", "row_idx_start": 4, "col_idx_start": 4, "split_id": 3},
]
for i, table in enumerate(result):
assert table.content == expected_tables[i]
assert table.meta == expected_meta[i]
def test_recursive_split_two_levels(self, splitter: CSVDocumentSplitter) -> None:
csv_content = """A,B,,,X,Y
1,2,,,7,8
,,,,M,N
,,,,9,10
P,Q,,,,
3,4,,,,
"""
doc = Document(content=csv_content, id="test_id")
result = splitter.run([doc])["documents"]
assert len(result) == 3
expected_tables = ["A,B\n1,2\n", "X,Y\n7,8\nM,N\n9,10\n", "P,Q\n3,4\n"]
expected_meta = [
{"source_id": "test_id", "row_idx_start": 0, "col_idx_start": 0, "split_id": 0},
{"source_id": "test_id", "row_idx_start": 0, "col_idx_start": 4, "split_id": 1},
{"source_id": "test_id", "row_idx_start": 4, "col_idx_start": 0, "split_id": 2},
]
for i, table in enumerate(result):
assert table.content == expected_tables[i]
assert table.meta == expected_meta[i]
def test_csv_with_blank_lines(self, splitter: CSVDocumentSplitter) -> None:
csv_data = """ID,LeftVal,,,RightVal,Extra
1,Hello,,,World,Joined
2,StillLeft,,,StillRight,Bridge
A,B,,,C,D
E,F,,,G,H
"""
splitter = CSVDocumentSplitter(row_split_threshold=1, column_split_threshold=1)
result = splitter.run([Document(content=csv_data, id="test_id")])
docs = result["documents"]
assert len(docs) == 4
expected_tables = [
"ID,LeftVal\n1,Hello\n2,StillLeft\n",
"RightVal,Extra\nWorld,Joined\nStillRight,Bridge\n",
"A,B\nE,F\n",
"C,D\nG,H\n",
]
expected_meta = [
{"source_id": "test_id", "row_idx_start": 0, "col_idx_start": 0, "split_id": 0},
{"source_id": "test_id", "row_idx_start": 0, "col_idx_start": 4, "split_id": 1},
{"source_id": "test_id", "row_idx_start": 4, "col_idx_start": 0, "split_id": 2},
{"source_id": "test_id", "row_idx_start": 4, "col_idx_start": 4, "split_id": 3},
]
for i, table in enumerate(docs):
assert table.content == expected_tables[i]
assert table.meta == expected_meta[i]
def test_sub_table_with_one_row(self):
splitter = CSVDocumentSplitter(row_split_threshold=1)
doc = Document(content="""A,B,C\n1,2,3\n,,\n4,5,6""")
split_result = splitter.run([doc])
assert len(split_result["documents"]) == 2
def test_threshold_no_effect(self, two_tables_sep_by_two_empty_rows: str) -> None:
splitter = CSVDocumentSplitter(row_split_threshold=3)
doc = Document(content=two_tables_sep_by_two_empty_rows)
result = splitter.run([doc])["documents"]
assert len(result) == 1
def test_empty_input(self, splitter: CSVDocumentSplitter) -> None:
csv_content = ""
doc = Document(content=csv_content)
result = splitter.run([doc])["documents"]
assert len(result) == 1
assert result[0].content == csv_content
def test_empty_documents(self, splitter: CSVDocumentSplitter) -> None:
result = splitter.run([])["documents"]
assert len(result) == 0
def test_to_dict_with_defaults(self) -> None:
splitter = CSVDocumentSplitter()
config_serialized = component_to_dict(splitter, name="CSVDocumentSplitter")
config = {
"type": "haystack.components.preprocessors.csv_document_splitter.CSVDocumentSplitter",
"init_parameters": {
"row_split_threshold": 2,
"column_split_threshold": 2,
"read_csv_kwargs": {},
"split_mode": "threshold",
},
}
assert config_serialized == config
def test_to_dict_non_defaults(self) -> None:
splitter = CSVDocumentSplitter(row_split_threshold=1, column_split_threshold=None, read_csv_kwargs={"sep": ";"})
config_serialized = component_to_dict(splitter, name="CSVDocumentSplitter")
config = {
"type": "haystack.components.preprocessors.csv_document_splitter.CSVDocumentSplitter",
"init_parameters": {
"row_split_threshold": 1,
"column_split_threshold": None,
"read_csv_kwargs": {"sep": ";"},
"split_mode": "threshold",
},
}
assert config_serialized == config
def test_from_dict_defaults(self) -> None:
splitter = component_from_dict(
CSVDocumentSplitter,
data={
"type": "haystack.components.preprocessors.csv_document_splitter.CSVDocumentSplitter",
"init_parameters": {},
},
name="CSVDocumentSplitter",
)
assert splitter.row_split_threshold == 2
assert splitter.column_split_threshold == 2
assert splitter.read_csv_kwargs == {}
assert splitter.split_mode == "threshold"
def test_from_dict_non_defaults(self) -> None:
splitter = component_from_dict(
CSVDocumentSplitter,
data={
"type": "haystack.components.preprocessors.csv_document_splitter.CSVDocumentSplitter",
"init_parameters": {
"row_split_threshold": 1,
"column_split_threshold": None,
"read_csv_kwargs": {"sep": ";"},
"split_mode": "row-wise",
},
},
name="CSVDocumentSplitter",
)
assert splitter.row_split_threshold == 1
assert splitter.column_split_threshold is None
assert splitter.read_csv_kwargs == {"sep": ";"}
assert splitter.split_mode == "row-wise"
def test_split_by_row(self, csv_with_four_rows: str) -> None:
splitter = CSVDocumentSplitter(split_mode="row-wise")
doc = Document(content=csv_with_four_rows)
result = splitter.run([doc])["documents"]
assert len(result) == 4
assert result[0].content == "A,B,C\n"
assert result[1].content == "1,2,3\n"
assert result[2].content == "X,Y,Z\n"
def test_split_by_row_with_empty_rows(self, caplog) -> None:
splitter = CSVDocumentSplitter(split_mode="row-wise")
doc = Document(content="")
with caplog.at_level(logging.ERROR):
result = splitter.run([doc])["documents"]
assert len(result) == 1
assert result[0].content == ""
def test_incorrect_split_mode(self) -> None:
with pytest.raises(ValueError, match="not recognized"):
CSVDocumentSplitter(split_mode="incorrect_mode")
@@ -0,0 +1,331 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
import logging
import pytest
from haystack import Document
from haystack.components.preprocessors import DocumentCleaner
from haystack.dataclasses import ByteStream, SparseEmbedding
class TestDocumentCleaner:
def test_init(self):
cleaner = DocumentCleaner()
assert cleaner.remove_empty_lines is True
assert cleaner.remove_extra_whitespaces is True
assert cleaner.remove_repeated_substrings is False
assert cleaner.remove_substrings is None
assert cleaner.remove_regex is None
assert cleaner.keep_id is False
def test_non_text_document(self, caplog):
with caplog.at_level(logging.WARNING):
cleaner = DocumentCleaner()
cleaner.run(documents=[Document()])
assert "DocumentCleaner only cleans text documents but document.content for document ID" in caplog.text
def test_single_document(self):
with pytest.raises(TypeError, match="DocumentCleaner expects a List of Documents as input."):
cleaner = DocumentCleaner()
cleaner.run(documents=Document())
def test_empty_list(self):
cleaner = DocumentCleaner()
result = cleaner.run(documents=[])
assert result == {"documents": []}
def test_remove_empty_lines(self):
cleaner = DocumentCleaner(remove_extra_whitespaces=False)
result = cleaner.run(
documents=[
Document(
content="This is a text with some words. \f"
""
"There is a second sentence. "
""
"And there is a third sentence."
)
]
)
assert len(result["documents"]) == 1
assert (
result["documents"][0].content
== "This is a text with some words. \fThere is a second sentence. And there is a third sentence."
)
def test_remove_whitespaces(self):
cleaner = DocumentCleaner(remove_empty_lines=False)
result = cleaner.run(
documents=[
Document(
content=" This is a text with some words. "
""
"There is a second sentence. "
""
"And there is a third sentence.\f "
)
]
)
assert len(result["documents"]) == 1
assert result["documents"][0].content == (
"This is a text with some words. There is a second sentence. And there is a third sentence.\f"
)
def test_remove_substrings(self):
cleaner = DocumentCleaner(remove_substrings=["This", "A", "words", "🪲"])
result = cleaner.run(documents=[Document(content="This is a text with some words.\f🪲")])
assert len(result["documents"]) == 1
assert result["documents"][0].content == " is a text with some .\f"
def test_remove_regex(self):
cleaner = DocumentCleaner(remove_regex=r"\s\s+")
result = cleaner.run(documents=[Document(content="This is a text \f with some words.")])
assert len(result["documents"]) == 1
assert result["documents"][0].content == "This is a text\fwith some words."
def test_remove_repeated_substrings(self):
cleaner = DocumentCleaner(
remove_empty_lines=False, remove_extra_whitespaces=False, remove_repeated_substrings=True
)
text = """First Page\f This is a header.
Page of
2
4
Lorem ipsum dolor sit amet
This is a footer number 1
This is footer number 2 This is a header.
Page of
3
4
Sid ut perspiciatis unde
This is a footer number 1
This is footer number 2 This is a header.
Page of
4
4
Sed do eiusmod tempor.
This is a footer number 1
This is footer number 2"""
expected_text = """First Page\f 2
4
Lorem ipsum dolor sit amet 3
4
Sid ut perspiciatis unde 4
4
Sed do eiusmod tempor."""
result = cleaner.run(documents=[Document(content=text)])
assert result["documents"][0].content == expected_text
def test_remove_repeated_substrings_preserves_unique_middle_page(self):
cleaner = DocumentCleaner(
remove_empty_lines=False, remove_extra_whitespaces=False, remove_repeated_substrings=True
)
text = "PAGE ONE\fThe quick brown fox jumps high\fPAGE THREE"
result = cleaner.run(documents=[Document(content=text)])["documents"][0]
assert result.content.split("\f")[1] == "The quick brown fox jumps high"
# With no genuine repeated header/footer, all three pages must round-trip unchanged and in order.
assert result.content.split("\f") == ["PAGE ONE", "The quick brown fox jumps high", "PAGE THREE"]
assert result.content == text
def test_copy_metadata(self):
cleaner = DocumentCleaner()
documents = [
Document(content="Text. ", meta={"name": "doc 0"}),
Document(content="Text. ", meta={"name": "doc 1"}),
]
result = cleaner.run(documents=documents)
assert len(result["documents"]) == 2
assert result["documents"][0].id != result["documents"][1].id
for doc, cleaned_doc in zip(documents, result["documents"], strict=True):
assert doc.meta == cleaned_doc.meta
assert cleaned_doc.content == "Text."
def test_keep_id_does_not_alter_document_ids(self):
cleaner = DocumentCleaner(keep_id=True)
documents = [Document(content="Text. ", id="1"), Document(content="Text. ", id="2")]
result = cleaner.run(documents=documents)
assert len(result["documents"]) == 2
assert result["documents"][0].id == "1"
assert result["documents"][1].id == "2"
def test_unicode_normalization(self):
text = """\
アイウエオ
Comment ça va
مرحبا بالعالم
emSpace"""
expected_text_NFC = """\
アイウエオ
Comment ça va
مرحبا بالعالم
emSpace"""
expected_text_NFD = """\
アイウエオ
Comment ça va
مرحبا بالعالم
emSpace"""
expected_text_NFKC = """\
アイウエオ
Comment ça va
مرحبا بالعالم
em Space"""
expected_text_NFKD = """\
アイウエオ
Comment ça va
مرحبا بالعالم
em Space"""
nfc_cleaner = DocumentCleaner(unicode_normalization="NFC", remove_extra_whitespaces=False)
nfd_cleaner = DocumentCleaner(unicode_normalization="NFD", remove_extra_whitespaces=False)
nfkc_cleaner = DocumentCleaner(unicode_normalization="NFKC", remove_extra_whitespaces=False)
nfkd_cleaner = DocumentCleaner(unicode_normalization="NFKD", remove_extra_whitespaces=False)
nfc_result = nfc_cleaner.run(documents=[Document(content=text)])
nfd_result = nfd_cleaner.run(documents=[Document(content=text)])
nfkc_result = nfkc_cleaner.run(documents=[Document(content=text)])
nfkd_result = nfkd_cleaner.run(documents=[Document(content=text)])
assert nfc_result["documents"][0].content == expected_text_NFC
assert nfd_result["documents"][0].content == expected_text_NFD
assert nfkc_result["documents"][0].content == expected_text_NFKC
assert nfkd_result["documents"][0].content == expected_text_NFKD
def test_ascii_only(self):
text = """\
アイウエオ
Comment ça va
Á
مرحبا بالعالم
emSpace"""
expected_text = """\
\n\
Comment ca va
A
\n\
em Space"""
cleaner = DocumentCleaner(ascii_only=True, remove_extra_whitespaces=False, remove_empty_lines=False)
result = cleaner.run(documents=[Document(content=text)])
assert result["documents"][0].content == expected_text
def test_other_document_fields_are_not_lost(self):
cleaner = DocumentCleaner(keep_id=True)
document = Document(
content="This is a text with some words. \nThere is a second sentence. \nAnd there is a third sentence.\n",
blob=ByteStream.from_string("some_data"),
meta={"data": 1},
score=0.1,
embedding=[0.1, 0.2, 0.3],
sparse_embedding=SparseEmbedding([0, 2], [0.1, 0.3]),
)
res = cleaner.run(documents=[document])
assert len(res) == 1
assert len(res["documents"])
assert res["documents"][0].id == document.id
assert res["documents"][0].content == (
"This is a text with some words. There is a second sentence. And there is a third sentence."
)
assert res["documents"][0].blob == document.blob
assert res["documents"][0].meta == document.meta
assert res["documents"][0].score == document.score
assert res["documents"][0].embedding == document.embedding
assert res["documents"][0].sparse_embedding == document.sparse_embedding
def test_strip_whitespaces(self):
"""Test that strip_whitespaces removes only leading and trailing whitespace."""
cleaner = DocumentCleaner(remove_empty_lines=False, remove_extra_whitespaces=False, strip_whitespaces=True)
result = cleaner.run(documents=[Document(content=" \n\nHello World\n\n Some text here \n\n ")])
assert len(result["documents"]) == 1
# strip_whitespaces should only remove leading/trailing whitespace, preserving internal whitespace
assert result["documents"][0].content == "Hello World\n\n Some text here"
def test_strip_whitespaces_preserves_internal_formatting(self):
"""Test that strip_whitespaces preserves internal whitespace like markdown formatting."""
cleaner = DocumentCleaner(remove_empty_lines=False, remove_extra_whitespaces=False, strip_whitespaces=True)
markdown_content = """
# Header
This is a paragraph.
- Item 1
- Item 2
"""
result = cleaner.run(documents=[Document(content=markdown_content)])
assert len(result["documents"]) == 1
expected = """# Header
This is a paragraph.
- Item 1
- Item 2"""
assert result["documents"][0].content == expected
def test_replace_regexes_single_pattern(self):
"""Test replace_regexes with a single pattern."""
cleaner = DocumentCleaner(
remove_empty_lines=False, remove_extra_whitespaces=False, replace_regexes={r"\n\n+": "\n"}
)
result = cleaner.run(documents=[Document(content="Line 1\n\n\n\nLine 2\n\nLine 3")])
assert len(result["documents"]) == 1
assert result["documents"][0].content == "Line 1\nLine 2\nLine 3"
def test_replace_regexes_multiple_patterns(self):
"""Test replace_regexes with multiple patterns."""
cleaner = DocumentCleaner(
remove_empty_lines=False, remove_extra_whitespaces=False, replace_regexes={r"\n\n+": "\n", r"\s{2,}": " "}
)
result = cleaner.run(documents=[Document(content="Hello World\n\n\nGoodbye")])
assert len(result["documents"]) == 1
assert result["documents"][0].content == "Hello World\nGoodbye"
def test_replace_regexes_custom_replacement(self):
"""Test replace_regexes with custom replacement strings."""
cleaner = DocumentCleaner(
remove_empty_lines=False,
remove_extra_whitespaces=False,
replace_regexes={r"\[REDACTED\]": "***", r"(\d{4})-(\d{2})-(\d{2})": r"\2/\3/\1"},
)
result = cleaner.run(documents=[Document(content="Name: [REDACTED], Date: 2024-01-15")])
assert len(result["documents"]) == 1
assert result["documents"][0].content == "Name: ***, Date: 01/15/2024"
def test_strip_whitespaces_and_replace_regexes_combined(self):
"""Test using both strip_whitespaces and replace_regexes together."""
cleaner = DocumentCleaner(
remove_empty_lines=False,
remove_extra_whitespaces=False,
strip_whitespaces=True,
replace_regexes={r"\n\n+": "\n"},
)
result = cleaner.run(documents=[Document(content="\n\n Hello\n\n\nWorld \n\n")])
assert len(result["documents"]) == 1
assert result["documents"][0].content == "Hello\nWorld"
def test_init_with_new_params(self):
"""Test that new parameters are properly initialized."""
cleaner = DocumentCleaner(strip_whitespaces=True, replace_regexes={r"\n+": "\n"})
assert cleaner.strip_whitespaces is True
assert cleaner.replace_regexes == {r"\n+": "\n"}
def test_replace_regexes_with_page_breaks(self):
"""Test replace_regexes with page breaks (form feed character)."""
cleaner = DocumentCleaner(
remove_empty_lines=False, remove_extra_whitespaces=False, replace_regexes={r"Page \d+": ""}
)
content = "Page 1 content.\fPage 2 content."
result = cleaner.run(documents=[Document(content=content)])
assert len(result["documents"]) == 1
assert result["documents"][0].content == " content.\f content."
@@ -0,0 +1,134 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
from unittest.mock import patch
import pytest
from haystack import Document, Pipeline
from haystack.components.preprocessors.document_preprocessor import DocumentPreprocessor
class TestDocumentPreprocessor:
@pytest.fixture
def preprocessor(self) -> DocumentPreprocessor:
return DocumentPreprocessor(
# Cleaner parameters
remove_empty_lines=True,
remove_extra_whitespaces=True,
remove_repeated_substrings=False,
keep_id=True,
# Splitter parameters
split_by="word",
split_length=3,
split_overlap=1,
respect_sentence_boundary=False,
language="en",
)
def test_init(self, preprocessor: DocumentPreprocessor) -> None:
assert isinstance(preprocessor.pipeline, Pipeline)
assert preprocessor.input_mapping == {"documents": ["splitter.documents"]}
assert preprocessor.output_mapping == {"cleaner.documents": "documents"}
cleaner = preprocessor.pipeline.get_component("cleaner")
assert cleaner.remove_empty_lines is True
assert cleaner.remove_extra_whitespaces is True
assert cleaner.remove_repeated_substrings is False
assert cleaner.keep_id is True
splitter = preprocessor.pipeline.get_component("splitter")
assert splitter.split_by == "word"
assert splitter.split_length == 3
assert splitter.split_overlap == 1
assert splitter.respect_sentence_boundary is False
assert splitter.language == "en"
def test_from_dict(self) -> None:
data = {
"init_parameters": {
"remove_empty_lines": True,
"remove_extra_whitespaces": True,
"remove_repeated_substrings": False,
"keep_id": True,
"remove_substrings": None,
"remove_regex": None,
"unicode_normalization": None,
"ascii_only": False,
"split_by": "word",
"split_length": 3,
"split_overlap": 1,
"split_threshold": 0,
"splitting_function": None,
"respect_sentence_boundary": False,
"language": "en",
"use_split_rules": True,
"extend_abbreviations": True,
},
"type": "haystack.components.preprocessors.document_preprocessor.DocumentPreprocessor",
}
preprocessor = DocumentPreprocessor.from_dict(data)
assert isinstance(preprocessor, DocumentPreprocessor)
def test_to_dict(self, preprocessor: DocumentPreprocessor) -> None:
expected = {
"init_parameters": {
"remove_empty_lines": True,
"remove_extra_whitespaces": True,
"remove_repeated_substrings": False,
"keep_id": True,
"remove_substrings": None,
"remove_regex": None,
"unicode_normalization": None,
"ascii_only": False,
"split_by": "word",
"split_length": 3,
"split_overlap": 1,
"split_threshold": 0,
"splitting_function": None,
"respect_sentence_boundary": False,
"language": "en",
"use_split_rules": True,
"extend_abbreviations": True,
},
"type": "haystack.components.preprocessors.document_preprocessor.DocumentPreprocessor",
}
assert preprocessor.to_dict() == expected
def test_warm_up(self, preprocessor: DocumentPreprocessor) -> None:
with patch.object(preprocessor.pipeline, "warm_up") as mock_warm_up:
preprocessor.warm_up()
mock_warm_up.assert_called_once()
def test_run(self, preprocessor: DocumentPreprocessor) -> None:
documents = [
Document(content="This is a test document. It has multiple sentences."),
Document(content="Another test document with some content."),
]
result = preprocessor.run(documents=documents)
# Check that we got processed documents back
assert "documents" in result
processed_docs = result["documents"]
assert len(processed_docs) > len(documents) # Should have more docs due to splitting
# Check that the content was cleaned and split
for doc in processed_docs:
assert doc.content.strip() == doc.content
assert len(doc.content.split()) <= 3 # Split length of 3 words
assert doc.id is not None
def test_run_with_custom_splitting_function(self) -> None:
def custom_split(text: str) -> list[str]:
return [t for t in text.split(".") if t.strip() != ""]
preprocessor = DocumentPreprocessor(split_by="function", splitting_function=custom_split, split_length=1)
documents = [Document(content="First sentence. Second sentence. Third sentence.")]
result = preprocessor.run(documents=documents)
processed_docs = result["documents"]
assert len(processed_docs) == 3 # Should be split into 3 sentences
assert all("." not in doc.content for doc in processed_docs) # Each doc should be a single sentence
@@ -0,0 +1,836 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
import re
import pytest
from haystack import Document
from haystack.components.preprocessors import DocumentSplitter
from haystack.utils import deserialize_callable, serialize_callable
# custom split function for testing
def custom_split(text):
return text.split(".")
def merge_documents(documents):
"""Merge a list of doc chunks into a single doc by concatenating their content, eliminating overlapping content."""
sorted_docs = sorted(documents, key=lambda doc: doc.meta["split_idx_start"])
merged_text = ""
last_idx_end = 0
for doc in sorted_docs:
start = doc.meta["split_idx_start"] # start of the current content
# if the start of the current content is before the end of the last appended content, adjust it
start = max(start, last_idx_end)
# append the non-overlapping part to the merged text
merged_text += doc.content[start - doc.meta["split_idx_start"] :]
# update the last end index
last_idx_end = doc.meta["split_idx_start"] + len(doc.content)
return merged_text
class TestSplittingByFunctionOrCharacterRegex:
def test_non_text_document(self, caplog):
with pytest.raises(
ValueError, match="DocumentSplitter only works with text documents but content for document ID"
):
splitter = DocumentSplitter()
splitter.run(documents=[Document()])
assert "DocumentSplitter only works with text documents but content for document ID" in caplog.text
def test_single_doc(self):
with pytest.raises(TypeError, match="DocumentSplitter expects a List of Documents as input."):
splitter = DocumentSplitter()
splitter.run(documents=Document())
def test_empty_list(self):
splitter = DocumentSplitter()
res = splitter.run(documents=[])
assert res == {"documents": []}
def test_unsupported_split_by(self):
with pytest.raises(ValueError, match="split_by must be one of "):
DocumentSplitter(split_by="unsupported")
def test_undefined_function(self):
with pytest.raises(ValueError, match="When 'split_by' is set to 'function', a valid 'splitting_function'"):
DocumentSplitter(split_by="function", splitting_function=None)
def test_unsupported_split_length(self):
with pytest.raises(ValueError, match="split_length must be greater than 0."):
DocumentSplitter(split_length=0)
def test_unsupported_split_overlap(self):
with pytest.raises(ValueError, match="split_overlap must be greater than or equal to 0."):
DocumentSplitter(split_overlap=-1)
def test_split_overlap_not_less_than_split_length(self):
# split_overlap == split_length makes the window step 0, and a larger
# overlap makes it negative — both crash deep in `windowed` at run time.
# Fail fast at init with a clear error instead.
with pytest.raises(ValueError, match="split_overlap must be less than split_length."):
DocumentSplitter(split_length=3, split_overlap=3)
with pytest.raises(ValueError, match="split_overlap must be less than split_length."):
DocumentSplitter(split_length=3, split_overlap=5)
def test_split_by_word(self):
splitter = DocumentSplitter(split_by="word", split_length=10)
text = "This is a text with some words. There is a second sentence. And there is a third sentence."
result = splitter.run(documents=[Document(content=text)])
docs = result["documents"]
assert len(docs) == 2
assert docs[0].content == "This is a text with some words. There is a "
assert docs[0].meta["split_id"] == 0
assert docs[0].meta["split_idx_start"] == text.index(docs[0].content)
assert docs[1].content == "second sentence. And there is a third sentence."
assert docs[1].meta["split_id"] == 1
assert docs[1].meta["split_idx_start"] == text.index(docs[1].content)
def test_split_by_word_with_threshold(self):
splitter = DocumentSplitter(split_by="word", split_length=15, split_threshold=10)
result = splitter.run(
documents=[
Document(
content="This is a text with some words. There is a second sentence. And there is a third sentence."
)
]
)
assert len(result["documents"]) == 1
assert (
result["documents"][0].content
== "This is a text with some words. There is a second sentence. And there is a third sentence."
)
def test_split_by_word_multiple_input_docs(self):
splitter = DocumentSplitter(split_by="word", split_length=10)
text1 = "This is a text with some words. There is a second sentence. And there is a third sentence."
text2 = (
"This is a different text with some words. There is a second sentence. And there is a third sentence. "
"And there is a fourth sentence."
)
result = splitter.run(documents=[Document(content=text1), Document(content=text2)])
docs = result["documents"]
assert len(docs) == 5
# doc 0
assert docs[0].content == "This is a text with some words. There is a "
assert docs[0].meta["split_id"] == 0
assert docs[0].meta["split_idx_start"] == text1.index(docs[0].content)
# doc 1
assert docs[1].content == "second sentence. And there is a third sentence."
assert docs[1].meta["split_id"] == 1
assert docs[1].meta["split_idx_start"] == text1.index(docs[1].content)
# doc 2
assert docs[2].content == "This is a different text with some words. There is "
assert docs[2].meta["split_id"] == 0
assert docs[2].meta["split_idx_start"] == text2.index(docs[2].content)
# doc 3
assert docs[3].content == "a second sentence. And there is a third sentence. And "
assert docs[3].meta["split_id"] == 1
assert docs[3].meta["split_idx_start"] == text2.index(docs[3].content)
# doc 4
assert docs[4].content == "there is a fourth sentence."
assert docs[4].meta["split_id"] == 2
assert docs[4].meta["split_idx_start"] == text2.index(docs[4].content)
def test_split_by_period(self):
splitter = DocumentSplitter(split_by="period", split_length=1)
text = "This is a text with some words. There is a second sentence. And there is a third sentence."
result = splitter.run(documents=[Document(content=text)])
docs = result["documents"]
assert len(docs) == 3
assert docs[0].content == "This is a text with some words."
assert docs[0].meta["split_id"] == 0
assert docs[0].meta["split_idx_start"] == text.index(docs[0].content)
assert docs[1].content == " There is a second sentence."
assert docs[1].meta["split_id"] == 1
assert docs[1].meta["split_idx_start"] == text.index(docs[1].content)
assert docs[2].content == " And there is a third sentence."
assert docs[2].meta["split_id"] == 2
assert docs[2].meta["split_idx_start"] == text.index(docs[2].content)
def test_split_by_passage(self):
splitter = DocumentSplitter(split_by="passage", split_length=1)
text = (
"This is a text with some words. There is a second sentence.\n\nAnd there is a third sentence.\n\n "
"And another passage."
)
result = splitter.run(documents=[Document(content=text)])
docs = result["documents"]
assert len(docs) == 3
assert docs[0].content == "This is a text with some words. There is a second sentence.\n\n"
assert docs[0].meta["split_id"] == 0
assert docs[0].meta["split_idx_start"] == text.index(docs[0].content)
assert docs[1].content == "And there is a third sentence.\n\n"
assert docs[1].meta["split_id"] == 1
assert docs[1].meta["split_idx_start"] == text.index(docs[1].content)
assert docs[2].content == " And another passage."
assert docs[2].meta["split_id"] == 2
assert docs[2].meta["split_idx_start"] == text.index(docs[2].content)
def test_split_by_page(self):
splitter = DocumentSplitter(split_by="page", split_length=1)
text = (
"This is a text with some words. There is a second sentence.\f And there is a third sentence.\f And "
"another passage."
)
result = splitter.run(documents=[Document(content=text)])
docs = result["documents"]
assert len(docs) == 3
assert docs[0].content == "This is a text with some words. There is a second sentence.\f"
assert docs[0].meta["split_id"] == 0
assert docs[0].meta["split_idx_start"] == text.index(docs[0].content)
assert docs[0].meta["page_number"] == 1
assert docs[1].content == " And there is a third sentence.\f"
assert docs[1].meta["split_id"] == 1
assert docs[1].meta["split_idx_start"] == text.index(docs[1].content)
assert docs[1].meta["page_number"] == 2
assert docs[2].content == " And another passage."
assert docs[2].meta["split_id"] == 2
assert docs[2].meta["split_idx_start"] == text.index(docs[2].content)
assert docs[2].meta["page_number"] == 3
def test_split_by_function(self):
splitting_function = lambda s: s.split(".")
splitter = DocumentSplitter(split_by="function", splitting_function=splitting_function)
text = "This.Is.A.Test"
result = splitter.run(documents=[Document(id="1", content=text, meta={"key": "value"})])
docs = result["documents"]
assert len(docs) == 4
assert docs[0].content == "This"
assert docs[0].meta == {"key": "value", "source_id": "1"}
assert docs[1].content == "Is"
assert docs[1].meta == {"key": "value", "source_id": "1"}
assert docs[2].content == "A"
assert docs[2].meta == {"key": "value", "source_id": "1"}
assert docs[3].content == "Test"
assert docs[3].meta == {"key": "value", "source_id": "1"}
splitting_function = lambda s: re.split(r"[\s]{2,}", s)
splitter = DocumentSplitter(split_by="function", splitting_function=splitting_function)
text = "This Is\n A Test"
result = splitter.run(documents=[Document(id="1", content=text, meta={"key": "value"})])
docs = result["documents"]
assert len(docs) == 4
assert docs[0].content == "This"
assert docs[0].meta == {"key": "value", "source_id": "1"}
assert docs[1].content == "Is"
assert docs[1].meta == {"key": "value", "source_id": "1"}
assert docs[2].content == "A"
assert docs[2].meta == {"key": "value", "source_id": "1"}
assert docs[3].content == "Test"
assert docs[3].meta == {"key": "value", "source_id": "1"}
def test_split_by_word_with_overlap(self):
splitter = DocumentSplitter(split_by="word", split_length=10, split_overlap=2)
text = "This is a text with some words. There is a second sentence. And there is a third sentence."
result = splitter.run(documents=[Document(content=text)])
docs = result["documents"]
assert len(docs) == 2
# doc 0
assert docs[0].content == "This is a text with some words. There is a "
assert docs[0].meta["split_id"] == 0
assert docs[0].meta["split_idx_start"] == text.index(docs[0].content)
assert docs[0].meta["_split_overlap"][0]["range"] == (0, 5)
assert docs[1].content[0:5] == "is a "
# doc 1
assert docs[1].content == "is a second sentence. And there is a third sentence."
assert docs[1].meta["split_id"] == 1
assert docs[1].meta["split_idx_start"] == text.index(docs[1].content)
assert docs[1].meta["_split_overlap"][0]["range"] == (38, 43)
assert docs[0].content[38:43] == "is a "
def test_split_by_line(self):
splitter = DocumentSplitter(split_by="line", split_length=1)
text = "This is a text with some words.\nThere is a second sentence.\nAnd there is a third sentence."
result = splitter.run(documents=[Document(content=text)])
docs = result["documents"]
assert len(docs) == 3
assert docs[0].content == "This is a text with some words.\n"
assert docs[0].meta["split_id"] == 0
assert docs[0].meta["split_idx_start"] == text.index(docs[0].content)
assert docs[1].content == "There is a second sentence.\n"
assert docs[1].meta["split_id"] == 1
assert docs[1].meta["split_idx_start"] == text.index(docs[1].content)
assert docs[2].content == "And there is a third sentence."
assert docs[2].meta["split_id"] == 2
assert docs[2].meta["split_idx_start"] == text.index(docs[2].content)
def test_source_id_stored_in_metadata(self):
splitter = DocumentSplitter(split_by="word", split_length=10)
doc1 = Document(content="This is a text with some words.")
doc2 = Document(content="This is a different text with some words.")
result = splitter.run(documents=[doc1, doc2])
assert result["documents"][0].meta["source_id"] == doc1.id
assert result["documents"][1].meta["source_id"] == doc2.id
def test_copy_metadata(self):
splitter = DocumentSplitter(split_by="word", split_length=10)
documents = [
Document(content="Text.", meta={"name": "doc 0"}),
Document(content="Text.", meta={"name": "doc 1"}),
]
result = splitter.run(documents=documents)
assert len(result["documents"]) == 2
assert result["documents"][0].id != result["documents"][1].id
for doc, split_doc in zip(documents, result["documents"], strict=True):
assert doc.meta.items() <= split_doc.meta.items()
assert split_doc.content == "Text."
def test_add_page_number_to_metadata_with_no_overlap_word_split(self):
splitter = DocumentSplitter(split_by="word", split_length=2)
doc1 = Document(content="This is some text.\f This text is on another page.")
doc2 = Document(content="This content has two.\f\f page brakes.")
result = splitter.run(documents=[doc1, doc2])
expected_pages = [1, 1, 2, 2, 2, 1, 1, 3]
for doc, p in zip(result["documents"], expected_pages, strict=True):
assert doc.meta["page_number"] == p
def test_add_page_number_to_metadata_with_no_overlap_period_split(self):
splitter = DocumentSplitter(split_by="period", split_length=1)
doc1 = Document(content="This is some text.\f This text is on another page.")
doc2 = Document(content="This content has two.\f\f page brakes.")
result = splitter.run(documents=[doc1, doc2])
expected_pages = [1, 1, 1, 1]
for doc, p in zip(result["documents"], expected_pages, strict=True):
assert doc.meta["page_number"] == p
def test_add_page_number_to_metadata_with_no_overlap_passage_split(self):
splitter = DocumentSplitter(split_by="passage", split_length=1)
doc1 = Document(
content="This is a text with some words.\f There is a second sentence.\n\nAnd there is a third sentence."
"\n\nAnd more passages.\n\n\f And another passage."
)
result = splitter.run(documents=[doc1])
expected_pages = [1, 2, 2, 2]
for doc, p in zip(result["documents"], expected_pages, strict=True):
assert doc.meta["page_number"] == p
def test_add_page_number_to_metadata_with_no_overlap_page_split(self):
splitter = DocumentSplitter(split_by="page", split_length=1)
doc1 = Document(
content="This is a text with some words. There is a second sentence.\f And there is a third sentence.\f "
"And another passage."
)
result = splitter.run(documents=[doc1])
expected_pages = [1, 2, 3]
for doc, p in zip(result["documents"], expected_pages, strict=True):
assert doc.meta["page_number"] == p
splitter = DocumentSplitter(split_by="page", split_length=2)
doc1 = Document(
content="This is a text with some words. There is a second sentence.\f And there is a third sentence.\f "
"And another passage."
)
result = splitter.run(documents=[doc1])
expected_pages = [1, 3]
for doc, p in zip(result["documents"], expected_pages, strict=True):
assert doc.meta["page_number"] == p
def test_add_page_number_to_metadata_with_overlap_word_split(self):
splitter = DocumentSplitter(split_by="word", split_length=3, split_overlap=1)
doc1 = Document(content="This is some text. And\f this text is on another page.")
doc2 = Document(content="This content has two.\f\f page brakes.")
result = splitter.run(documents=[doc1, doc2])
expected_pages = [1, 1, 1, 2, 2, 1, 1, 3]
for doc, p in zip(result["documents"], expected_pages, strict=True):
assert doc.meta["page_number"] == p
def test_add_page_number_to_metadata_with_overlap_period_split(self):
splitter = DocumentSplitter(split_by="period", split_length=2, split_overlap=1)
doc1 = Document(content="This is some text. And this is more text.\f This text is on another page. End.")
doc2 = Document(content="This content has two.\f\f page brakes. More text.")
result = splitter.run(documents=[doc1, doc2])
expected_pages = [1, 1, 1, 2, 1, 1, 3]
for doc, p in zip(result["documents"], expected_pages, strict=True):
assert doc.meta["page_number"] == p
def test_add_page_number_to_metadata_with_overlap_passage_split(self):
splitter = DocumentSplitter(split_by="passage", split_length=2, split_overlap=1)
doc1 = Document(
content="This is a text with some words.\f There is a second sentence.\n\nAnd there is a third sentence."
"\n\nAnd more passages.\n\n\f And another passage."
)
result = splitter.run(documents=[doc1])
expected_pages = [1, 2, 2]
for doc, p in zip(result["documents"], expected_pages, strict=True):
assert doc.meta["page_number"] == p
def test_add_page_number_to_metadata_with_overlap_page_split(self):
splitter = DocumentSplitter(split_by="page", split_length=2, split_overlap=1)
doc1 = Document(
content="This is a text with some words. There is a second sentence.\f And there is a third sentence.\f "
"And another passage."
)
result = splitter.run(documents=[doc1])
expected_pages = [1, 2]
for doc, p in zip(result["documents"], expected_pages, strict=True):
assert doc.meta["page_number"] == p
def test_add_split_overlap_information(self):
splitter = DocumentSplitter(split_length=10, split_overlap=5, split_by="word")
text = "This is a text with some words. There is a second sentence. And a third sentence."
doc = Document(content="This is a text with some words. There is a second sentence. And a third sentence.")
docs = splitter.run(documents=[doc])["documents"]
# check split_overlap is added to all the documents
assert len(docs) == 3
# doc 0
assert docs[0].content == "This is a text with some words. There is a "
assert docs[0].meta["split_id"] == 0
assert docs[0].meta["split_idx_start"] == text.index(docs[0].content) # 0
assert docs[0].meta["_split_overlap"][0]["range"] == (0, 23)
assert docs[1].content[0:23] == "some words. There is a "
# doc 1
assert docs[1].content == "some words. There is a second sentence. And a third "
assert docs[1].meta["split_id"] == 1
assert docs[1].meta["split_idx_start"] == text.index(docs[1].content) # 20
assert docs[1].meta["_split_overlap"][0]["range"] == (20, 43)
assert docs[1].meta["_split_overlap"][1]["range"] == (0, 29)
assert docs[0].content[20:43] == "some words. There is a "
assert docs[2].content[0:29] == "second sentence. And a third "
# doc 2
assert docs[2].content == "second sentence. And a third sentence."
assert docs[2].meta["split_id"] == 2
assert docs[2].meta["split_idx_start"] == text.index(docs[2].content) # 43
assert docs[2].meta["_split_overlap"][0]["range"] == (23, 52)
assert docs[1].content[23:52] == "second sentence. And a third "
# reconstruct the original document content from the split documents
assert doc.content == merge_documents(docs)
def test_to_dict(self):
"""
Test the to_dict method of the DocumentSplitter class.
"""
splitter = DocumentSplitter(split_by="word", split_length=10, split_overlap=2, split_threshold=5)
serialized = splitter.to_dict()
assert serialized["type"] == "haystack.components.preprocessors.document_splitter.DocumentSplitter"
assert serialized["init_parameters"]["split_by"] == "word"
assert serialized["init_parameters"]["split_length"] == 10
assert serialized["init_parameters"]["split_overlap"] == 2
assert serialized["init_parameters"]["split_threshold"] == 5
assert serialized["init_parameters"]["skip_empty_documents"]
assert "splitting_function" not in serialized["init_parameters"]
def test_to_dict_with_splitting_function(self):
"""
Test the to_dict method of the DocumentSplitter class when a custom splitting function is provided.
"""
splitter = DocumentSplitter(split_by="function", splitting_function=custom_split)
serialized = splitter.to_dict()
assert serialized["type"] == "haystack.components.preprocessors.document_splitter.DocumentSplitter"
assert serialized["init_parameters"]["split_by"] == "function"
assert "splitting_function" in serialized["init_parameters"]
assert serialized["init_parameters"]["skip_empty_documents"]
assert callable(deserialize_callable(serialized["init_parameters"]["splitting_function"]))
def test_from_dict(self):
"""
Test the from_dict class method of the DocumentSplitter class.
"""
data = {
"type": "haystack.components.preprocessors.document_splitter.DocumentSplitter",
"init_parameters": {
"split_by": "word",
"split_length": 10,
"split_overlap": 2,
"split_threshold": 5,
"skip_empty_documents": False,
},
}
splitter = DocumentSplitter.from_dict(data)
assert splitter.split_by == "word"
assert splitter.split_length == 10
assert splitter.split_overlap == 2
assert splitter.split_threshold == 5
assert splitter.splitting_function is None
assert splitter.skip_empty_documents is False
def test_from_dict_with_splitting_function(self):
"""
Test the from_dict class method of the DocumentSplitter class when a custom splitting function is provided.
"""
data = {
"type": "haystack.components.preprocessors.document_splitter.DocumentSplitter",
"init_parameters": {"split_by": "function", "splitting_function": serialize_callable(custom_split)},
}
splitter = DocumentSplitter.from_dict(data)
assert splitter.split_by == "function"
assert callable(splitter.splitting_function)
assert splitter.splitting_function("a.b.c") == ["a", "b", "c"]
def test_roundtrip_serialization(self):
"""
Test the round-trip serialization of the DocumentSplitter class.
"""
original_splitter = DocumentSplitter(split_by="word", split_length=10, split_overlap=2, split_threshold=5)
serialized = original_splitter.to_dict()
deserialized_splitter = DocumentSplitter.from_dict(serialized)
assert original_splitter.split_by == deserialized_splitter.split_by
assert original_splitter.split_length == deserialized_splitter.split_length
assert original_splitter.split_overlap == deserialized_splitter.split_overlap
assert original_splitter.split_threshold == deserialized_splitter.split_threshold
def test_roundtrip_serialization_with_splitting_function(self):
"""
Test the round-trip serialization of the DocumentSplitter class when a custom splitting function is provided.
"""
original_splitter = DocumentSplitter(split_by="function", splitting_function=custom_split)
serialized = original_splitter.to_dict()
deserialized_splitter = DocumentSplitter.from_dict(serialized)
assert original_splitter.split_by == deserialized_splitter.split_by
assert callable(deserialized_splitter.splitting_function)
assert deserialized_splitter.splitting_function("a.b.c") == ["a", "b", "c"]
def test_run_empty_document_with_skip_empty_documents_true(self):
"""
Test if the component runs correctly with an empty document.
"""
splitter = DocumentSplitter()
doc = Document(content="")
results = splitter.run([doc])
assert results["documents"] == []
def test_run_empty_document_with_skip_empty_documents_false(self):
splitter = DocumentSplitter(skip_empty_documents=False)
doc = Document(content="")
results = splitter.run([doc])
assert len(results["documents"]) == 1
assert results["documents"][0].content == ""
def test_run_document_only_whitespaces(self):
"""
Test if the component runs correctly with a document containing only whitespaces.
"""
splitter = DocumentSplitter()
doc = Document(content=" ")
results = splitter.run([doc])
assert results["documents"][0].content == " "
class TestSplittingNLTKSentenceSplitter:
@pytest.mark.parametrize(
"sentences, expected_num_sentences",
[
(["The sun set.", "Moonlight shimmered softly, wolves howled nearby, night enveloped everything."], 0),
(["The sun set.", "It was a dark night ..."], 0),
(["The sun set.", " The moon was full."], 1),
(["The sun.", " The moon."], 1), # Ignores the first sentence
(["Sun", "Moon"], 1), # Ignores the first sentence even if its inclusion would be < split_overlap
],
)
def test_number_of_sentences_to_keep(self, sentences: list[str], expected_num_sentences: int) -> None:
num_sentences = DocumentSplitter._number_of_sentences_to_keep(
sentences=sentences, split_length=5, split_overlap=2
)
assert num_sentences == expected_num_sentences
def test_number_of_sentences_to_keep_split_overlap_zero(self) -> None:
sentences = [
"Moonlight shimmered softly, wolves howled nearby, night enveloped everything.",
" It was a dark night ...",
" The moon was full.",
]
num_sentences = DocumentSplitter._number_of_sentences_to_keep(
sentences=sentences, split_length=5, split_overlap=0
)
assert num_sentences == 0
def test_run_split_by_sentence_1(self) -> None:
document_splitter = DocumentSplitter(
split_by="sentence",
split_length=2,
split_overlap=0,
split_threshold=0,
language="en",
use_split_rules=True,
extend_abbreviations=True,
)
text = (
"Moonlight shimmered softly, wolves howled nearby, night enveloped everything. It was a dark night ... "
"The moon was full."
)
documents = document_splitter.run(documents=[Document(content=text)])["documents"]
assert len(documents) == 2
assert (
documents[0].content == "Moonlight shimmered softly, wolves howled nearby, night enveloped "
"everything. It was a dark night ... "
)
assert documents[1].content == "The moon was full."
def test_run_split_by_sentence_2(self) -> None:
document_splitter = DocumentSplitter(
split_by="sentence",
split_length=1,
split_overlap=0,
split_threshold=0,
language="en",
use_split_rules=False,
extend_abbreviations=True,
)
text = (
"This is a test sentence with many many words that exceeds the split length and should not be repeated. "
"This is another test sentence. (This is a third test sentence.) "
"This is the last test sentence."
)
documents = document_splitter.run(documents=[Document(content=text)])["documents"]
assert len(documents) == 4
assert (
documents[0].content
== "This is a test sentence with many many words that exceeds the split length and should not be repeated. "
)
assert documents[0].meta["page_number"] == 1
assert documents[0].meta["split_id"] == 0
assert documents[0].meta["split_idx_start"] == text.index(documents[0].content)
assert documents[1].content == "This is another test sentence. "
assert documents[1].meta["page_number"] == 1
assert documents[1].meta["split_id"] == 1
assert documents[1].meta["split_idx_start"] == text.index(documents[1].content)
assert documents[2].content == "(This is a third test sentence.) "
assert documents[2].meta["page_number"] == 1
assert documents[2].meta["split_id"] == 2
assert documents[2].meta["split_idx_start"] == text.index(documents[2].content)
assert documents[3].content == "This is the last test sentence."
assert documents[3].meta["page_number"] == 1
assert documents[3].meta["split_id"] == 3
assert documents[3].meta["split_idx_start"] == text.index(documents[3].content)
def test_run_split_by_sentence_3(self) -> None:
document_splitter = DocumentSplitter(
split_by="sentence",
split_length=1,
split_overlap=0,
split_threshold=0,
language="en",
use_split_rules=True,
extend_abbreviations=True,
)
text = "Sentence on page 1.\fSentence on page 2. \fSentence on page 3. \f\f Sentence on page 5."
documents = document_splitter.run(documents=[Document(content=text)])["documents"]
assert len(documents) == 4
assert documents[0].content == "Sentence on page 1.\f"
assert documents[0].meta["page_number"] == 1
assert documents[0].meta["split_id"] == 0
assert documents[0].meta["split_idx_start"] == text.index(documents[0].content)
assert documents[1].content == "Sentence on page 2. \f"
assert documents[1].meta["page_number"] == 2
assert documents[1].meta["split_id"] == 1
assert documents[1].meta["split_idx_start"] == text.index(documents[1].content)
assert documents[2].content == "Sentence on page 3. \f\f "
assert documents[2].meta["page_number"] == 3
assert documents[2].meta["split_id"] == 2
assert documents[2].meta["split_idx_start"] == text.index(documents[2].content)
assert documents[3].content == "Sentence on page 5."
assert documents[3].meta["page_number"] == 5
assert documents[3].meta["split_id"] == 3
assert documents[3].meta["split_idx_start"] == text.index(documents[3].content)
def test_run_split_by_sentence_4(self) -> None:
document_splitter = DocumentSplitter(
split_by="sentence",
split_length=2,
split_overlap=1,
split_threshold=0,
language="en",
use_split_rules=True,
extend_abbreviations=True,
)
text = "Sentence on page 1.\fSentence on page 2. \fSentence on page 3. \f\f Sentence on page 5."
documents = document_splitter.run(documents=[Document(content=text)])["documents"]
assert len(documents) == 3
assert documents[0].content == "Sentence on page 1.\fSentence on page 2. \f"
assert documents[0].meta["page_number"] == 1
assert documents[0].meta["split_id"] == 0
assert documents[0].meta["split_idx_start"] == text.index(documents[0].content)
assert documents[1].content == "Sentence on page 2. \fSentence on page 3. \f\f "
assert documents[1].meta["page_number"] == 2
assert documents[1].meta["split_id"] == 1
assert documents[1].meta["split_idx_start"] == text.index(documents[1].content)
assert documents[2].content == "Sentence on page 3. \f\f Sentence on page 5."
assert documents[2].meta["page_number"] == 3
assert documents[2].meta["split_id"] == 2
assert documents[2].meta["split_idx_start"] == text.index(documents[2].content)
def test_run_split_by_word_respect_sentence_boundary(self) -> None:
document_splitter = DocumentSplitter(
split_by="word",
split_length=3,
split_overlap=0,
split_threshold=0,
language="en",
respect_sentence_boundary=True,
)
text = (
"Moonlight shimmered softly, wolves howled nearby, night enveloped everything. It was a dark night.\f"
"The moon was full."
)
documents = document_splitter.run(documents=[Document(content=text)])["documents"]
assert len(documents) == 3
assert documents[0].content == "Moonlight shimmered softly, wolves howled nearby, night enveloped everything. "
assert documents[0].meta["page_number"] == 1
assert documents[0].meta["split_id"] == 0
assert documents[0].meta["split_idx_start"] == text.index(documents[0].content)
assert documents[1].content == "It was a dark night.\f"
assert documents[1].meta["page_number"] == 1
assert documents[1].meta["split_id"] == 1
assert documents[1].meta["split_idx_start"] == text.index(documents[1].content)
assert documents[2].content == "The moon was full."
assert documents[2].meta["page_number"] == 2
assert documents[2].meta["split_id"] == 2
assert documents[2].meta["split_idx_start"] == text.index(documents[2].content)
def test_run_split_by_word_respect_sentence_boundary_no_repeats(self) -> None:
document_splitter = DocumentSplitter(
split_by="word",
split_length=13,
split_overlap=3,
split_threshold=0,
language="en",
respect_sentence_boundary=True,
use_split_rules=False,
extend_abbreviations=False,
)
text = (
"This is a test sentence with many many words that exceeds the split length and should not be repeated. "
"This is another test sentence. (This is a third test sentence.) "
"This is the last test sentence."
)
documents = document_splitter.run([Document(content=text)])["documents"]
assert len(documents) == 3
assert (
documents[0].content
== "This is a test sentence with many many words that exceeds the split length and should not be repeated. "
)
assert "This is a test sentence with many many words" not in documents[1].content
assert "This is a test sentence with many many words" not in documents[2].content
def test_run_split_by_word_respect_sentence_boundary_with_split_overlap_and_page_breaks(self) -> None:
document_splitter = DocumentSplitter(
split_by="word",
split_length=8,
split_overlap=1,
split_threshold=0,
language="en",
use_split_rules=True,
extend_abbreviations=True,
respect_sentence_boundary=True,
)
text = (
"Sentence on page 1. Another on page 1.\fSentence on page 2. Another on page 2.\f"
"Sentence on page 3. Another on page 3.\f\f Sentence on page 5."
)
documents = document_splitter.run(documents=[Document(content=text)])["documents"]
assert len(documents) == 6
assert documents[0].content == "Sentence on page 1. Another on page 1.\f"
assert documents[0].meta["page_number"] == 1
assert documents[0].meta["split_id"] == 0
assert documents[0].meta["split_idx_start"] == text.index(documents[0].content)
assert documents[1].content == "Another on page 1.\fSentence on page 2. "
assert documents[1].meta["page_number"] == 1
assert documents[1].meta["split_id"] == 1
assert documents[1].meta["split_idx_start"] == text.index(documents[1].content)
assert documents[2].content == "Sentence on page 2. Another on page 2.\f"
assert documents[2].meta["page_number"] == 2
assert documents[2].meta["split_id"] == 2
assert documents[2].meta["split_idx_start"] == text.index(documents[2].content)
assert documents[3].content == "Another on page 2.\fSentence on page 3. "
assert documents[3].meta["page_number"] == 2
assert documents[3].meta["split_id"] == 3
assert documents[3].meta["split_idx_start"] == text.index(documents[3].content)
assert documents[4].content == "Sentence on page 3. Another on page 3.\f\f "
assert documents[4].meta["page_number"] == 3
assert documents[4].meta["split_id"] == 4
assert documents[4].meta["split_idx_start"] == text.index(documents[4].content)
assert documents[5].content == "Another on page 3.\f\f Sentence on page 5."
assert documents[5].meta["page_number"] == 3
assert documents[5].meta["split_id"] == 5
assert documents[5].meta["split_idx_start"] == text.index(documents[5].content)
def test_respect_sentence_boundary_checks(self):
# this combination triggers the warning
splitter = DocumentSplitter(split_by="sentence", split_length=10, respect_sentence_boundary=True)
assert splitter.respect_sentence_boundary is False
def test_sentence_serialization(self):
"""Test serialization with NLTK sentence splitting configuration and using non-default values"""
splitter = DocumentSplitter(
split_by="sentence",
language="de",
use_split_rules=False,
extend_abbreviations=False,
respect_sentence_boundary=False,
)
serialized = splitter.to_dict()
deserialized = DocumentSplitter.from_dict(serialized)
assert deserialized.split_by == "sentence"
assert hasattr(deserialized, "sentence_splitter")
assert deserialized.language == "de"
assert deserialized.use_split_rules is False
assert deserialized.extend_abbreviations is False
assert deserialized.respect_sentence_boundary is False
def test_nltk_serialization_roundtrip(self):
"""Test complete serialization roundtrip with actual document splitting"""
splitter = DocumentSplitter(
split_by="sentence",
language="de",
use_split_rules=False,
extend_abbreviations=False,
respect_sentence_boundary=False,
)
serialized = splitter.to_dict()
deserialized_splitter = DocumentSplitter.from_dict(serialized)
assert splitter.split_by == deserialized_splitter.split_by
def test_respect_sentence_boundary_serialization(self):
"""Test serialization with respect_sentence_boundary option"""
splitter = DocumentSplitter(split_by="word", respect_sentence_boundary=True, language="de")
serialized = splitter.to_dict()
deserialized = DocumentSplitter.from_dict(serialized)
assert deserialized.respect_sentence_boundary is True
assert hasattr(deserialized, "sentence_splitter")
assert deserialized.language == "de"
def test_duplicate_pages_get_different_doc_id(self):
splitter = DocumentSplitter(split_by="page", split_length=1)
doc1 = Document(content="This is some text.\fThis is some text.\fThis is some text.\fThis is some text.")
result = splitter.run(documents=[doc1])
assert len({doc.id for doc in result["documents"]}) == 4
@@ -0,0 +1,833 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
import os
from dataclasses import replace
from unittest.mock import AsyncMock, Mock, patch
import pytest
from haystack import Document
from haystack.components.embedders import OpenAIDocumentEmbedder
from haystack.components.preprocessors import EmbeddingBasedDocumentSplitter
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# disable tqdm entirely for tests
from tqdm import tqdm
tqdm.disable = True
class TestEmbeddingBasedDocumentSplitter:
def test_init(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(
document_embedder=mock_embedder, sentences_per_group=2, percentile=0.9, min_length=50, max_length=1000
)
assert splitter.document_embedder == mock_embedder
assert splitter.sentences_per_group == 2
assert splitter.percentile == 0.9
assert splitter.min_length == 50
assert splitter.max_length == 1000
def test_init_invalid_sentences_per_group(self):
mock_embedder = Mock()
with pytest.raises(ValueError, match="sentences_per_group must be greater than 0"):
EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder, sentences_per_group=0)
def test_init_invalid_percentile(self):
mock_embedder = Mock()
with pytest.raises(ValueError, match="percentile must be between 0.0 and 1.0"):
EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder, percentile=1.5)
def test_init_invalid_min_length(self):
mock_embedder = Mock()
with pytest.raises(ValueError, match="min_length must be greater than or equal to 0"):
EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder, min_length=-1)
def test_init_invalid_max_length(self):
mock_embedder = Mock()
with pytest.raises(ValueError, match="max_length must be greater than min_length"):
EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder, min_length=100, max_length=50)
def test_run_invalid_input(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
splitter.sentence_splitter = Mock()
with pytest.raises(TypeError, match="expects a List of Documents"):
splitter.run(documents="not a list")
@pytest.mark.asyncio
async def test_run_invalid_input_async(self) -> None:
mock_embedder = AsyncMock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
splitter.sentence_splitter = AsyncMock()
with pytest.raises(TypeError, match="expects a List of Documents"):
await splitter.run_async(documents="not a list")
def test_run_document_with_none_content(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
splitter.sentence_splitter = Mock()
with pytest.raises(ValueError, match="content for document ID"):
splitter.run(documents=[Document(content=None)])
@pytest.mark.asyncio
async def test_run_document_with_none_content_async(self) -> None:
mock_embedder = AsyncMock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
splitter.sentence_splitter = AsyncMock()
with pytest.raises(ValueError, match="content for document ID"):
await splitter.run_async(documents=[Document(content=None)])
def test_run_empty_document(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
splitter.sentence_splitter = Mock()
result = splitter.run(documents=[Document(content="")])
assert result["documents"] == []
@pytest.mark.asyncio
async def test_run_empty_document_async(self) -> None:
mock_embedder = AsyncMock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
splitter.sentence_splitter = AsyncMock()
result = await splitter.run_async(documents=[Document(content="")])
assert result["documents"] == []
def test_group_sentences_single(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder, sentences_per_group=1)
sentences = ["Sentence 1.", "Sentence 2.", "Sentence 3."]
groups = splitter._group_sentences(sentences)
assert groups == sentences
def test_group_sentences_multiple(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder, sentences_per_group=2)
sentences = ["Sentence 1. ", "Sentence 2. ", "Sentence 3. ", "Sentence 4."]
groups = splitter._group_sentences(sentences)
assert groups == ["Sentence 1. Sentence 2. ", "Sentence 3. Sentence 4."]
def test_cosine_distance(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
# Test with identical vectors
embedding1 = [1.0, 0.0, 0.0]
embedding2 = [1.0, 0.0, 0.0]
distance = splitter._cosine_distance(embedding1, embedding2)
assert distance == 0.0
# Test with orthogonal vectors
embedding1 = [1.0, 0.0, 0.0]
embedding2 = [0.0, 1.0, 0.0]
distance = splitter._cosine_distance(embedding1, embedding2)
assert distance == 1.0
# Test with zero vectors
embedding1 = [0.0, 0.0, 0.0]
embedding2 = [1.0, 0.0, 0.0]
distance = splitter._cosine_distance(embedding1, embedding2)
assert distance == 1.0
def test_find_split_points_empty(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
split_points = splitter._find_split_points([])
assert split_points == []
split_points = splitter._find_split_points([[1.0, 0.0]])
assert split_points == []
def test_find_split_points(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder, percentile=0.5)
# Create embeddings where the second pair has high distance
embeddings = [
[1.0, 0.0, 0.0], # Similar to next
[0.9, 0.1, 0.0], # Similar to previous
[0.0, 1.0, 0.0], # Very different from next
[0.1, 0.9, 0.0], # Similar to previous
]
split_points = splitter._find_split_points(embeddings)
# Should find a split point after the second embedding (index 2)
assert 2 in split_points
def test_create_splits_from_points(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
sentence_groups = ["Group 1 ", "Group 2 ", "Group 3 ", "Group 4"]
split_points = [2] # Split after index 1
splits = splitter._create_splits_from_points(sentence_groups, split_points)
assert splits == ["Group 1 Group 2 ", "Group 3 Group 4"]
def test_create_splits_from_points_no_points(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
sentence_groups = ["Group 1 ", "Group 2 ", "Group 3"]
split_points = []
splits = splitter._create_splits_from_points(sentence_groups, split_points)
assert splits == ["Group 1 Group 2 Group 3"]
def test_merge_small_splits(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder, min_length=10)
splits = ["Short ", "Also short ", "Long enough text ", "Another short"]
merged = splitter._merge_small_splits(splits)
assert len(merged) == 3
assert merged[0] == "Short Also short "
assert merged[1] == "Long enough text "
assert merged[2] == "Another short"
def test_merge_small_splits_respect_max_length(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder, min_length=10, max_length=15)
splits = ["123456", "123456789", "1234"]
merged = splitter._merge_small_splits(splits=splits)
assert len(merged) == 2
# First split remains beneath min_length b/c next split is too long
assert merged[0] == "123456"
# Second split is merged with third split to get above min_length and still beneath max_length
assert merged[1] == "1234567891234"
def test_create_documents_from_splits(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
original_doc = Document(content="test", meta={"key": "value"})
splits = ["Split 1", "Split 2"]
documents = splitter._create_documents_from_splits(splits, original_doc)
assert len(documents) == 2
assert documents[0].content == "Split 1"
assert documents[0].meta["source_id"] == original_doc.id
assert documents[0].meta["split_id"] == 0
assert documents[0].meta["key"] == "value"
assert documents[1].content == "Split 2"
assert documents[1].meta["split_id"] == 1
def test_create_documents_from_splits_with_page_numbers(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
original_doc = Document(content="Page 1 content.\fPage 2 content.\f\fPage 4 content.", meta={"key": "value"})
splits = ["Page 1 content.\f", "Page 2 content.\f\f", "Page 4 content."]
documents = splitter._create_documents_from_splits(splits, original_doc)
assert len(documents) == 3
assert documents[0].content == "Page 1 content.\f"
assert documents[0].meta["page_number"] == 1
assert documents[1].content == "Page 2 content.\f\f"
assert documents[1].meta["page_number"] == 2
assert documents[2].content == "Page 4 content."
assert documents[2].meta["page_number"] == 4
def test_create_documents_from_splits_with_consecutive_page_breaks(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
# Test with consecutive page breaks at the end
original_doc = Document(content="Page 1 content.\fPage 2 content.\f\f\f", meta={"key": "value"})
splits = ["Page 1 content.\f", "Page 2 content.\f\f\f"]
documents = splitter._create_documents_from_splits(splits, original_doc)
assert len(documents) == 2
assert documents[0].content == "Page 1 content.\f"
assert documents[0].meta["page_number"] == 1
assert documents[1].content == "Page 2 content.\f\f\f"
# Should be page 2, not 4, because consecutive page breaks at the end are adjusted
assert documents[1].meta["page_number"] == 2
def test_calculate_embeddings(self):
mock_embedder = Mock()
# Mock the document embedder to return documents with embeddings
def mock_run(documents):
return {"documents": [replace(doc, embedding=[1.0, 2.0, 3.0]) for doc in documents]}
mock_embedder.run = Mock(side_effect=mock_run)
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
sentence_groups = ["Group 1", "Group 2", "Group 3"]
embeddings = splitter._calculate_embeddings(sentence_groups)
assert len(embeddings) == 3
assert all(embedding == [1.0, 2.0, 3.0] for embedding in embeddings)
mock_embedder.run.assert_called_once()
@pytest.mark.asyncio
async def test_calculate_embeddings_async(self) -> None:
mock_embedder = AsyncMock()
# Mock the document embedder to return documents with embeddings
async def mock_run_async(documents):
return {"documents": [replace(doc, embedding=[1.0, 2.0, 3.0]) for doc in documents]}
mock_embedder.run_async = AsyncMock(side_effect=mock_run_async)
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
sentence_groups = ["Group 1", "Group 2", "Group 3"]
embeddings = await splitter._calculate_embeddings_async(sentence_groups)
assert len(embeddings) == 3
assert all(embedding == [1.0, 2.0, 3.0] for embedding in embeddings)
mock_embedder.run_async.assert_called_once()
def test_to_dict(self):
mock_embedder = Mock()
mock_embedder.to_dict.return_value = {"type": "MockEmbedder"}
splitter = EmbeddingBasedDocumentSplitter(
document_embedder=mock_embedder, sentences_per_group=2, percentile=0.9, min_length=50, max_length=1000
)
result = splitter.to_dict()
assert "EmbeddingBasedDocumentSplitter" in result["type"]
assert result["init_parameters"]["sentences_per_group"] == 2
assert result["init_parameters"]["percentile"] == 0.9
assert result["init_parameters"]["min_length"] == 50
assert result["init_parameters"]["max_length"] == 1000
assert "document_embedder" in result["init_parameters"]
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
def test_split_document_with_multiple_topics(self):
embedder = OpenAIDocumentEmbedder(model="text-embedding-3-small")
splitter = EmbeddingBasedDocumentSplitter(
document_embedder=embedder, sentences_per_group=2, percentile=0.9, min_length=30, max_length=300
)
# A document with multiple topics
text = (
"The weather today is beautiful. The sun is shining brightly. The temperature is perfect for a walk. "
"Machine learning has revolutionized many industries. Neural networks can process vast amounts of data. "
"Deep learning models achieve remarkable accuracy on complex tasks. "
"Cooking is both an art and a science. Fresh ingredients make all the difference. "
"Proper seasoning enhances the natural flavors of food. "
"The history of ancient civilizations fascinates researchers. Archaeological discoveries reveal new insights. " # noqa: E501
"Ancient texts provide valuable information about past societies."
)
doc = Document(content=text)
result = splitter.run(documents=[doc])
split_docs = result["documents"]
# There should be more than one split
assert len(split_docs) > 1
# Each split should be non-empty and respect min_length
for split_doc in split_docs:
assert split_doc.content.strip() != ""
assert len(split_doc.content) >= 30
# The splits should cover the original text
combined = "".join([d.content for d in split_docs])
original = text
assert combined in original or original in combined
@pytest.mark.asyncio
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
async def test_split_document_with_multiple_topics_async(self) -> None:
embedder = OpenAIDocumentEmbedder(model="text-embedding-3-small")
splitter = EmbeddingBasedDocumentSplitter(
document_embedder=embedder, sentences_per_group=2, percentile=0.9, min_length=30, max_length=300
)
# A document with multiple topics
text = (
"The weather today is beautiful. The sun is shining brightly. The temperature is perfect for a walk. "
"Machine learning has revolutionized many industries. Neural networks can process vast amounts of data. "
"Deep learning models achieve remarkable accuracy on complex tasks. "
"Cooking is both an art and a science. Fresh ingredients make all the difference. "
"Proper seasoning enhances the natural flavors of food. "
"The history of ancient civilizations fascinates researchers. Archaeological discoveries reveal new insights. " # noqa: E501
"Ancient texts provide valuable information about past societies."
)
doc = Document(content=text)
result = await splitter.run_async(documents=[doc])
split_docs = result["documents"]
# There should be more than one split
assert len(split_docs) > 1
# Each split should be non-empty and respect min_length
for split_doc in split_docs:
assert split_doc.content.strip() != ""
assert len(split_doc.content) >= 30
# The splits should cover the original text
combined = "".join([d.content for d in split_docs])
original = text
assert combined in original or original in combined
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
def test_trailing_whitespace_is_preserved(self):
embedder = OpenAIDocumentEmbedder(model="text-embedding-3-small")
splitter = EmbeddingBasedDocumentSplitter(document_embedder=embedder, sentences_per_group=1)
# Normal trailing whitespace
text = "The weather today is beautiful. "
result = splitter.run(documents=[Document(content=text)])
assert result["documents"][0].content == text
# Newline at the end
text = "The weather today is beautiful.\n"
result = splitter.run(documents=[Document(content=text)])
assert result["documents"][0].content == text
# Page break at the end
text = "The weather today is beautiful.\f"
result = splitter.run(documents=[Document(content=text)])
assert result["documents"][0].content == text
@pytest.mark.asyncio
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
async def test_trailing_whitespace_is_preserved_async(self) -> None:
embedder = OpenAIDocumentEmbedder(model="text-embedding-3-small")
splitter = EmbeddingBasedDocumentSplitter(document_embedder=embedder, sentences_per_group=1)
# Normal trailing whitespace
text = "The weather today is beautiful. "
result = await splitter.run_async(documents=[Document(content=text)])
assert result["documents"][0].content == text
# Newline at the end
text = "The weather today is beautiful.\n"
result = await splitter.run_async(documents=[Document(content=text)])
assert result["documents"][0].content == text
# Page break at the end
text = "The weather today is beautiful.\f"
result = await splitter.run_async(documents=[Document(content=text)])
assert result["documents"][0].content == text
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
def test_no_extra_whitespaces_between_sentences(self):
embedder = OpenAIDocumentEmbedder(model="text-embedding-3-small")
splitter = EmbeddingBasedDocumentSplitter(
document_embedder=embedder, sentences_per_group=1, percentile=0.9, min_length=10, max_length=500
)
text = (
"The weather today is beautiful. The sun is shining brightly. The temperature is perfect for a walk. "
"There are no clouds and no rain. Machine learning has revolutionized many industries. "
"Neural networks can process vast amounts of data. Deep learning models achieve remarkable accuracy on complex tasks." # noqa: E501
)
doc = Document(content=text)
result = splitter.run(documents=[doc])
split_docs = result["documents"]
assert len(split_docs) == 2
# Expect the original whitespace structure with trailing spaces where they exist
assert (
split_docs[0].content
== "The weather today is beautiful. The sun is shining brightly. The temperature is perfect for a walk. There are no clouds and no rain. " # noqa: E501
) # noqa: E501
assert (
split_docs[1].content
== "Machine learning has revolutionized many industries. Neural networks can process vast amounts of data. Deep learning models achieve remarkable accuracy on complex tasks." # noqa: E501
) # noqa: E501
@pytest.mark.asyncio
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
async def test_no_extra_whitespaces_between_sentences_async(self) -> None:
embedder = OpenAIDocumentEmbedder(model="text-embedding-3-small")
splitter = EmbeddingBasedDocumentSplitter(
document_embedder=embedder, sentences_per_group=1, percentile=0.9, min_length=10, max_length=500
)
text = (
"The weather today is beautiful. The sun is shining brightly. The temperature is perfect for a walk. "
"There are no clouds and no rain. Machine learning has revolutionized many industries. "
"Neural networks can process vast amounts of data. Deep learning models achieve remarkable accuracy on complex tasks." # noqa: E501
)
doc = Document(content=text)
result = await splitter.run_async(documents=[doc])
split_docs = result["documents"]
assert len(split_docs) == 2
# Expect the original whitespace structure with trailing spaces where they exist
assert (
split_docs[0].content
== "The weather today is beautiful. The sun is shining brightly. The temperature is perfect for a walk. There are no clouds and no rain. " # noqa: E501
) # noqa: E501
assert (
split_docs[1].content
== "Machine learning has revolutionized many industries. Neural networks can process vast amounts of data. Deep learning models achieve remarkable accuracy on complex tasks." # noqa: E501
) # noqa: E501
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
def test_split_large_splits_recursion(self):
"""
Test that _split_large_splits() works correctly without infinite loops.
This test uses a longer text that will trigger the recursive splitting logic.
If the chunk cannot be split further, it is allowed to be larger than max_length.
"""
embedder = OpenAIDocumentEmbedder(model="text-embedding-3-small")
semantic_chunker = EmbeddingBasedDocumentSplitter(
document_embedder=embedder, sentences_per_group=5, percentile=0.95, min_length=50, max_length=1000
)
text = """# Artificial intelligence and its Impact on Society
## Article from Wikipedia, the free encyclopedia
### Introduction to Artificial Intelligence
Artificial intelligence (AI) is the capability of computational systems to perform tasks typically associated with human intelligence, such as learning, reasoning, problem-solving, perception, and decision-making. It is a field of research in computer science that develops and studies methods and software that enable machines to perceive their environment and use learning and intelligence to take actions that maximize their chances of achieving defined goals.
### The History of Software
The history of software is closely tied to the development of digital computers in the mid-20th century. Early programs were written in the machine language specific to the hardware. The introduction of high-level programming languages in 1958 allowed for more human-readable instructions, making software development easier and more portable across different computer architectures. Software in a programming language is run through a compiler or interpreter to execute on the architecture's hardware. Over time, software has become complex, owing to developments in networking, operating systems, and databases.""" # noqa: E501
doc = Document(content=text)
result = semantic_chunker.run(documents=[doc])
split_docs = result["documents"]
assert len(split_docs) == 1
# If the chunk cannot be split further, it is allowed to be larger than max_length
# At least one split should be larger than max_length in this test case
assert any(len(split_doc.content) > 1000 for split_doc in split_docs)
# Verify that the splits cover the original content
combined_content = "".join([d.content for d in split_docs])
assert combined_content == text
for i, split_doc in enumerate(split_docs):
assert split_doc.meta["source_id"] == doc.id
assert split_doc.meta["split_id"] == i
assert "page_number" in split_doc.meta
@pytest.mark.asyncio
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
async def test_split_large_splits_recursion_async(self) -> None:
"""
Test that _split_large_splits() works correctly without infinite loops.
This test uses a longer text that will trigger the recursive splitting logic.
If the chunk cannot be split further, it is allowed to be larger than max_length.
"""
embedder = OpenAIDocumentEmbedder(model="text-embedding-3-small")
semantic_chunker = EmbeddingBasedDocumentSplitter(
document_embedder=embedder, sentences_per_group=5, percentile=0.95, min_length=50, max_length=1000
)
text = """# Artificial intelligence and its Impact on Society
## Article from Wikipedia, the free encyclopedia
### Introduction to Artificial Intelligence
Artificial intelligence (AI) is the capability of computational systems to perform tasks typically associated with human intelligence, such as learning, reasoning, problem-solving, perception, and decision-making. It is a field of research in computer science that develops and studies methods and software that enable machines to perceive their environment and use learning and intelligence to take actions that maximize their chances of achieving defined goals.
### The History of Software
The history of software is closely tied to the development of digital computers in the mid-20th century. Early programs were written in the machine language specific to the hardware. The introduction of high-level programming languages in 1958 allowed for more human-readable instructions, making software development easier and more portable across different computer architectures. Software in a programming language is run through a compiler or interpreter to execute on the architecture's hardware. Over time, software has become complex, owing to developments in networking, operating systems, and databases.""" # noqa: E501
doc = Document(content=text)
result = await semantic_chunker.run_async(documents=[doc])
split_docs = result["documents"]
assert len(split_docs) == 1
# If the chunk cannot be split further, it is allowed to be larger than max_length
# At least one split should be larger than max_length in this test case
assert any(len(split_doc.content) > 1000 for split_doc in split_docs)
# Verify that the splits cover the original content
combined_content = "".join([d.content for d in split_docs])
assert combined_content == text
for i, split_doc in enumerate(split_docs):
assert split_doc.meta["source_id"] == doc.id
assert split_doc.meta["split_id"] == i
assert "page_number" in split_doc.meta
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
def test_split_large_splits_actually_splits(self):
"""
Test that _split_large_splits() actually works and can split long texts into multiple chunks.
This test uses a very long text that should be split into multiple chunks.
"""
embedder = OpenAIDocumentEmbedder(model="text-embedding-3-small")
semantic_chunker = EmbeddingBasedDocumentSplitter(
document_embedder=embedder,
sentences_per_group=3,
percentile=0.85, # Lower percentile to create more splits
min_length=100,
max_length=500, # Smaller max_length to force more splits
)
# Create a very long text with multiple paragraphs and topics
text = """# Comprehensive Guide to Machine Learning and Artificial Intelligence
## Introduction to Machine Learning
Machine learning is a subset of artificial intelligence that focuses on the development of computer programs that can access data and use it to learn for themselves. The process of learning begins with observations or data, such as examples, direct experience, or instruction, in order to look for patterns in data and make better decisions in the future based on the examples that we provide. The primary aim is to allow the computers learn automatically without human intervention or assistance and adjust actions accordingly.
## Types of Machine Learning
There are several types of machine learning algorithms, each with their own strengths and weaknesses. Supervised learning involves training a model on a labeled dataset, where the correct answers are provided. The model learns to map inputs to outputs based on these examples. Unsupervised learning, on the other hand, deals with unlabeled data and seeks to find hidden patterns or structures within the data. Reinforcement learning is a type of learning where an agent learns to behave in an environment by performing certain actions and receiving rewards or penalties.
## Deep Learning and Neural Networks
Deep learning is a subset of machine learning that uses neural networks with multiple layers to model and understand complex patterns. Neural networks are inspired by the human brain and consist of interconnected nodes or neurons. Each connection between neurons has a weight that is adjusted during training. The network learns by adjusting these weights based on the error between predicted and actual outputs. Deep learning has been particularly successful in areas such as computer vision, natural language processing, and speech recognition.
\f
## Natural Language Processing
Natural Language Processing (NLP) is a field of artificial intelligence that focuses on the interaction between computers and human language. It involves developing algorithms and models that can understand, interpret, and generate human language. NLP applications include machine translation, sentiment analysis, text summarization, and question answering systems. Recent advances in deep learning have significantly improved the performance of NLP systems, leading to more accurate and sophisticated language models.
## Computer Vision and Image Recognition
Computer vision is another important area of artificial intelligence that deals with how computers can gain high-level understanding from digital images or videos. It involves developing algorithms that can identify and understand visual information from the world. Applications include facial recognition, object detection, medical image analysis, and autonomous vehicle navigation. Deep learning models, particularly convolutional neural networks (CNNs), have revolutionized computer vision by achieving human-level performance on many tasks.
## The Future of Artificial Intelligence
The future of artificial intelligence holds immense potential for transforming various industries and aspects of human life. We can expect to see more sophisticated AI systems that can handle complex reasoning tasks, understand context better, and interact more naturally with humans. However, this rapid advancement also brings challenges related to ethics, privacy, and the impact on employment. It's crucial to develop AI systems that are not only powerful but also safe, fair, and beneficial to society as a whole.
\f
## Ethical Considerations in AI
As artificial intelligence becomes more prevalent, ethical considerations become increasingly important. Issues such as bias in AI systems, privacy concerns, and the potential for misuse need to be carefully addressed. AI systems can inherit biases from their training data, leading to unfair outcomes for certain groups. Privacy concerns arise from the vast amounts of data required to train AI systems. Additionally, there are concerns about the potential for AI to be used maliciously or to replace human workers in certain industries.
## Applications in Healthcare
Artificial intelligence has the potential to revolutionize healthcare by improving diagnosis, treatment planning, and patient care. Machine learning algorithms can analyze medical images to detect diseases earlier and more accurately than human doctors. AI systems can also help in drug discovery by predicting the effectiveness of potential treatments. In addition, AI-powered chatbots and virtual assistants can provide basic healthcare information and support to patients, reducing the burden on healthcare professionals.
## AI in Finance and Banking
The financial industry has been quick to adopt artificial intelligence for various applications. AI systems can analyze market data to make investment decisions, detect fraudulent transactions, and provide personalized financial advice. Machine learning algorithms can assess credit risk more accurately than traditional methods, leading to better lending decisions. Additionally, AI-powered chatbots can handle customer service inquiries, reducing costs and improving customer satisfaction.
\f
## Transportation and Autonomous Vehicles
Autonomous vehicles represent one of the most visible applications of artificial intelligence in transportation. Self-driving cars use a combination of sensors, cameras, and AI algorithms to navigate roads safely. These systems can detect obstacles, read traffic signs, and make decisions about speed and direction. Beyond autonomous cars, AI is also being used in logistics and supply chain management to optimize routes and reduce delivery times.
## Education and Personalized Learning
Artificial intelligence is transforming education by enabling personalized learning experiences. AI systems can adapt to individual student needs, providing customized content and pacing. Intelligent tutoring systems can provide immediate feedback and support to students, helping them learn more effectively. Additionally, AI can help educators by automating administrative tasks and providing insights into student performance and learning patterns.""" # noqa: E501
doc = Document(content=text)
result = semantic_chunker.run(documents=[doc])
split_docs = result["documents"]
assert len(split_docs) == 11
# Verify that the splits cover the original content
combined_content = "".join([d.content for d in split_docs])
assert combined_content == text
for i, split_doc in enumerate(split_docs):
assert split_doc.meta["source_id"] == doc.id
assert split_doc.meta["split_id"] == i
assert "page_number" in split_doc.meta
if i in [0, 1, 2, 3]:
assert split_doc.meta["page_number"] == 1
if i in [4, 5, 6]:
assert split_doc.meta["page_number"] == 2
if i in [7, 8]:
assert split_doc.meta["page_number"] == 3
if i in [9, 10]:
assert split_doc.meta["page_number"] == 4
@pytest.mark.asyncio
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
async def test_split_large_splits_actually_splits_async(self) -> None:
"""
Test that _split_large_splits() actually works and can split long texts into multiple chunks.
This test uses a very long text that should be split into multiple chunks.
"""
embedder = OpenAIDocumentEmbedder(model="text-embedding-3-small")
semantic_chunker = EmbeddingBasedDocumentSplitter(
document_embedder=embedder,
sentences_per_group=3,
percentile=0.85, # Lower percentile to create more splits
min_length=100,
max_length=500, # Smaller max_length to force more splits
)
# Create a very long text with multiple paragraphs and topics
text = """# Comprehensive Guide to Machine Learning and Artificial Intelligence
## Introduction to Machine Learning
Machine learning is a subset of artificial intelligence that focuses on the development of computer programs that can access data and use it to learn for themselves. The process of learning begins with observations or data, such as examples, direct experience, or instruction, in order to look for patterns in data and make better decisions in the future based on the examples that we provide. The primary aim is to allow the computers learn automatically without human intervention or assistance and adjust actions accordingly.
## Types of Machine Learning
There are several types of machine learning algorithms, each with their own strengths and weaknesses. Supervised learning involves training a model on a labeled dataset, where the correct answers are provided. The model learns to map inputs to outputs based on these examples. Unsupervised learning, on the other hand, deals with unlabeled data and seeks to find hidden patterns or structures within the data. Reinforcement learning is a type of learning where an agent learns to behave in an environment by performing certain actions and receiving rewards or penalties.
## Deep Learning and Neural Networks
Deep learning is a subset of machine learning that uses neural networks with multiple layers to model and understand complex patterns. Neural networks are inspired by the human brain and consist of interconnected nodes or neurons. Each connection between neurons has a weight that is adjusted during training. The network learns by adjusting these weights based on the error between predicted and actual outputs. Deep learning has been particularly successful in areas such as computer vision, natural language processing, and speech recognition.
\f
## Natural Language Processing
Natural Language Processing (NLP) is a field of artificial intelligence that focuses on the interaction between computers and human language. It involves developing algorithms and models that can understand, interpret, and generate human language. NLP applications include machine translation, sentiment analysis, text summarization, and question answering systems. Recent advances in deep learning have significantly improved the performance of NLP systems, leading to more accurate and sophisticated language models.
## Computer Vision and Image Recognition
Computer vision is another important area of artificial intelligence that deals with how computers can gain high-level understanding from digital images or videos. It involves developing algorithms that can identify and understand visual information from the world. Applications include facial recognition, object detection, medical image analysis, and autonomous vehicle navigation. Deep learning models, particularly convolutional neural networks (CNNs), have revolutionized computer vision by achieving human-level performance on many tasks.
## The Future of Artificial Intelligence
The future of artificial intelligence holds immense potential for transforming various industries and aspects of human life. We can expect to see more sophisticated AI systems that can handle complex reasoning tasks, understand context better, and interact more naturally with humans. However, this rapid advancement also brings challenges related to ethics, privacy, and the impact on employment. It's crucial to develop AI systems that are not only powerful but also safe, fair, and beneficial to society as a whole.
\f
## Ethical Considerations in AI
As artificial intelligence becomes more prevalent, ethical considerations become increasingly important. Issues such as bias in AI systems, privacy concerns, and the potential for misuse need to be carefully addressed. AI systems can inherit biases from their training data, leading to unfair outcomes for certain groups. Privacy concerns arise from the vast amounts of data required to train AI systems. Additionally, there are concerns about the potential for AI to be used maliciously or to replace human workers in certain industries.
## Applications in Healthcare
Artificial intelligence has the potential to revolutionize healthcare by improving diagnosis, treatment planning, and patient care. Machine learning algorithms can analyze medical images to detect diseases earlier and more accurately than human doctors. AI systems can also help in drug discovery by predicting the effectiveness of potential treatments. In addition, AI-powered chatbots and virtual assistants can provide basic healthcare information and support to patients, reducing the burden on healthcare professionals.
## AI in Finance and Banking
The financial industry has been quick to adopt artificial intelligence for various applications. AI systems can analyze market data to make investment decisions, detect fraudulent transactions, and provide personalized financial advice. Machine learning algorithms can assess credit risk more accurately than traditional methods, leading to better lending decisions. Additionally, AI-powered chatbots can handle customer service inquiries, reducing costs and improving customer satisfaction.
\f
## Transportation and Autonomous Vehicles
Autonomous vehicles represent one of the most visible applications of artificial intelligence in transportation. Self-driving cars use a combination of sensors, cameras, and AI algorithms to navigate roads safely. These systems can detect obstacles, read traffic signs, and make decisions about speed and direction. Beyond autonomous cars, AI is also being used in logistics and supply chain management to optimize routes and reduce delivery times.
## Education and Personalized Learning
Artificial intelligence is transforming education by enabling personalized learning experiences. AI systems can adapt to individual student needs, providing customized content and pacing. Intelligent tutoring systems can provide immediate feedback and support to students, helping them learn more effectively. Additionally, AI can help educators by automating administrative tasks and providing insights into student performance and learning patterns.""" # noqa: E501
doc = Document(content=text)
result = await semantic_chunker.run_async(documents=[doc])
split_docs = result["documents"]
assert len(split_docs) == 11
# Verify that the splits cover the original content
combined_content = "".join([d.content for d in split_docs])
assert combined_content == text
for i, split_doc in enumerate(split_docs):
assert split_doc.meta["source_id"] == doc.id
assert split_doc.meta["split_id"] == i
assert "page_number" in split_doc.meta
if i in [0, 1, 2, 3]:
assert split_doc.meta["page_number"] == 1
if i in [4, 5, 6]:
assert split_doc.meta["page_number"] == 2
if i in [7, 8]:
assert split_doc.meta["page_number"] == 3
if i in [9, 10]:
assert split_doc.meta["page_number"] == 4
class TestComponentLifecycle:
def test_warm_up_builds_splitter_and_delegates_to_embedder(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
with patch(
"haystack.components.preprocessors.embedding_based_document_splitter.SentenceSplitter"
) as mock_splitter_class:
splitter.warm_up()
assert splitter.sentence_splitter is mock_splitter_class.return_value
mock_splitter_class.assert_called_once()
mock_embedder.warm_up.assert_called_once()
def test_warm_up_builds_splitter_once(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
with patch(
"haystack.components.preprocessors.embedding_based_document_splitter.SentenceSplitter"
) as mock_splitter_class:
splitter.warm_up()
first_splitter = splitter.sentence_splitter
splitter.warm_up()
mock_splitter_class.assert_called_once()
assert splitter.sentence_splitter is first_splitter
@pytest.mark.asyncio
async def test_warm_up_async_delegates_to_embedder_async(self) -> None:
mock_embedder = Mock()
mock_embedder.warm_up_async = AsyncMock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
with patch("haystack.components.preprocessors.embedding_based_document_splitter.SentenceSplitter"):
await splitter.warm_up_async()
mock_embedder.warm_up_async.assert_awaited_once()
@pytest.mark.asyncio
async def test_warm_up_async_falls_back_to_sync_warm_up(self) -> None:
mock_embedder = Mock(spec=["warm_up"])
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
with patch("haystack.components.preprocessors.embedding_based_document_splitter.SentenceSplitter"):
await splitter.warm_up_async()
mock_embedder.warm_up.assert_called_once()
def test_close_delegates_to_embedder(self):
mock_embedder = Mock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
splitter.close()
mock_embedder.close.assert_called_once()
@pytest.mark.asyncio
async def test_close_async_delegates_to_embedder(self) -> None:
mock_embedder = Mock()
mock_embedder.close_async = AsyncMock()
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
await splitter.close_async()
mock_embedder.close_async.assert_awaited_once()
@pytest.mark.asyncio
async def test_close_async_falls_back_to_sync_close(self) -> None:
mock_embedder = Mock(spec=["close"])
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
await splitter.close_async()
mock_embedder.close.assert_called_once()
def test_lifecycle_is_safe_when_embedder_lacks_methods(self):
mock_embedder = Mock(spec=[])
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
with patch("haystack.components.preprocessors.embedding_based_document_splitter.SentenceSplitter"):
splitter.warm_up()
splitter.close()
@pytest.mark.asyncio
async def test_lifecycle_is_safe_when_embedder_lacks_methods_async(self) -> None:
mock_embedder = Mock(spec=[])
splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
with patch("haystack.components.preprocessors.embedding_based_document_splitter.SentenceSplitter"):
await splitter.warm_up_async()
await splitter.close_async()
@@ -0,0 +1,269 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
import pytest
from haystack import Document, Pipeline
from haystack.components.preprocessors import HierarchicalDocumentSplitter
from haystack.components.writers import DocumentWriter
class TestHierarchicalDocumentSplitter:
def test_init_with_default_params(self):
builder = HierarchicalDocumentSplitter(block_sizes={100, 200, 300})
assert builder.block_sizes == [300, 200, 100]
assert builder.split_overlap == 0
assert builder.split_by == "word"
def test_init_with_custom_params(self):
builder = HierarchicalDocumentSplitter(block_sizes={100, 200, 300}, split_overlap=25, split_by="word")
assert builder.block_sizes == [300, 200, 100]
assert builder.split_overlap == 25
assert builder.split_by == "word"
def test_init_with_empty_block_sizes_raises(self):
with pytest.raises(ValueError, match="block_sizes must not be empty"):
HierarchicalDocumentSplitter(block_sizes=set())
def test_init_with_negative_split_overlap_raises(self):
with pytest.raises(ValueError, match="split_overlap must be greater than or equal to 0"):
HierarchicalDocumentSplitter(block_sizes={10, 5, 2}, split_overlap=-1)
def test_init_with_split_overlap_equal_to_smallest_block_size_raises(self):
with pytest.raises(ValueError, match="split_overlap .* must be less than the smallest value in block_sizes"):
HierarchicalDocumentSplitter(block_sizes={10, 5, 2}, split_overlap=2)
def test_init_with_split_overlap_greater_than_smallest_block_size_raises(self):
with pytest.raises(ValueError, match="split_overlap .* must be less than the smallest value in block_sizes"):
HierarchicalDocumentSplitter(block_sizes={10, 5, 2}, split_overlap=3)
def test_to_dict(self):
builder = HierarchicalDocumentSplitter(block_sizes={100, 200, 300}, split_overlap=25, split_by="word")
expected = builder.to_dict()
assert expected == {
"type": "haystack.components.preprocessors.hierarchical_document_splitter.HierarchicalDocumentSplitter",
"init_parameters": {"block_sizes": [300, 200, 100], "split_overlap": 25, "split_by": "word"},
}
def test_from_dict(self):
data = {
"type": "haystack.components.preprocessors.hierarchical_document_splitter.HierarchicalDocumentSplitter",
"init_parameters": {"block_sizes": [10, 5, 2], "split_overlap": 0, "split_by": "word"},
}
builder = HierarchicalDocumentSplitter.from_dict(data)
assert builder.block_sizes == [10, 5, 2]
assert builder.split_overlap == 0
assert builder.split_by == "word"
def test_run(self):
builder = HierarchicalDocumentSplitter(block_sizes={10, 5, 2}, split_overlap=0, split_by="word")
text = "one two three four five six seven eight nine ten"
doc = Document(content=text)
output = builder.run([doc])
docs = output["documents"]
builder.run([doc])
assert len(docs) == 9
assert docs[0].content == "one two three four five six seven eight nine ten"
# level 1 - root node
assert docs[0].meta["__level"] == 0
assert len(docs[0].meta["__children_ids"]) == 2
# level 2 -left branch
assert docs[1].meta["__parent_id"] == docs[0].id
assert docs[1].meta["__level"] == 1
assert len(docs[1].meta["__children_ids"]) == 3
# level 2 - right branch
assert docs[2].meta["__parent_id"] == docs[0].id
assert docs[2].meta["__level"] == 1
assert len(docs[2].meta["__children_ids"]) == 3
# level 3 - left branch - leaf nodes
assert docs[3].meta["__parent_id"] == docs[1].id
assert docs[4].meta["__parent_id"] == docs[1].id
assert docs[5].meta["__parent_id"] == docs[1].id
assert docs[3].meta["__level"] == 2
assert docs[4].meta["__level"] == 2
assert docs[5].meta["__level"] == 2
assert len(docs[3].meta["__children_ids"]) == 0
assert len(docs[4].meta["__children_ids"]) == 0
assert len(docs[5].meta["__children_ids"]) == 0
# level 3 - right branch - leaf nodes
assert docs[6].meta["__parent_id"] == docs[2].id
assert docs[7].meta["__parent_id"] == docs[2].id
assert docs[8].meta["__parent_id"] == docs[2].id
assert docs[6].meta["__level"] == 2
assert docs[7].meta["__level"] == 2
assert docs[8].meta["__level"] == 2
assert len(docs[6].meta["__children_ids"]) == 0
assert len(docs[7].meta["__children_ids"]) == 0
assert len(docs[8].meta["__children_ids"]) == 0
def test_run_does_not_mutate_input_document(self):
builder = HierarchicalDocumentSplitter(block_sizes={5, 2}, split_overlap=0, split_by="word")
doc = Document(content="one two three four five six seven eight", meta={"source": "test"})
builder.run([doc])
# the caller's Document should be untouched
assert doc.meta == {"source": "test"}
for key in ("__block_size", "__parent_id", "__children_ids", "__level"):
assert key not in doc.meta
def test_to_dict_in_pipeline(self, in_memory_doc_store):
pipeline = Pipeline()
hierarchical_doc_builder = HierarchicalDocumentSplitter(block_sizes={10, 5, 2})
doc_writer = DocumentWriter(document_store=in_memory_doc_store)
pipeline.add_component(name="hierarchical_doc_splitter", instance=hierarchical_doc_builder)
pipeline.add_component(name="doc_writer", instance=doc_writer)
pipeline.connect("hierarchical_doc_splitter", "doc_writer")
expected = pipeline.to_dict()
assert expected.keys() == {
"connections",
"connection_type_validation",
"components",
"max_runs_per_component",
"metadata",
}
assert expected["components"].keys() == {"hierarchical_doc_splitter", "doc_writer"}
assert expected["components"]["hierarchical_doc_splitter"] == {
"type": "haystack.components.preprocessors.hierarchical_document_splitter.HierarchicalDocumentSplitter",
"init_parameters": {"block_sizes": [10, 5, 2], "split_overlap": 0, "split_by": "word"},
}
def test_from_dict_in_pipeline(self):
data = {
"metadata": {},
"max_runs_per_component": 100,
"components": {
"hierarchical_document_splitter": {
"type": "haystack.components.preprocessors.hierarchical_document_splitter.HierarchicalDocumentSplitter", # noqa: E501
"init_parameters": {"block_sizes": [10, 5, 2], "split_overlap": 0, "split_by": "word"},
},
"doc_writer": {
"type": "haystack.components.writers.document_writer.DocumentWriter",
"init_parameters": {
"document_store": {
"type": "haystack.document_stores.in_memory.document_store.InMemoryDocumentStore",
"init_parameters": {
"bm25_tokenization_regex": "(?u)\\b\\w\\w+\\b",
"bm25_algorithm": "BM25L",
"bm25_parameters": {},
"embedding_similarity_function": "dot_product",
"index": "f32ad5bf-43cb-4035-9823-1de1ae9853c1",
"shared": True,
},
},
"policy": "NONE",
},
},
},
"connections": [{"sender": "hierarchical_document_splitter.documents", "receiver": "doc_writer.documents"}],
}
assert Pipeline.from_dict(data)
@pytest.mark.integration
def test_example_in_pipeline(self, in_memory_doc_store):
pipeline = Pipeline()
hierarchical_doc_builder = HierarchicalDocumentSplitter(
block_sizes={10, 5, 2}, split_overlap=0, split_by="word"
)
doc_writer = DocumentWriter(document_store=in_memory_doc_store)
pipeline.add_component(name="hierarchical_doc_splitter", instance=hierarchical_doc_builder)
pipeline.add_component(name="doc_writer", instance=doc_writer)
pipeline.connect("hierarchical_doc_splitter.documents", "doc_writer")
text = "one two three four five six seven eight nine ten"
doc = Document(content=text)
docs = pipeline.run({"hierarchical_doc_splitter": {"documents": [doc]}})
assert docs["doc_writer"]["documents_written"] == 9
assert len(in_memory_doc_store.storage.values()) == 9
def test_serialization_deserialization_pipeline(self, in_memory_doc_store):
pipeline = Pipeline()
hierarchical_doc_builder = HierarchicalDocumentSplitter(
block_sizes={10, 5, 2}, split_overlap=0, split_by="word"
)
doc_writer = DocumentWriter(document_store=in_memory_doc_store)
pipeline.add_component(name="hierarchical_doc_splitter", instance=hierarchical_doc_builder)
pipeline.add_component(name="doc_writer", instance=doc_writer)
pipeline.connect("hierarchical_doc_splitter.documents", "doc_writer")
pipeline_dict = pipeline.to_dict()
new_pipeline = Pipeline.from_dict(pipeline_dict)
assert new_pipeline == pipeline
def test_split_by_sentence_assure_warm_up_was_called(self, in_memory_doc_store):
pipeline = Pipeline()
hierarchical_doc_builder = HierarchicalDocumentSplitter(
block_sizes={10, 5, 2}, split_overlap=0, split_by="sentence"
)
doc_writer = DocumentWriter(document_store=in_memory_doc_store)
pipeline.add_component(name="hierarchical_doc_splitter", instance=hierarchical_doc_builder)
pipeline.add_component(name="doc_writer", instance=doc_writer)
pipeline.connect("hierarchical_doc_splitter.documents", "doc_writer")
text = "This is one sentence. This is another sentence. This is the third sentence."
doc = Document(content=text)
docs = pipeline.run({"hierarchical_doc_splitter": {"documents": [doc]}})
assert docs["doc_writer"]["documents_written"] == 3
assert len(in_memory_doc_store.storage.values()) == 3
def test_hierarchical_splitter_multiple_block_sizes(self):
# Test with three different block sizes
doc = Document(
content="This is a simple test document with multiple sentences. It should be split into various sizes. "
"This helps test the hierarchy."
)
# Using three block sizes: 10, 5, 2 words
splitter = HierarchicalDocumentSplitter(block_sizes={10, 5, 2}, split_overlap=0, split_by="word")
result = splitter.run([doc])
documents = result["documents"]
# Verify root document
assert len(documents) > 1
root = documents[0]
assert root.meta["__level"] == 0
assert root.meta["__parent_id"] is None
# Verify level 1 documents (block_size=10)
level_1_docs = [d for d in documents if d.meta["__level"] == 1]
for doc in level_1_docs:
assert doc.meta["__block_size"] == 10
assert doc.meta["__parent_id"] == root.id
# Verify level 2 documents (block_size=5)
level_2_docs = [d for d in documents if d.meta["__level"] == 2]
for doc in level_2_docs:
assert doc.meta["__block_size"] == 5
assert doc.meta["__parent_id"] in [d.id for d in level_1_docs]
# Verify level 3 documents (block_size=2)
level_3_docs = [d for d in documents if d.meta["__level"] == 3]
for doc in level_3_docs:
assert doc.meta["__block_size"] == 2
assert doc.meta["__parent_id"] in [d.id for d in level_2_docs]
# Verify children references
for doc in documents:
if doc.meta["__children_ids"]:
child_ids = doc.meta["__children_ids"]
children = [d for d in documents if d.id in child_ids]
for child in children:
assert child.meta["__parent_id"] == doc.id
assert child.meta["__level"] == doc.meta["__level"] + 1
@@ -0,0 +1,840 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
from collections import defaultdict
from unittest.mock import ANY
import pytest
from haystack import Document
from haystack.components.preprocessors.markdown_header_splitter import MarkdownHeaderSplitter
# Fixtures
@pytest.fixture
def sample_text():
return (
"# Header 1\n"
"Content under header 1.\n"
"## Header 1.1\n"
"### Subheader 1.1.1\n"
"Content under sub-header 1.1.1\n"
"## Header 1.2\n"
"### Subheader 1.2.1\n"
"Content under header 1.2.1.\n"
"### Subheader 1.2.2\n"
"Content under header 1.2.2.\n"
"### Subheader 1.2.3\n"
"Content under header 1.2.3."
)
@pytest.fixture
def sample_text_with_page_breaks():
return (
"# Header 1\n"
"Content under header 1.\n\f\n"
"## Header 1.1\n"
"### Subheader 1.1.1\n"
"Content under sub-header 1.1.1\n\f\n"
"## Header 1.2\n"
"### Subheader 1.2.1\n"
"Content under header 1.2.1.\n\f\n"
"### Subheader 1.2.2\n"
"Content under header 1.2.2.\n\f\n"
"### Subheader 1.2.3\n"
"Content under header 1.2.3."
)
# Basic splitting and structure
def test_basic_split(sample_text):
splitter = MarkdownHeaderSplitter()
docs = [Document(content=sample_text)]
result = splitter.run(documents=docs)
split_docs = result["documents"]
# Check that content is present and correct
# Test first split
header1_doc = split_docs[0]
assert header1_doc.meta["split_id"] == 0
assert header1_doc.meta["page_number"] == 1
assert header1_doc.content == "# Header 1\nContent under header 1.\n"
# Test second split
subheader111_doc = split_docs[1]
assert subheader111_doc.meta["split_id"] == 1
assert subheader111_doc.meta["page_number"] == 1
assert subheader111_doc.content == "## Header 1.1\n### Subheader 1.1.1\nContent under sub-header 1.1.1\n"
# Test third split
subheader121_doc = split_docs[2]
assert subheader121_doc.meta["split_id"] == 2
assert subheader121_doc.meta["page_number"] == 1
assert subheader121_doc.content == "## Header 1.2\n### Subheader 1.2.1\nContent under header 1.2.1.\n"
# Test fourth split
subheader122_doc = split_docs[3]
assert subheader122_doc.meta["split_id"] == 3
assert subheader122_doc.meta["page_number"] == 1
assert subheader122_doc.content == "### Subheader 1.2.2\nContent under header 1.2.2.\n"
# Test fifth split
subheader123_doc = split_docs[4]
assert subheader123_doc.meta["split_id"] == 4
assert subheader123_doc.meta["page_number"] == 1
assert subheader123_doc.content == "### Subheader 1.2.3\nContent under header 1.2.3."
# Reconstruct original text
reconstructed_doc = "".join([doc.content for doc in split_docs])
assert reconstructed_doc == sample_text
def test_keep_headers_preserves_parent_headers_for_first_child():
text = (
"# Header 1\n"
"Intro text\n\n"
"## Header 1.1\n"
"Text 1\n\n"
"## Header 1.2\n"
"Text 2\n\n"
"### Header 1.2.1\n"
"Text 3\n\n"
"### Header 1.2.2\n"
"Text 4\n"
)
splitter = MarkdownHeaderSplitter(keep_headers=True)
split_docs = splitter.run(documents=[Document(content=text)])["documents"]
assert [(doc.meta["header"], doc.meta["parent_headers"]) for doc in split_docs] == [
("Header 1", []),
("Header 1.1", ["Header 1"]),
("Header 1.2", ["Header 1"]),
("Header 1.2.1", ["Header 1", "Header 1.2"]),
("Header 1.2.2", ["Header 1", "Header 1.2"]),
]
# reconstruct original text
reconstructed_text = "".join(doc.content for doc in split_docs)
assert reconstructed_text == text
def test_split_without_headers(sample_text):
splitter = MarkdownHeaderSplitter(keep_headers=False)
docs = [Document(content=sample_text)]
result = splitter.run(documents=docs)
split_docs = result["documents"]
# Should split into all headers with content
headers = [doc.meta["header"] for doc in split_docs]
assert "Header 1" in headers
assert "Subheader 1.1.1" in headers
assert "Subheader 1.2.1" in headers
assert "Subheader 1.2.2" in headers
assert "Subheader 1.2.3" in headers
# Check that content is present and correct
# Test first split
header1_doc = split_docs[0]
assert header1_doc.meta["header"] == "Header 1"
assert header1_doc.meta["split_id"] == 0
assert header1_doc.meta["page_number"] == 1
assert header1_doc.meta["parent_headers"] == []
assert header1_doc.content == "\nContent under header 1.\n"
# Test second split
subheader111_doc = split_docs[1]
assert subheader111_doc.meta["header"] == "Subheader 1.1.1"
assert subheader111_doc.meta["split_id"] == 1
assert subheader111_doc.meta["page_number"] == 1
assert subheader111_doc.meta["parent_headers"] == ["Header 1", "Header 1.1"]
assert subheader111_doc.content == "\nContent under sub-header 1.1.1\n"
# Test third split
subheader121_doc = split_docs[2]
assert subheader121_doc.meta["header"] == "Subheader 1.2.1"
assert subheader121_doc.meta["split_id"] == 2
assert subheader121_doc.meta["page_number"] == 1
assert subheader121_doc.meta["parent_headers"] == ["Header 1", "Header 1.2"]
assert subheader121_doc.content == "\nContent under header 1.2.1.\n"
# Test fourth split
subheader122_doc = split_docs[3]
assert subheader122_doc.meta["header"] == "Subheader 1.2.2"
assert subheader122_doc.meta["split_id"] == 3
assert subheader122_doc.meta["page_number"] == 1
assert subheader122_doc.meta["parent_headers"] == ["Header 1", "Header 1.2"]
assert subheader122_doc.content == "\nContent under header 1.2.2.\n"
# Test fifth split
subheader123_doc = split_docs[4]
assert subheader123_doc.meta["header"] == "Subheader 1.2.3"
assert subheader123_doc.meta["split_id"] == 4
assert subheader123_doc.meta["page_number"] == 1
assert subheader123_doc.meta["parent_headers"] == ["Header 1", "Header 1.2"]
assert subheader123_doc.content == "\nContent under header 1.2.3."
def test_split_no_headers():
splitter = MarkdownHeaderSplitter()
docs = [Document(content="No headers here."), Document(content="Just some text without headers.")]
result = splitter.run(documents=docs)
split_docs = result["documents"]
# Should return one doc per input, and no header key in meta
assert len(split_docs) == 2
for doc in split_docs:
assert "header" not in doc.meta
# Sanity Checks
assert split_docs[0].content == docs[0].content
assert split_docs[1].content == docs[1].content
def test_split_multiple_documents(sample_text):
splitter = MarkdownHeaderSplitter(keep_headers=False)
docs = [
Document(content=sample_text),
Document(content="# Another Header\nSome content."),
Document(content="# H1\nA"),
Document(content="# H2\nB"),
]
result = splitter.run(documents=docs)
split_docs = result["documents"]
assert len(split_docs) == 8
# First 5 splits are from sample_text
assert split_docs[5].meta["header"] == "Another Header"
assert split_docs[6].meta["header"] == "H1"
assert split_docs[7].meta["header"] == "H2"
# Verify that split_ids are per-parent-document
splits_by_source = defaultdict(list)
for doc in split_docs:
splits_by_source[doc.meta["source_id"]].append(doc.meta["split_id"])
# Each parent document should have split_ids starting from 0
for split_ids in splits_by_source.values():
assert split_ids == list(range(len(split_ids)))
def test_split_only_headers():
text = "# H1\n# H2\n# H3"
splitter = MarkdownHeaderSplitter()
docs = [Document(content=text)]
result = splitter.run(documents=docs)
split_docs = result["documents"]
# Return doc without content unchunked
assert len(split_docs) == 1
assert split_docs[0].content == text
# Metadata preservation
def test_preserve_document_metadata():
"""Test that document metadata is preserved through splitting."""
splitter = MarkdownHeaderSplitter(keep_headers=False) # keep_headers=True case is covered by this test too
docs = [Document(content="# Header\nContent", meta={"source": "test", "importance": "high", "custom_field": 123})]
result = splitter.run(documents=docs)
split_docs = result["documents"]
# Original metadata should be preserved
assert split_docs[0].meta["source"] == "test"
assert split_docs[0].meta["importance"] == "high"
assert split_docs[0].meta["custom_field"] == 123
# New metadata should be added
assert "header" in split_docs[0].meta
assert split_docs[0].meta["header"] == "Header"
assert "split_id" in split_docs[0].meta
assert split_docs[0].meta["split_id"] == 0
assert split_docs[0].content == "\nContent"
def test_secondary_split_keeps_content_before_embedded_header():
"""With keep_headers=False, prose before an embedded lower-level header must
not be dropped during the secondary split."""
splitter = MarkdownHeaderSplitter(
keep_headers=False, header_split_levels=[1], secondary_split="word", split_length=100
)
docs = splitter.run(documents=[Document(content="# Main\nintro paragraph text\n## Sub\nmore text\n")])["documents"]
combined = "".join(doc.content or "" for doc in docs)
assert "intro paragraph text" in combined
def test_secondary_split_keeps_content_before_code_fence_comment():
"""With keep_headers=False, prose before a '#' comment inside a fenced code block must
not be dropped during the secondary split (the comment is not a real header)."""
splitter = MarkdownHeaderSplitter(
keep_headers=False, header_split_levels=[1], secondary_split="word", split_length=100
)
content = "# Main\nsome intro text\n```python\n# a comment in code\n```\nmore text\n"
docs = splitter.run(documents=[Document(content=content)])["documents"]
combined = "".join(doc.content or "" for doc in docs)
assert "some intro text" in combined
# Error and edge case handling
def test_non_text_document():
"""Test that the component correctly handles non-text documents."""
splitter = MarkdownHeaderSplitter()
docs = [Document(content=None)]
# Should raise ValueError about text documents
with pytest.raises(ValueError, match="only works with text documents"):
splitter.run(documents=docs)
def test_empty_document_list():
"""Test handling of an empty document list."""
splitter = MarkdownHeaderSplitter()
result = splitter.run(documents=[])
assert result["documents"] == []
class TestHeaderSplitLevels:
def test_default_splits_on_all_levels(self, sample_text):
"""Default behaviour: all six header levels create split boundaries.
Note: empty headers (no content of their own) are folded into the next chunk via pending_headers rather
than appearing as standalone entries in meta.
"""
splitter = MarkdownHeaderSplitter() # header_split_levels defaults to [1,2,3,4,5,6]
docs = splitter.run(documents=[Document(content=sample_text)])["documents"]
# sample_text has 5 chunks with content; "Header 1.1" and "Header 1.2" are empty headers prepended to their
# first child and do not appear as their own split boundary
assert len(docs) == 5
assert docs[0].content == "# Header 1\nContent under header 1.\n"
assert docs[1].content == "## Header 1.1\n### Subheader 1.1.1\nContent under sub-header 1.1.1\n"
assert docs[2].content == "## Header 1.2\n### Subheader 1.2.1\nContent under header 1.2.1.\n"
assert docs[3].content == "### Subheader 1.2.2\nContent under header 1.2.2.\n"
assert docs[4].content == "### Subheader 1.2.3\nContent under header 1.2.3."
headers = [doc.meta["header"] for doc in docs]
assert headers == ["Header 1", "Subheader 1.1.1", "Subheader 1.2.1", "Subheader 1.2.2", "Subheader 1.2.3"]
def test_h1_and_h2_only(self, sample_text):
"""Only h1/h2 headers create splits; h3+ content is absorbed into the preceding chunk."""
splitter = MarkdownHeaderSplitter(header_split_levels=[1, 2])
docs = splitter.run(documents=[Document(content=sample_text)])["documents"]
assert len(docs) == 3
assert docs[0].content == "# Header 1\nContent under header 1.\n"
assert docs[1].content == "## Header 1.1\n### Subheader 1.1.1\nContent under sub-header 1.1.1\n"
assert docs[2].content == (
"## Header 1.2\n"
"### Subheader 1.2.1\nContent under header 1.2.1.\n"
"### Subheader 1.2.2\nContent under header 1.2.2.\n"
"### Subheader 1.2.3\nContent under header 1.2.3."
)
headers = [doc.meta["header"] for doc in docs]
assert headers == ["Header 1", "Header 1.1", "Header 1.2"]
# h3 headers must not appear as split boundaries
assert "Subheader 1.1.1" not in headers
assert "Subheader 1.2.1" not in headers
def test_single_level(self, sample_text):
"""Splitting on only h1 yields one chunk that is the full document."""
splitter = MarkdownHeaderSplitter(header_split_levels=[1])
docs = splitter.run(documents=[Document(content=sample_text)])["documents"]
assert len(docs) == 1
assert docs[0].meta["header"] == "Header 1"
# entire document is in one chunk — h1 is first, so content equals the full source text
assert docs[0].content == sample_text
def test_deep_levels_only(self):
"""Splitting on h3 only; h1/h2 headers above the first h3 are not captured in any chunk."""
text = (
"# Top Level\n"
"Ignored top content.\n"
"## Mid Level\n"
"Ignored mid content.\n"
"### Deep Section A\n"
"Content A.\n"
"### Deep Section B\n"
"Content B.\n"
)
splitter = MarkdownHeaderSplitter(header_split_levels=[3])
docs = splitter.run(documents=[Document(content=text)])["documents"]
assert len(docs) == 2
assert docs[0].content == "### Deep Section A\nContent A.\n"
assert docs[1].content == "### Deep Section B\nContent B.\n"
headers = [doc.meta["header"] for doc in docs]
assert "Top Level" not in headers
assert "Mid Level" not in headers
# text before the first h3 match is not absorbed — it is dropped entirely
assert "Ignored top content." not in docs[0].content
assert "Ignored mid content." not in docs[0].content
def test_non_contiguous_levels(self):
"""Non-contiguous level selection (e.g. [1, 3]) splits on h1 and h3 but not h2."""
text = "# H1 Title\n## H2 Ignored\nH2 content.\n### H3 Section\nH3 content.\n"
splitter = MarkdownHeaderSplitter(header_split_levels=[1, 3])
docs = splitter.run(documents=[Document(content=text)])["documents"]
assert len(docs) == 2
# h2 content sits between the h1 and h3 match boundaries, so it is absorbed into the h1 chunk
assert docs[0].content == "# H1 Title\n## H2 Ignored\nH2 content.\n"
assert docs[0].meta["header"] == "H1 Title"
assert docs[1].content == "### H3 Section\nH3 content.\n"
assert docs[1].meta["header"] == "H3 Section"
assert "H2 Ignored" not in [doc.meta["header"] for doc in docs]
def test_validation_empty_list(self):
with pytest.raises(ValueError, match="non-empty list"):
MarkdownHeaderSplitter(header_split_levels=[])
def test_validation_level_zero(self):
with pytest.raises(ValueError, match="invalid values"):
MarkdownHeaderSplitter(header_split_levels=[0, 1, 2])
def test_validation_level_seven(self):
with pytest.raises(ValueError, match="invalid values"):
MarkdownHeaderSplitter(header_split_levels=[1, 7])
def test_validation_non_integer(self):
with pytest.raises(ValueError, match="invalid values"):
MarkdownHeaderSplitter(header_split_levels=[1, "2"]) # type: ignore[list-item]
def test_validation_duplicate_levels(self):
with pytest.raises(ValueError, match="duplicate"):
MarkdownHeaderSplitter(header_split_levels=[1, 2, 2])
class TestCodeBlockExclusion:
"""Tests that hash lines inside fenced code blocks are not treated as headers."""
def test_backtick_fence(self):
"""Hash lines inside triple-backtick fences are ignored."""
text = (
"# Real Header\n"
"Some content.\n"
"```python\n"
"# this is a Python comment, not a header\n"
"## also not a header\n"
"x = 1\n"
"```\n"
"More content.\n"
"## Real Subheader\n"
"Subheader content.\n"
)
splitter = MarkdownHeaderSplitter()
docs = splitter.run(documents=[Document(content=text)])["documents"]
assert len(docs) == 2
assert docs[0].content == (
"# Real Header\n"
"Some content.\n"
"```python\n"
"# this is a Python comment, not a header\n"
"## also not a header\n"
"x = 1\n"
"```\n"
"More content.\n"
)
assert docs[0].meta["header"] == "Real Header"
assert docs[1].content == "## Real Subheader\nSubheader content.\n"
assert docs[1].meta["header"] == "Real Subheader"
assert "this is a Python comment, not a header" not in [doc.meta["header"] for doc in docs]
assert "also not a header" not in [doc.meta["header"] for doc in docs]
def test_tilde_fence(self):
"""Hash lines inside triple-tilde fences are ignored."""
text = "# Real Header\n~~~bash\n# shell comment\necho hello\n~~~\n## Real Subheader\nContent.\n"
splitter = MarkdownHeaderSplitter()
docs = splitter.run(documents=[Document(content=text)])["documents"]
assert len(docs) == 2
assert docs[0].content == "# Real Header\n~~~bash\n# shell comment\necho hello\n~~~\n"
assert docs[0].meta["header"] == "Real Header"
assert docs[1].content == "## Real Subheader\nContent.\n"
assert docs[1].meta["header"] == "Real Subheader"
assert "shell comment" not in [doc.meta["header"] for doc in docs]
def test_multiple_code_blocks(self):
"""Multiple fenced code blocks in one document are all excluded."""
text = (
"# Section One\n"
"Intro text.\n"
"```\n"
"# fake header A\n"
"```\n"
"Middle text.\n"
"```python\n"
"# fake header B\n"
"```\n"
"## Section Two\n"
"More content.\n"
)
splitter = MarkdownHeaderSplitter()
docs = splitter.run(documents=[Document(content=text)])["documents"]
assert len(docs) == 2
assert docs[0].content == (
"# Section One\nIntro text.\n```\n# fake header A\n```\nMiddle text.\n```python\n# fake header B\n```\n"
)
assert docs[0].meta["header"] == "Section One"
assert docs[1].content == "## Section Two\nMore content.\n"
assert docs[1].meta["header"] == "Section Two"
assert "fake header A" not in [doc.meta["header"] for doc in docs]
assert "fake header B" not in [doc.meta["header"] for doc in docs]
def test_longer_fence_delimiters(self):
"""Fences with more than three backticks/tildes are also recognised."""
text = "# Real Header\n````python\n# not a header\n````\n## Real Subheader\nContent.\n"
splitter = MarkdownHeaderSplitter()
docs = splitter.run(documents=[Document(content=text)])["documents"]
assert len(docs) == 2
assert docs[0].content == "# Real Header\n````python\n# not a header\n````\n"
assert docs[0].meta["header"] == "Real Header"
assert docs[1].content == "## Real Subheader\nContent.\n"
assert docs[1].meta["header"] == "Real Subheader"
assert "not a header" not in [doc.meta["header"] for doc in docs]
def test_code_block_with_no_real_headers(self):
"""If the only hash lines are inside code blocks, the document is returned unchunked."""
text = "Plain text before code.\n```\n# entirely fake\n```\nPlain text after code.\n"
splitter = MarkdownHeaderSplitter()
docs = splitter.run(documents=[Document(content=text)])["documents"]
assert len(docs) == 1
assert docs[0].content == text
assert "header" not in docs[0].meta
def test_invalid_secondary_split_at_init():
"""Test that an invalid secondary split type raises an error at initialization time."""
with pytest.raises(ValueError, match="split_by must be one of"):
MarkdownHeaderSplitter(secondary_split="invalid_split_type")
def test_invalid_split_parameters_at_init():
"""Test invalid split parameter validation at initialization time."""
# Test split_length validation
with pytest.raises(ValueError, match="split_length must be greater than 0"):
MarkdownHeaderSplitter(secondary_split="word", split_length=0)
# Test split_overlap validation
with pytest.raises(ValueError, match="split_overlap must be greater than or equal to 0"):
MarkdownHeaderSplitter(secondary_split="word", split_overlap=-1)
def test_empty_content_handling():
"""Test handling of documents with empty content."""
splitter_skip = MarkdownHeaderSplitter() # skip empty documents by default
docs = [Document(content="")]
result = splitter_skip.run(documents=docs)
assert len(result["documents"]) == 0
splitter_no_skip = MarkdownHeaderSplitter(skip_empty_documents=False)
docs = [Document(content="")]
result = splitter_no_skip.run(documents=docs)
assert len(result["documents"]) == 1
def test_split_id_sequentiality_primary_and_secondary(sample_text):
# Test primary splitting with single document
splitter = MarkdownHeaderSplitter()
docs = [Document(content=sample_text)]
result = splitter.run(documents=docs)
split_docs = result["documents"]
# Test number of documents
assert len(split_docs) == 5
# Check that split_ids are sequential from 0 for this single parent document
split_ids = [doc.meta["split_id"] for doc in split_docs]
assert split_ids == list(range(len(split_ids)))
# Test secondary splitting with single document
splitter = MarkdownHeaderSplitter(secondary_split="word", split_length=3)
docs = [Document(content=sample_text)]
result = splitter.run(documents=docs)
split_docs = result["documents"]
# Test number of documents
assert len(split_docs) == 12
# Check that split_ids are sequential from 0 for this single parent document
split_ids = [doc.meta["split_id"] for doc in split_docs]
assert split_ids == list(range(len(split_ids)))
# Test with multiple input documents; each should have its own split_id sequence
splitter = MarkdownHeaderSplitter(secondary_split="word", split_length=3) # Use fresh instance
docs = [Document(content=sample_text), Document(content="# Another Header\nSome more content here.")]
result = splitter.run(documents=docs)
split_docs = result["documents"]
# Test number of documents
assert len(split_docs) == 14
# Verify split_ids are per-parent-document
splits_by_source = defaultdict(list)
for doc in split_docs:
splits_by_source[doc.meta["source_id"]].append(doc.meta["split_id"])
# Each parent document should have split_ids starting from 0
for split_ids in splits_by_source.values():
assert split_ids == list(range(len(split_ids)))
def test_secondary_split_with_overlap():
text = (
"# Introduction\n"
"This is the introduction section with some words for testing overlap splitting. "
"It should be split into chunks with overlap.\n"
"## Details\n"
"Here are more details about the topic. "
"Splitting should work across multiple headers and content blocks.\n"
"### Subsection\n"
"This subsection contains additional information and should also be split with overlap."
)
# keep_headers=False
splitter = MarkdownHeaderSplitter(secondary_split="word", split_length=8, split_overlap=3, keep_headers=False)
docs = [Document(content=text)]
result = splitter.run(documents=docs)
split_docs = result["documents"]
assert len(split_docs) == 9
# Verify exact content and metadata of each split
# intro (4 docs)
assert split_docs[0].content == "\nThis is the introduction section with some words "
assert split_docs[0].meta["header"] == "Introduction"
assert split_docs[0].meta["split_id"] == 0
assert split_docs[1].content == "with some words for testing overlap splitting. It "
assert split_docs[1].meta["header"] == "Introduction"
assert split_docs[1].meta["split_id"] == 1
assert split_docs[2].content == "overlap splitting. It should be split into chunks "
assert split_docs[2].meta["header"] == "Introduction"
assert split_docs[2].meta["split_id"] == 2
assert split_docs[3].content == "split into chunks with overlap.\n"
assert split_docs[3].meta["header"] == "Introduction"
assert split_docs[3].meta["split_id"] == 3
# details (3 docs)
assert split_docs[4].content == "\nHere are more details about the topic. Splitting "
assert split_docs[4].meta["header"] == "Details"
assert split_docs[4].meta["split_id"] == 4
assert split_docs[5].content == "the topic. Splitting should work across multiple headers "
assert split_docs[5].meta["header"] == "Details"
assert split_docs[5].meta["split_id"] == 5
assert split_docs[6].content == "across multiple headers and content blocks.\n"
assert split_docs[6].meta["header"] == "Details"
assert split_docs[6].meta["split_id"] == 6
# subsection (2 docs)
assert split_docs[7].content == "\nThis subsection contains additional information and should also "
assert split_docs[7].meta["header"] == "Subsection"
assert split_docs[7].meta["split_id"] == 7
assert split_docs[8].content == "and should also be split with overlap."
assert split_docs[8].meta["header"] == "Subsection"
assert split_docs[8].meta["split_id"] == 8
# verify 3-word overlap behavior (split_overlap=3)
# consecutive pairs within a header should share the 3 words at their boundary
# intro
assert split_docs[0].content.split()[-3:] == split_docs[1].content.split()[:3]
assert split_docs[1].content.split()[-3:] == split_docs[2].content.split()[:3]
assert split_docs[2].content.split()[-3:] == split_docs[3].content.split()[:3]
# details
assert split_docs[4].content.split()[-3:] == split_docs[5].content.split()[:3]
assert split_docs[5].content.split()[-3:] == split_docs[6].content.split()[:3]
# subsection
assert split_docs[7].content.split()[-3:] == split_docs[8].content.split()[:3]
# re-run with keep_headers=True, change split_length and split_overlap
splitter = MarkdownHeaderSplitter(secondary_split="word", split_length=4, split_overlap=2)
docs = [Document(content=text)]
result = splitter.run(documents=docs)
split_docs = result["documents"]
assert len(split_docs) == 24
assert split_docs[0].content.startswith("# Introduction")
assert all("header" in doc.meta for doc in split_docs)
def test_secondary_split_with_threshold():
text = "# Header\n" + " ".join([f"word{i}" for i in range(1, 11)])
# keep_headers=True
splitter = MarkdownHeaderSplitter(secondary_split="word", split_length=3, split_threshold=2, keep_headers=True)
docs = [Document(content=text)]
result = splitter.run(documents=docs)
split_docs = result["documents"]
# Explicitly test each split
assert len(split_docs) == 4
assert split_docs[0].content == "# Header\nword1 word2 "
assert split_docs[0].meta["split_id"] == 0
assert split_docs[1].content == "word3 word4 word5 "
assert split_docs[1].meta["split_id"] == 1
assert split_docs[2].content == "word6 word7 word8 "
assert split_docs[2].meta["split_id"] == 2
assert split_docs[3].content == "word9 word10"
assert split_docs[3].meta["split_id"] == 3
# keep_headers=False
splitter = MarkdownHeaderSplitter(secondary_split="word", split_length=3, split_threshold=2, keep_headers=False)
docs = [Document(content=text)]
result = splitter.run(documents=docs)
split_docs = result["documents"]
# Explicitly test each split
assert len(split_docs) == 3
assert split_docs[0].content == "\nword1 word2 word3 "
assert split_docs[0].meta["split_id"] == 0
assert split_docs[1].content == "word4 word5 word6 "
assert split_docs[1].meta["split_id"] == 1
assert split_docs[2].content == "word7 word8 word9 word10" # 4 words (due to threshold, not possible to split 3-1)
assert split_docs[2].meta["split_id"] == 2
def test_page_break_handling_in_secondary_split():
text = "# Header\nFirst page\f Second page\f Third page"
splitter = MarkdownHeaderSplitter(secondary_split="word", split_length=1)
docs = [Document(content=text)]
result = splitter.run(documents=docs)
split_docs = result["documents"]
expected_page_numbers = [1, 1, 1, 2, 2, 3, 3]
actual_page_numbers = [doc.meta.get("page_number") for doc in split_docs]
assert actual_page_numbers == expected_page_numbers
def test_page_break_handling_with_multiple_headers(sample_text_with_page_breaks):
splitter = MarkdownHeaderSplitter(secondary_split="word", split_length=3)
docs = [Document(content=sample_text_with_page_breaks)]
result = splitter.run(documents=docs)
split_docs = result["documents"]
assert len(split_docs) == 12
assert split_docs[0].content == "# Header 1\nContent "
assert split_docs[0].meta == {
"source_id": ANY,
"split_id": 0,
"page_number": 1,
"split_idx_start": 0,
"header": "Header 1",
"parent_headers": [],
}
assert split_docs[1].content == "under header 1.\n\f\n"
assert split_docs[1].meta == {
"source_id": ANY,
"split_id": 1,
"page_number": 1,
"split_idx_start": 19,
"header": "Header 1",
"parent_headers": [],
}
assert split_docs[2].content == "## Header 1.1\n### "
assert split_docs[2].meta == {
"source_id": ANY,
"split_id": 2,
"page_number": 2,
"split_idx_start": 0,
"header": "Subheader 1.1.1",
"parent_headers": ["Header 1", "Header 1.1"],
}
assert split_docs[3].content == "Subheader 1.1.1\nContent under "
assert split_docs[3].meta == {
"source_id": ANY,
"split_id": 3,
"page_number": 2,
"split_idx_start": 18,
"header": "Subheader 1.1.1",
"parent_headers": ["Header 1", "Header 1.1"],
}
assert split_docs[4].content == "sub-header 1.1.1\n\f\n"
assert split_docs[4].meta == {
"source_id": ANY,
"split_id": 4,
"page_number": 2,
"split_idx_start": 48,
"header": "Subheader 1.1.1",
"parent_headers": ["Header 1", "Header 1.1"],
}
assert split_docs[5].content == "## Header 1.2\n### "
assert split_docs[5].meta == {
"source_id": ANY,
"split_id": 5,
"page_number": 3,
"split_idx_start": 0,
"header": "Subheader 1.2.1",
"parent_headers": ["Header 1", "Header 1.2"],
}
assert split_docs[6].content == "Subheader 1.2.1\nContent under "
assert split_docs[6].meta == {
"source_id": ANY,
"split_id": 6,
"page_number": 3,
"split_idx_start": 18,
"header": "Subheader 1.2.1",
"parent_headers": ["Header 1", "Header 1.2"],
}
assert split_docs[7].content == "header 1.2.1.\n\f\n"
assert split_docs[7].meta == {
"source_id": ANY,
"split_id": 7,
"page_number": 3,
"split_idx_start": 48,
"header": "Subheader 1.2.1",
"parent_headers": ["Header 1", "Header 1.2"],
}
assert split_docs[8].content == "### Subheader 1.2.2\nContent "
assert split_docs[8].meta == {
"source_id": ANY,
"split_id": 8,
"page_number": 4,
"split_idx_start": 0,
"header": "Subheader 1.2.2",
"parent_headers": ["Header 1", "Header 1.2"],
}
assert split_docs[9].content == "under header 1.2.2.\n\f\n"
assert split_docs[9].meta == {
"source_id": ANY,
"split_id": 9,
"page_number": 4,
"split_idx_start": 28,
"header": "Subheader 1.2.2",
"parent_headers": ["Header 1", "Header 1.2"],
}
assert split_docs[10].content == "### Subheader 1.2.3\nContent "
assert split_docs[10].meta == {
"source_id": ANY,
"split_id": 10,
"page_number": 5,
"split_idx_start": 0,
"header": "Subheader 1.2.3",
"parent_headers": ["Header 1", "Header 1.2"],
}
assert split_docs[11].content == "under header 1.2.3."
assert split_docs[11].meta == {
"source_id": ANY,
"split_id": 11,
"page_number": 5,
"split_idx_start": 28,
"header": "Subheader 1.2.3",
"parent_headers": ["Header 1", "Header 1.2"],
}
# reconstruct original
reconstructed_text = "".join(doc.content for doc in split_docs)
assert reconstructed_text == sample_text_with_page_breaks
@@ -0,0 +1,874 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
import textwrap
import pytest
from haystack import Document
from haystack.components.preprocessors import PythonCodeSplitter
@pytest.fixture
def simple_module_source():
return textwrap.dedent(
'''
"""Example module docstring."""
import os
import sys
from math import sqrt
def add(a, b):
"""Add two numbers."""
return a + b
def subtract(a, b):
"""Subtract two numbers."""
return a - b
'''
).lstrip()
@pytest.fixture
def class_source():
return textwrap.dedent(
'''
"""Geometry helpers."""
from math import pi
class Shape:
"""Base shape."""
kind = "shape"
def __init__(self, name: str) -> None:
self.name = name
def describe(self) -> str:
return f"shape {self.name}"
class Circle(Shape, metaclass=type):
"""A circle."""
def __init__(self, r: float) -> None:
super().__init__("circle")
self.r = r
@staticmethod
def pi_value() -> float:
return pi
@classmethod
def unit(cls) -> "Circle":
return cls(1.0)
def area(self) -> float:
return pi * self.r * self.r
'''
).lstrip()
@pytest.fixture
def oversized_function_source():
"""A function long enough to trigger the secondary line-based split."""
body_lines = "\n".join(f" x_{i} = {i}" for i in range(200))
return f"def giant():\n{body_lines}\n return x_0\n"
class TestInitValidation:
def test_defaults(self):
splitter = PythonCodeSplitter()
assert splitter.min_effective_lines == 20
assert splitter.max_effective_lines == 100
assert splitter.expected_chars_per_line == 45
assert splitter.oversized_factor == 3
assert splitter.strip_docstrings is False
assert splitter.preserve_class_definition is True
assert splitter.secondary_split_overlap == 5
assert splitter.secondary_split_length is None
def test_custom_values(self):
splitter = PythonCodeSplitter(
min_effective_lines=2,
max_effective_lines=10,
expected_chars_per_line=80,
oversized_factor=4,
strip_docstrings=True,
preserve_class_definition=False,
secondary_split_overlap=2,
secondary_split_length=15,
)
assert splitter.min_effective_lines == 2
assert splitter.max_effective_lines == 10
assert splitter.expected_chars_per_line == 80
assert splitter.oversized_factor == 4
assert splitter.strip_docstrings is True
assert splitter.preserve_class_definition is False
assert splitter.secondary_split_overlap == 2
assert splitter.secondary_split_length == 15
@pytest.mark.parametrize(
"kwargs",
[
{"min_effective_lines": 0},
{"min_effective_lines": -1},
{"max_effective_lines": 0},
{"max_effective_lines": -3},
{"min_effective_lines": 10, "max_effective_lines": 5},
{"expected_chars_per_line": 0},
{"oversized_factor": 0},
{"secondary_split_overlap": -1},
{"secondary_split_length": -1},
],
)
def test_invalid_init_raises(self, kwargs):
with pytest.raises(ValueError):
PythonCodeSplitter(**kwargs)
class TestRunInputValidation:
def test_none_content_raises_value_error(self):
splitter = PythonCodeSplitter()
doc = Document(content=None)
with pytest.raises(ValueError):
splitter.run(documents=[doc])
def test_non_string_content_raises_type_error(self):
splitter = PythonCodeSplitter()
doc = Document(content="placeholder")
doc.content = 12345 # type: ignore[assignment]
with pytest.raises(TypeError):
splitter.run(documents=[doc])
def test_invalid_syntax_raises(self):
splitter = PythonCodeSplitter()
doc = Document(content="def broken(:\n pass\n")
with pytest.raises(SyntaxError):
splitter.run(documents=[doc])
def test_empty_documents_list(self):
splitter = PythonCodeSplitter()
result = splitter.run(documents=[])
assert result == {"documents": []}
class TestBasicOutput:
def test_returns_dict_with_documents(self, simple_module_source):
splitter = PythonCodeSplitter()
result = splitter.run(documents=[Document(content=simple_module_source)])
assert isinstance(result, dict)
assert "documents" in result
assert isinstance(result["documents"], list)
assert len(result["documents"]) >= 1
for chunk in result["documents"]:
assert isinstance(chunk, Document)
assert isinstance(chunk.content, str)
assert chunk.content
def test_split_id_starts_at_zero_and_increments(self, class_source):
splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5)
result = splitter.run(documents=[Document(content=class_source)])
chunks = result["documents"]
ids = [c.meta["split_id"] for c in chunks]
assert ids == list(range(len(chunks)))
def test_source_id_consistent_within_one_document(self, class_source):
splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5)
result = splitter.run(documents=[Document(content=class_source)])
chunks = result["documents"]
source_ids = {c.meta["source_id"] for c in chunks}
assert len(source_ids) == 1
def test_source_id_differs_between_documents(self, simple_module_source, class_source):
splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=10)
docs = [Document(content=simple_module_source), Document(content=class_source)]
result = splitter.run(documents=docs)
source_ids = {c.meta["source_id"] for c in result["documents"]}
assert len(source_ids) == 2
def test_chunks_have_required_meta_fields(self, simple_module_source):
splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=10)
result = splitter.run(documents=[Document(content=simple_module_source)])
for chunk in result["documents"]:
assert "source_id" in chunk.meta
assert "split_id" in chunk.meta
assert "start_line" in chunk.meta
assert "end_line" in chunk.meta
assert "unit_kinds" in chunk.meta
assert isinstance(chunk.meta["unit_kinds"], list)
assert chunk.meta["start_line"] >= 1
assert chunk.meta["end_line"] >= chunk.meta["start_line"]
def test_unit_kinds_lists_what_was_merged(self, simple_module_source):
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=200)
result = splitter.run(documents=[Document(content=simple_module_source)])
all_kinds = set()
for chunk in result["documents"]:
for kind in chunk.meta["unit_kinds"]:
all_kinds.add(kind)
text = " ".join(all_kinds).lower()
assert any("import" in t for t in all_kinds) or "import" in text
assert any("func" in t or "method" in t for t in all_kinds) or any("func" in t for t in text.split())
def test_multiple_documents_each_produces_chunks(self, simple_module_source, class_source):
splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=10)
docs = [
Document(content=simple_module_source, meta={"file_name": "a.py"}),
Document(content=class_source, meta={"file_name": "b.py"}),
]
result = splitter.run(documents=docs)
file_names = {c.meta.get("file_name") for c in result["documents"]}
assert file_names == {"a.py", "b.py"}
def test_split_id_resets_per_document(self, simple_module_source, class_source):
splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=10)
docs = [
Document(content=simple_module_source, meta={"file_name": "a.py"}),
Document(content=class_source, meta={"file_name": "b.py"}),
]
result = splitter.run(documents=docs)
per_file_ids: dict[str, list[int]] = {"a.py": [], "b.py": []}
for chunk in result["documents"]:
per_file_ids[chunk.meta["file_name"]].append(chunk.meta["split_id"])
for ids in per_file_ids.values():
assert ids == sorted(ids)
assert ids[0] == 0
class TestOrderingAndLineRanges:
def test_chunks_are_in_source_order(self, class_source):
splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5)
result = splitter.run(documents=[Document(content=class_source)])
prev_start = 0
for chunk in result["documents"]:
assert chunk.meta["start_line"] >= prev_start
assert chunk.meta["end_line"] >= chunk.meta["start_line"]
prev_start = chunk.meta["start_line"]
def test_chunks_dont_overlap_in_primary_split(self, class_source):
splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5)
chunks = splitter.run(documents=[Document(content=class_source)])["documents"]
for prev, nxt in zip(chunks, chunks[1:], strict=False):
assert nxt.meta["start_line"] > prev.meta["end_line"]
def test_chunks_read_top_to_bottom(self, class_source):
# With the default preserve_class_definition=True, chunks may have a class
# signature prefixed but the source slice itself must still appear verbatim.
splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5)
result = splitter.run(documents=[Document(content=class_source)])
source_lines = class_source.splitlines(keepends=True)
for chunk in result["documents"]:
expected = "".join(source_lines[chunk.meta["start_line"] - 1 : chunk.meta["end_line"]])
assert expected in (chunk.content or "")
def test_chunks_equal_source_slice_without_class_preservation(self, class_source):
splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5, preserve_class_definition=False)
result = splitter.run(documents=[Document(content=class_source)])
source_lines = class_source.splitlines(keepends=True)
for chunk in result["documents"]:
expected = "".join(source_lines[chunk.meta["start_line"] - 1 : chunk.meta["end_line"]])
assert (chunk.content or "").endswith(expected)
class TestFileNamePropagation:
def test_file_name_propagated_to_all_chunks(self, simple_module_source):
splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5)
result = splitter.run(documents=[Document(content=simple_module_source, meta={"file_name": "sample.py"})])
for chunk in result["documents"]:
assert chunk.meta["file_name"] == "sample.py"
def test_no_file_name_when_absent(self, simple_module_source):
splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5)
result = splitter.run(documents=[Document(content=simple_module_source)])
for chunk in result["documents"]:
assert "file_name" not in chunk.meta
def test_other_meta_is_propagated(self, simple_module_source):
splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5)
result = splitter.run(
documents=[Document(content=simple_module_source, meta={"file_name": "x.py", "project": "haystack"})]
)
for chunk in result["documents"]:
assert chunk.meta["project"] == "haystack"
class TestDecorators:
def test_decorators_metadata_present(self):
source = textwrap.dedent(
"""
class A:
@staticmethod
def s():
return 1
@classmethod
def c(cls):
return 2
"""
).lstrip()
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=2)
result = splitter.run(documents=[Document(content=source)])
all_decorators = []
for chunk in result["documents"]:
all_decorators.extend(chunk.meta.get("decorators") or [])
assert any("staticmethod" in d for d in all_decorators)
assert any("classmethod" in d for d in all_decorators)
def test_decorator_lines_included_in_chunk_content(self):
source = textwrap.dedent(
"""
class A:
@staticmethod
def s():
return 1
"""
).lstrip()
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=2)
result = splitter.run(documents=[Document(content=source)])
chunks_with_s = [c for c in result["documents"] if "def s" in (c.content or "")]
assert chunks_with_s
for chunk in chunks_with_s:
assert "@staticmethod" in (chunk.content or "")
def test_decorators_deduped_in_chunk(self):
source = textwrap.dedent(
"""
class A:
@staticmethod
def one():
return 1
@staticmethod
def two():
return 2
"""
).lstrip()
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=20)
result = splitter.run(documents=[Document(content=source)])
for chunk in result["documents"]:
decorators = chunk.meta.get("decorators") or []
assert len(decorators) == len(set(decorators))
def test_function_with_three_decorators_lists_all(self):
source = textwrap.dedent(
"""
def deco_a(fn):
return fn
def deco_b(fn):
return fn
def deco_c(fn):
return fn
@deco_a
@deco_b
@deco_c
def triple():
return 1
"""
).lstrip()
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=200)
result = splitter.run(documents=[Document(content=source)])
chunks_with_triple = [c for c in result["documents"] if "def triple" in (c.content or "")]
assert chunks_with_triple
decorators = [d for c in chunks_with_triple for d in (c.meta.get("decorators") or [])]
joined = " ".join(decorators)
assert "deco_a" in joined
assert "deco_b" in joined
assert "deco_c" in joined
def test_function_with_three_decorators_all_lines_in_content(self):
source = textwrap.dedent(
"""
@deco_a
@deco_b
@deco_c
def triple():
return 1
"""
).lstrip()
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=50)
result = splitter.run(documents=[Document(content=source)])
chunks_with_triple = [c for c in result["documents"] if "def triple" in (c.content or "")]
assert chunks_with_triple
for chunk in chunks_with_triple:
content = chunk.content or ""
assert "@deco_a" in content
assert "@deco_b" in content
assert "@deco_c" in content
class TestIncludeClassesMeta:
@pytest.mark.parametrize("class_name", ["Shape", "Circle"])
def test_class_methods_carry_include_classes(self, class_source, class_name):
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=3)
result = splitter.run(documents=[Document(content=class_source)])
chunks = [c for c in result["documents"] if class_name in (c.meta.get("include_classes") or [])]
assert chunks
def test_include_classes_absent_when_no_class_involved(self, simple_module_source):
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=200)
result = splitter.run(documents=[Document(content=simple_module_source)])
for chunk in result["documents"]:
include_classes = chunk.meta.get("include_classes")
assert not include_classes
def test_include_classes_is_deduplicated(self):
source = textwrap.dedent(
"""
class A:
def one(self):
return 1
def two(self):
return 2
def three(self):
return 3
"""
).lstrip()
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=200)
result = splitter.run(documents=[Document(content=source)])
for chunk in result["documents"]:
include_classes = chunk.meta.get("include_classes") or []
assert len(include_classes) == len(set(include_classes))
def test_include_classes_preserves_source_order(self):
source = textwrap.dedent(
"""
class First:
def f(self):
return 1
class Second:
def g(self):
return 2
class Third:
def h(self):
return 3
"""
).lstrip()
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=500)
result = splitter.run(documents=[Document(content=source)])
for chunk in result["documents"]:
include_classes = chunk.meta.get("include_classes") or []
if len(include_classes) >= 2:
expected_order = [c for c in ["First", "Second", "Third"] if c in include_classes]
assert include_classes == expected_order
class TestPreserveClassDefinition:
@pytest.fixture
def multi_method_class_source(self):
return textwrap.dedent(
'''
class Greeter:
"""A friendly greeter."""
kind = "greeter"
def __init__(self, name: str) -> None:
self.name = name
def hello(self) -> str:
return f"hello {self.name}"
def bye(self) -> str:
return f"bye {self.name}"
def shout(self) -> str:
return f"HELLO {self.name.upper()}"
def whisper(self) -> str:
return f"hello {self.name}..."
'''
).lstrip()
def test_default_preserve_class_definition_is_true(self):
splitter = PythonCodeSplitter()
assert splitter.preserve_class_definition is True
def test_class_signature_prepended_to_later_chunks(self, multi_method_class_source):
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=3, preserve_class_definition=True)
chunks = splitter.run(documents=[Document(content=multi_method_class_source)])["documents"]
assert len(chunks) >= 2
for chunk in chunks:
if "Greeter" in (chunk.meta.get("include_classes") or []):
assert "class Greeter" in (chunk.content or "")
def test_disabled_does_not_prepend_class_signature(self, multi_method_class_source):
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=3, preserve_class_definition=False)
chunks = splitter.run(documents=[Document(content=multi_method_class_source)])["documents"]
chunks_with_header = [c for c in chunks if "class Greeter" in (c.content or "")]
assert len(chunks_with_header) == 1
def test_preserve_does_not_duplicate_header_in_original_chunk(self, multi_method_class_source):
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=3, preserve_class_definition=True)
result = splitter.run(documents=[Document(content=multi_method_class_source)])
for chunk in result["documents"]:
assert (chunk.content or "").count("class Greeter") <= 1
def test_preserve_keeps_inheritance_and_metaclass(self):
source = textwrap.dedent(
'''
class Base:
pass
class Meta(type):
pass
class Child(Base, metaclass=Meta):
"""A child class."""
def one(self):
return 1
def two(self):
return 2
def three(self):
return 3
def four(self):
return 4
'''
).lstrip()
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=3, preserve_class_definition=True)
result = splitter.run(documents=[Document(content=source)])
child_chunks = [c for c in result["documents"] if "Child" in (c.meta.get("include_classes") or [])]
assert len(child_chunks) >= 2
for chunk in child_chunks:
assert "class Child(Base, metaclass=Meta):" in (chunk.content or "")
def test_preserve_keeps_decorators_on_class(self):
source = textwrap.dedent(
"""
def reg(cls):
return cls
@reg
class Decorated:
def one(self):
return 1
def two(self):
return 2
def three(self):
return 3
def four(self):
return 4
"""
).lstrip()
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=3, preserve_class_definition=True)
result = splitter.run(documents=[Document(content=source)])
decorated_chunks = [c for c in result["documents"] if "Decorated" in (c.meta.get("include_classes") or [])]
assert len(decorated_chunks) >= 2
for chunk in decorated_chunks:
content = chunk.content or ""
assert "class Decorated" in content
assert "@reg" in content
def test_preserve_handles_multiple_classes(self):
source = textwrap.dedent(
"""
class A:
def a1(self):
return 1
def a2(self):
return 2
def a3(self):
return 3
class B:
def b1(self):
return 1
def b2(self):
return 2
def b3(self):
return 3
"""
).lstrip()
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=3, preserve_class_definition=True)
result = splitter.run(documents=[Document(content=source)])
for chunk in result["documents"]:
content = chunk.content or ""
for cls in chunk.meta.get("include_classes") or []:
assert f"class {cls}" in content
class TestDocstringStripping:
def test_default_keeps_docstrings_in_content(self, class_source):
splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=8)
result = splitter.run(documents=[Document(content=class_source)])
contents = "\n".join(c.content or "" for c in result["documents"])
assert "A circle." in contents
for chunk in result["documents"]:
assert not chunk.meta.get("docstrings")
def test_strip_docstrings_moves_them_to_meta(self):
source = textwrap.dedent(
'''
def foo():
"""Foo docstring."""
return 1
def bar():
"""Bar docstring."""
return 2
'''
).lstrip()
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=10, strip_docstrings=True)
result = splitter.run(documents=[Document(content=source)])
all_docstrings = [d for c in result["documents"] for d in (c.meta.get("docstrings") or [])]
joined_docstrings = " | ".join(all_docstrings)
assert "Foo docstring." in joined_docstrings
assert "Bar docstring." in joined_docstrings
joined_content = "\n".join(c.content or "" for c in result["documents"])
assert "Foo docstring." not in joined_content
assert "Bar docstring." not in joined_content
def test_strip_docstrings_preserves_module_docstring(self):
source = textwrap.dedent(
'''
"""Module-level docstring."""
def foo():
"""Inner."""
return 1
'''
).lstrip()
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=20, strip_docstrings=True)
result = splitter.run(documents=[Document(content=source)])
joined = "\n".join(c.content or "" for c in result["documents"])
assert "Module-level docstring." in joined
def test_strip_class_header_docstring_moves_to_meta(self):
source = textwrap.dedent(
'''
class MyClass:
"""Class-level docstring."""
class_var = 42
def method(self):
return self.class_var
'''
).lstrip()
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=10, strip_docstrings=True)
result = splitter.run(documents=[Document(content=source)])
header_chunks = [c for c in result["documents"] if "class_header" in c.meta.get("unit_kinds", [])]
assert header_chunks
header = header_chunks[0]
assert "Class-level docstring." not in (header.content or "")
assert "Class-level docstring." in " | ".join(header.meta.get("docstrings") or [])
class TestTopLevelStatements:
@pytest.fixture
def rich_module_source(self):
return textwrap.dedent(
'''
"""Utility helpers for the pipeline."""
import os
import sys
from pathlib import Path
MAX_RETRIES = 3
DEFAULT_TIMEOUT = 30.0
LOG_PREFIX = "app"
def process(data):
"""Process data."""
result = data.strip()
return result
def validate(value):
"""Validate a value."""
if value is None:
raise ValueError("value cannot be None")
return True
class Manager:
"""Resource manager."""
def __init__(self):
self.items = []
def add(self, item):
self.items.append(item)
def remove(self, item):
self.items.remove(item)
if __name__ == "__main__":
mgr = Manager()
mgr.add(process("test"))
'''
).lstrip()
@pytest.mark.parametrize("expected_kind", ["statement", "imports", "module_docstring"])
def test_unit_kind_present(self, rich_module_source, expected_kind):
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=50)
result = splitter.run(documents=[Document(content=rich_module_source)])
all_kinds = [k for c in result["documents"] for k in c.meta.get("unit_kinds", [])]
assert expected_kind in all_kinds
def test_first_chunk_contains_preamble_statements(self, rich_module_source):
# max_effective_lines=6 matches the preamble (module docstring + 3 imports + 3
# assignments) so the greedy merger flushes before the first function.
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=6)
result = splitter.run(documents=[Document(content=rich_module_source)])
assert len(result["documents"]) >= 2
first = result["documents"][0]
content = first.content or ""
assert '"""Utility helpers for the pipeline."""' in content
assert "import os" in content
assert "import sys" in content
assert "from pathlib import Path" in content
assert "MAX_RETRIES = 3" in content
assert "DEFAULT_TIMEOUT = 30.0" in content
assert 'LOG_PREFIX = "app"' in content
assert "def process" not in content
assert "def validate" not in content
assert "class Manager" not in content
def test_if_main_produces_statement_unit(self, rich_module_source):
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=50)
result = splitter.run(documents=[Document(content=rich_module_source)])
all_kinds = [k for c in result["documents"] for k in c.meta.get("unit_kinds", [])]
assert "statement" in all_kinds
joined = "\n".join(c.content or "" for c in result["documents"])
assert 'if __name__ == "__main__"' in joined
def test_all_imports_appear_in_output(self, rich_module_source):
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=50)
result = splitter.run(documents=[Document(content=rich_module_source)])
joined = "\n".join(c.content or "" for c in result["documents"])
assert "import os" in joined
assert "import sys" in joined
assert "from pathlib import Path" in joined
def test_module_docstring_text_preserved(self, rich_module_source):
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=50)
result = splitter.run(documents=[Document(content=rich_module_source)])
joined = "\n".join(c.content or "" for c in result["documents"])
assert "Utility helpers for the pipeline." in joined
class TestOversizedFallback:
def test_warns_on_oversized_function(self, oversized_function_source, caplog):
import logging
splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5, oversized_factor=3)
with caplog.at_level(logging.WARNING):
result = splitter.run(documents=[Document(content=oversized_function_source)])
text = caplog.text.lower()
assert "oversiz" in text or "secondary" in text
assert len(result["documents"]) >= 2
def test_oversized_chunks_marked_with_secondary_split(self, oversized_function_source):
splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5, oversized_factor=3)
result = splitter.run(documents=[Document(content=oversized_function_source)])
assert [c for c in result["documents"] if c.meta.get("secondary_split")]
def test_oversized_chunks_belong_to_same_function(self, oversized_function_source):
splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5, oversized_factor=3)
result = splitter.run(documents=[Document(content=oversized_function_source)])
secondary = [c for c in result["documents"] if c.meta.get("secondary_split")]
for chunk in secondary:
assert chunk.meta["start_line"] >= 1
assert chunk.meta["end_line"] <= oversized_function_source.count("\n") + 1
def test_small_function_does_not_trigger_secondary(self, simple_module_source):
splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=50, oversized_factor=3)
result = splitter.run(documents=[Document(content=simple_module_source)])
for chunk in result["documents"]:
assert not chunk.meta.get("secondary_split")
def test_long_lines_count_as_more_effective_lines(self):
long_line_source = (
textwrap.dedent(
"""
def short():
return 1
def longer():
return "{padding}"
"""
)
.lstrip()
.format(padding="x" * 500)
)
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=2, expected_chars_per_line=10)
result = splitter.run(documents=[Document(content=long_line_source)])
chunks_with_short = [c for c in result["documents"] if "def short" in (c.content or "")]
chunks_with_long = [c for c in result["documents"] if "def longer" in (c.content or "")]
assert chunks_with_short and chunks_with_long
assert chunks_with_short[0] is not chunks_with_long[0]
class TestEdgeCases:
def test_module_with_only_docstring(self):
source = '"""Just a docstring."""\n'
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=5)
result = splitter.run(documents=[Document(content=source)])
assert len(result["documents"]) == 1
assert "Just a docstring." in (result["documents"][0].content or "")
def test_module_with_only_imports(self):
source = "import os\nimport sys\nfrom math import pi\n"
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=5)
result = splitter.run(documents=[Document(content=source)])
joined = "\n".join(c.content or "" for c in result["documents"])
assert "import os" in joined
assert "import sys" in joined
assert "from math import pi" in joined
def test_module_with_only_one_function(self):
source = "def f():\n return 1\n"
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=5)
result = splitter.run(documents=[Document(content=source)])
assert len(result["documents"]) == 1
assert "def f" in (result["documents"][0].content or "")
def test_empty_string_source(self):
# Empty source is valid Python; the splitter must not raise.
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=5)
result = splitter.run(documents=[Document(content="")])
assert isinstance(result["documents"], list)
def test_invalid_syntax_raises_syntax_error(self):
splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=5)
doc = Document(content="class Broken(\n pass\n")
with pytest.raises(SyntaxError):
splitter.run(documents=[doc])
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,117 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
import time
from pathlib import Path
from unittest.mock import patch
import pytest
from pytest import LogCaptureFixture
from haystack.components.preprocessors.sentence_tokenizer import QUOTE_SPANS_RE, SentenceSplitter
def test_apply_split_rules_no_join() -> None:
text = "This is a test. This is another test. And a third test."
spans = [(0, 15), (16, 36), (37, 54)]
result = SentenceSplitter._apply_split_rules(text, spans)
assert len(result) == 3
assert result == [(0, 15), (16, 36), (37, 54)]
def test_apply_split_rules_join_case_1():
text = 'He said "This is sentence one. This is sentence two." Then he left.'
result = SentenceSplitter._apply_split_rules(text, [(0, 30), (31, 53), (54, 67)])
assert len(result) == 2
assert result == [(0, 53), (54, 67)]
def test_apply_split_rules_join_case_3():
splitter = SentenceSplitter(language="en", use_split_rules=True)
text = """
1. First item
2. Second item
3. Third item."""
spans = [(0, 7), (8, 25), (26, 44), (45, 56)]
result = splitter._apply_split_rules(text, spans)
assert len(result) == 1
assert result == [(0, 56)]
def test_apply_split_rules_join_case_4() -> None:
text = "This is a test. (With a parenthetical statement.) And another sentence."
spans = [(0, 15), (16, 50), (51, 74)]
result = SentenceSplitter._apply_split_rules(text, spans)
assert len(result) == 2
assert result == [(0, 50), (51, 74)]
@pytest.fixture
def mock_file_content():
return "Mr.\nDr.\nProf."
def test_read_abbreviations_existing_file(tmp_path, mock_file_content):
abbrev_dir = tmp_path / "data" / "abbreviations"
abbrev_dir.mkdir(parents=True)
abbrev_file = abbrev_dir / "en.txt"
abbrev_file.write_text(mock_file_content)
with patch("haystack.components.preprocessors.sentence_tokenizer.Path") as mock_path:
mock_path.return_value.parent.parent.parent = tmp_path
result = SentenceSplitter._read_abbreviations("en")
assert result == ["Mr.", "Dr.", "Prof."]
def test_read_abbreviations_missing_file(caplog: LogCaptureFixture):
with patch("haystack.components.preprocessors.sentence_tokenizer.Path") as mock_path:
mock_path.return_value.parent.parent = Path("/nonexistent")
result = SentenceSplitter._read_abbreviations("pt")
assert result == []
assert "No abbreviations file found for pt. Using default abbreviations." in caplog.text
def test_quote_spans_regex():
# double quotes
text1 = 'He said "Hello world" and left.'
matches1 = list(QUOTE_SPANS_RE.finditer(text1))
assert len(matches1) == 1
assert matches1[0].group() == '"Hello world"'
# single quotes
text2 = "She replied 'Goodbye world' and smiled."
matches2 = list(QUOTE_SPANS_RE.finditer(text2))
assert len(matches2) == 1
assert matches2[0].group() == "'Goodbye world'"
# multiple quotes
text3 = 'First "quote" and second "quote" in same text.'
matches3 = list(QUOTE_SPANS_RE.finditer(text3))
assert len(matches3) == 2
assert matches3[0].group() == '"quote"'
assert matches3[1].group() == '"quote"'
# quotes containing newlines
text4 = 'Text with "quote\nspanning\nmultiple\nlines"'
matches4 = list(QUOTE_SPANS_RE.finditer(text4))
assert len(matches4) == 1
assert matches4[0].group() == '"quote\nspanning\nmultiple\nlines"'
# no quotes
text5 = "This text has no quotes."
matches5 = list(QUOTE_SPANS_RE.finditer(text5))
assert len(matches5) == 0
def test_split_sentences_performance() -> None:
# make sure our regex is not vulnerable to Regex Denial of Service (ReDoS)
# https://owasp.org/www-community/attacks/Regular_expression_Denial_of_Service_-_ReDoS
# this is a very long string, roughly 50 MB, but it should not take more than 2 seconds to process
splitter = SentenceSplitter()
text = " " + '"' * 20 + "A" * 50000000 + "B"
start = time.time()
_ = splitter.split_sentences(text)
end = time.time()
assert end - start < 2, f"Execution time exceeded 2 seconds: {end - start:.2f} seconds"
@@ -0,0 +1,74 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
from haystack.components.preprocessors import TextCleaner
def test_init_default():
cleaner = TextCleaner()
assert cleaner._remove_regexps is None
assert not cleaner._convert_to_lowercase
assert not cleaner._remove_punctuation
assert not cleaner._remove_numbers
assert cleaner._regex is None
assert cleaner._translator is None
def test_run():
cleaner = TextCleaner()
texts = ["Some text", "Some other text", "Yet another text"]
result = cleaner.run(texts=texts)
assert len(result) == 1
assert result["texts"] == texts
def test_run_with_empty_inputs():
cleaner = TextCleaner()
result = cleaner.run(texts=[])
assert len(result) == 1
assert result["texts"] == []
def test_run_with_regex():
cleaner = TextCleaner(remove_regexps=[r"\d+"])
result = cleaner.run(texts=["Open123 Source", "HaystackAI"])
assert len(result) == 1
assert result["texts"] == ["Open Source", "HaystackAI"]
def test_run_with_multiple_regexps():
cleaner = TextCleaner(remove_regexps=[r"\d+", r"[^\w\s]"])
result = cleaner.run(texts=["Open123! Source", "Haystack.AI"])
assert len(result) == 1
assert result["texts"] == ["Open Source", "HaystackAI"]
def test_run_with_convert_to_lowercase():
cleaner = TextCleaner(convert_to_lowercase=True)
result = cleaner.run(texts=["Open123! Source", "Haystack.AI"])
assert len(result) == 1
assert result["texts"] == ["open123! source", "haystack.ai"]
def test_run_with_remove_punctuation():
cleaner = TextCleaner(remove_punctuation=True)
result = cleaner.run(texts=["Open123! Source", "Haystack.AI"])
assert len(result) == 1
assert result["texts"] == ["Open123 Source", "HaystackAI"]
def test_run_with_remove_numbers():
cleaner = TextCleaner(remove_numbers=True)
result = cleaner.run(texts=["Open123! Source", "Haystack.AI"])
assert len(result) == 1
assert result["texts"] == ["Open! Source", "Haystack.AI"]
def test_run_with_multiple_parameters():
cleaner = TextCleaner(
remove_regexps=[r"\d+", r"[^\w\s]"], convert_to_lowercase=True, remove_punctuation=True, remove_numbers=True
)
result = cleaner.run(texts=["Open%123. !$Source", "Haystack.AI##"])
assert len(result) == 1
assert result["texts"] == ["open source", "haystackai"]