diff --git a/haystack/components/preprocessors/embedding_based_document_splitter.py b/haystack/components/preprocessors/embedding_based_document_splitter.py index 17acefd3b7..34e74cfa09 100644 --- a/haystack/components/preprocessors/embedding_based_document_splitter.py +++ b/haystack/components/preprocessors/embedding_based_document_splitter.py @@ -445,6 +445,16 @@ def _merge_small_splits(self, splits: list[str]) -> list[str]: # Don't forget the last split merged.append(current_split) + # The loop only merges forward, so the final split can still be below min_length. Merge it backwards, + # subject to the same max_length limit as forward merges. + if ( + len(merged) > 1 + and len(merged[-1]) < self.min_length + and len(merged[-2]) + len(merged[-1]) < self.max_length + ): + trailing_split = merged.pop() + merged[-1] += trailing_split + return merged def _split_large_splits(self, splits: list[str]) -> list[str]: diff --git a/releasenotes/notes/merge-short-trailing-split-8f5b640fdbc3ccbe.yaml b/releasenotes/notes/merge-short-trailing-split-8f5b640fdbc3ccbe.yaml new file mode 100644 index 0000000000..456a4d4603 --- /dev/null +++ b/releasenotes/notes/merge-short-trailing-split-8f5b640fdbc3ccbe.yaml @@ -0,0 +1,7 @@ +--- +fixes: + - | + Fixed ``EmbeddingBasedDocumentSplitter`` emitting a final split shorter than ``min_length``. Small splits + were only merged forward, so the last one had nothing left to absorb and was returned as its own document. + It is now merged into the preceding split, unless doing so would reach ``max_length``, the same limit that + already governs forward merges. diff --git a/test/components/preprocessors/test_embedding_based_document_splitter.py b/test/components/preprocessors/test_embedding_based_document_splitter.py index 50db2a8aa8..968a770e35 100644 --- a/test/components/preprocessors/test_embedding_based_document_splitter.py +++ b/test/components/preprocessors/test_embedding_based_document_splitter.py @@ -215,6 +215,36 @@ def test_merge_small_splits_respect_max_length(self): # Second split is merged with third split to get above min_length and still beneath max_length assert merged[1] == "1234567891234" + def test_merge_small_splits_merges_short_trailing_split(self): + mock_embedder = Mock() + splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder, min_length=10) + + # The loop only merges forward, so the final accumulator has nothing left to absorb. + splits = ["Long enough text ", "Ok."] + merged = splitter._merge_small_splits(splits=splits) + + assert merged == ["Long enough text Ok."] + + def test_merge_small_splits_keeps_short_trailing_split_when_max_length_blocks(self): + mock_embedder = Mock() + splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder, min_length=10, max_length=15) + + # Merging backwards would reach max_length, so the short tail stays on its own, + # matching how a blocked forward merge already behaves. + splits = ["123456789012", "1234"] + merged = splitter._merge_small_splits(splits=splits) + + assert merged == ["123456789012", "1234"] + + def test_merge_small_splits_keeps_a_lone_short_split(self): + mock_embedder = Mock() + splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder, min_length=10) + + # Nothing to merge into. + merged = splitter._merge_small_splits(splits=["Ok."]) + + assert merged == ["Ok."] + def test_create_documents_from_splits(self): mock_embedder = Mock() splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)