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
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:
@@ -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\fThis is a header.
|
||||
Page of
|
||||
2
|
||||
4
|
||||
Lorem ipsum dolor sit amet
|
||||
This is a footer number 1
|
||||
This is footer number 2This is a header.
|
||||
Page of
|
||||
3
|
||||
4
|
||||
Sid ut perspiciatis unde
|
||||
This is a footer number 1
|
||||
This is footer number 2This 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
|
||||
مرحبا بالعالم
|
||||
em Space"""
|
||||
|
||||
expected_text_NFC = """\
|
||||
アイウエオ
|
||||
Comment ça va
|
||||
مرحبا بالعالم
|
||||
em Space"""
|
||||
|
||||
expected_text_NFD = """\
|
||||
アイウエオ
|
||||
Comment ça va
|
||||
مرحبا بالعالم
|
||||
em Space"""
|
||||
|
||||
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
|
||||
Á
|
||||
مرحبا بالعالم
|
||||
em Space"""
|
||||
|
||||
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"]
|
||||
Reference in New Issue
Block a user