fix: EmbeddingBasedDocumentSplitter can't split exactly-two-group docs

haystack/components/preprocessors/embedding_based_document_splitter.py

@@ -79,14 +79,17 @@ def __init__(
         language: Language = "en",
         use_split_rules: bool = True,
         extend_abbreviations: bool = True,
+        min_distance_for_two_groups: float | None = None,
     ) -> None:
         """
         Initialize EmbeddingBasedDocumentSplitter.
 
         :param document_embedder: The DocumentEmbedder to use for calculating embeddings.
         :param sentences_per_group: Number of sentences to group together before embedding.
         :param percentile: Percentile threshold for cosine distance. Distances above this percentile
-            are treated as break points.
+            are treated as break points. Does not apply to documents that tokenize into exactly two sentence
+            groups: with only one distance to evaluate there is no distribution to compute a percentile against,
+            so `min_distance_for_two_groups` is used instead for that case.
         :param min_length: Minimum length of splits in characters. Splits below this length will be merged.
         :param max_length: Maximum length of splits in characters. Splits above this length will be recursively split.
         :param language: Language for sentence tokenization.
@@ -95,6 +98,17 @@ def __init__(
         :param extend_abbreviations: If True, the abbreviations used by NLTK's PunktTokenizer are extended by a list
             of curated abbreviations. Currently supported languages are: en, de.
             If False, the default abbreviations are used.
+        :param min_distance_for_two_groups: Cosine distance threshold used only for documents that tokenize into
+            exactly two sentence groups. With a single distance to evaluate there is no distribution to compute
+            `percentile` against, so this fixed, absolute threshold is used instead: the lone gap becomes a split
+            point if its distance exceeds this value. There is no universally correct default: what counts as
+            "the same topic, reworded" versus "a genuine topic change" is specific to the embedding model in use,
+            and the two can differ in cosine distance by an order of magnitude or more between models (some
+            embedders compress all pairwise distances into a narrow band, others spread them out). Defaults to
+            `None`, meaning documents with exactly two sentence groups are never split, matching the behavior of
+            documents with only one group. To enable splitting for this case, measure cosine distances from your
+            own embedder on a few known same-topic and known different-topic sentence pairs, and set this to a
+            value between the two.
         """
         self.document_embedder = document_embedder
 
@@ -114,6 +128,10 @@ def __init__(
             raise ValueError("max_length must be greater than min_length.")
         self.max_length = max_length
 
+        if min_distance_for_two_groups is not None and not 0.0 <= min_distance_for_two_groups <= 2.0:
+            raise ValueError("min_distance_for_two_groups must be between 0.0 and 2.0.")
+        self.min_distance_for_two_groups = min_distance_for_two_groups
+
         self.language = language
         self.use_split_rules = use_split_rules
         self.extend_abbreviations = extend_abbreviations
@@ -367,6 +385,17 @@ def _find_split_points(self, embeddings: list[list[float]]) -> list[int]:
             )
             distances.append(distance)
 
+        # With a single distance (i.e. exactly two sentence groups) there is no distribution to compute a
+        # percentile against: np.percentile of a one-element list always returns that same element, so
+        # `distance > threshold` below would always compare the value to itself and could never be True.
+        # `percentile` cannot apply here, so we fall back to `min_distance_for_two_groups` instead, which
+        # defaults to `None` (never split) since there's no universally correct absolute distance threshold
+        # across embedding models -- see its docstring in `__init__`.
+        if len(distances) == 1:
+            if self.min_distance_for_two_groups is None:
+                return []
+            return [1] if distances[0] > self.min_distance_for_two_groups else []
+
         # Calculate threshold based on percentile
         threshold = np.percentile(distances, self.percentile * 100)
 
@@ -572,6 +601,7 @@ def to_dict(self) -> dict[str, Any]:
             language=self.language,
             use_split_rules=self.use_split_rules,
             extend_abbreviations=self.extend_abbreviations,
+            min_distance_for_two_groups=self.min_distance_for_two_groups,
         )
 
     @classmethod

releasenotes/notes/fix-embedding-splitter-two-groups-48201a43ca06f07b.yaml

@@ -0,0 +1,18 @@
+---
+fixes:
+  - |
+    `EmbeddingBasedDocumentSplitter` could never split a document that tokenized into exactly two sentence
+    groups, no matter how dissimilar the groups were or how low `percentile` was set. With a single distance
+    to compare, `numpy.percentile` always returns that same value, so the `distance > threshold` check was
+    comparing the value to itself and could never be `True`. `percentile` cannot decide this case, since it
+    requires a distribution of distances to be meaningful and a single gap doesn't provide one. Splitting a
+    two-group document is now available via the new `min_distance_for_two_groups` parameter (see below);
+    without it, two-group documents keep the previous, safe behavior of never being split.
+enhancements:
+  - |
+    Added `min_distance_for_two_groups` to `EmbeddingBasedDocumentSplitter`, a cosine-distance threshold used
+    only for documents that tokenize into exactly two sentence groups, where `percentile` cannot apply.
+    Defaults to `None` (never split this case). There is no universally correct default value: how far apart
+    "the same topic, reworded" and "a genuine topic change" sit in cosine distance is specific to the
+    embedding model in use and can vary by an order of magnitude or more between models, so this is left for
+    users to calibrate against their own embedder.

test/components/preprocessors/test_embedding_based_document_splitter.py

@@ -32,6 +32,7 @@ def test_init(self):
         assert splitter.percentile == 0.9
         assert splitter.min_length == 50
         assert splitter.max_length == 1000
+        assert splitter.min_distance_for_two_groups is None
 
     def test_init_invalid_sentences_per_group(self):
         mock_embedder = Mock()
@@ -43,6 +44,11 @@ def test_init_invalid_percentile(self):
         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_distance_for_two_groups(self):
+        mock_embedder = Mock()
+        with pytest.raises(ValueError, match="min_distance_for_two_groups must be between 0.0 and 2.0"):
+            EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder, min_distance_for_two_groups=2.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"):
@@ -170,6 +176,95 @@ def test_find_split_points(self):
         # Should find a split point after the second embedding (index 2)
         assert 2 in split_points
 
+    def test_find_split_points_two_groups_default_never_splits(self):
+        """
+        With exactly two sentence groups there is only one distance to compare, so np.percentile of that
+        single-element list always returns the value itself: `distance > threshold` is comparing the value
+        to itself and can never be True. `percentile` therefore cannot decide this case, and there is no
+        universally correct absolute distance threshold across embedding models to fall back on either (see
+        `min_distance_for_two_groups`'s docstring). The safe default is to never split here, matching the
+        pre-existing behavior for documents with a single sentence group -- even for maximally dissimilar
+        (orthogonal) embeddings, at the most aggressive percentile=0.0 setting.
+        """
+        mock_embedder = Mock()
+        splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder, percentile=0.0)
+        assert splitter.min_distance_for_two_groups is None
+
+        embeddings = [[1.0, 0.0], [0.0, 1.0]]  # orthogonal -> maximum possible cosine distance
+        assert splitter._find_split_points(embeddings) == []
+
+    def test_find_split_points_two_groups_splits_once_min_distance_for_two_groups_is_configured(self):
+        """
+        Splitting a two-group document is opt-in: once `min_distance_for_two_groups` is set, the lone
+        distance is compared against it directly, independent of `percentile` (which cannot apply here).
+        """
+        mock_embedder = Mock()
+        embeddings = [[1.0, 0.0], [0.0, 1.0]]  # orthogonal -> cosine distance of 1.0
+
+        for percentile in (0.0, 0.5, 0.95, 1.0):
+            splitter = EmbeddingBasedDocumentSplitter(
+                document_embedder=mock_embedder, percentile=percentile, min_distance_for_two_groups=0.5
+            )
+            assert splitter._find_split_points(embeddings) == [1], f"expected split at percentile={percentile}"
+
+    def test_find_split_points_two_identical_groups_do_not_split(self):
+        """Two identical groups (distance == 0) must not produce a split point even at the threshold's edge."""
+        mock_embedder = Mock()
+        splitter = EmbeddingBasedDocumentSplitter(
+            document_embedder=mock_embedder, percentile=0.0, min_distance_for_two_groups=0.0
+        )
+
+        embeddings = [[1.0, 0.0], [1.0, 0.0]]
+        assert splitter._find_split_points(embeddings) == []
+
+    def test_find_split_points_two_near_duplicate_groups_do_not_split_regardless_of_percentile(self):
+        """
+        Regression test: comparing the lone distance to itself (the original bug) or splitting on any
+        nonzero distance (an earlier, overly aggressive fix) both fail on near-duplicate groups. Two groups
+        that are almost, but not exactly, identical (a tiny cosine distance well under a realistic
+        `min_distance_for_two_groups`) should not split, and since `percentile` does not apply to the
+        two-group case, this must hold at every percentile setting, including the most aggressive (0.0).
+        """
+        mock_embedder = Mock()
+        embeddings = [[1.0, 0.0], [1.0, 0.001]]  # cosine distance ~5e-7: near-duplicate, not identical
+
+        for percentile in (0.0, 0.5, 0.95, 1.0):
+            splitter = EmbeddingBasedDocumentSplitter(
+                document_embedder=mock_embedder, percentile=percentile, min_distance_for_two_groups=0.3
+            )
+            assert splitter._find_split_points(embeddings) == [], f"unexpected split at percentile={percentile}"
+
+    def test_run_splits_document_with_exactly_two_sentence_groups(self):
+        """
+        End-to-end regression test: a document that tokenizes into exactly two sentence groups with
+        maximally dissimilar content is split by run() once `min_distance_for_two_groups` is configured,
+        even at the most aggressive percentile=0.0 setting.
+        """
+        text = (
+            "Astronomy studies the stars and galaxies of outer space. "
+            "Sourdough baking relies on wild yeast fermentation."
+        )
+
+        def mock_run(documents):
+            embeddings = [[1.0, 0.0] if i == 0 else [0.0, 1.0] for i in range(len(documents))]
+            return {"documents": [replace(doc, embedding=emb) for doc, emb in zip(documents, embeddings, strict=True)]}
+
+        mock_embedder = Mock()
+        mock_embedder.run = Mock(side_effect=mock_run)
+
+        splitter = EmbeddingBasedDocumentSplitter(
+            document_embedder=mock_embedder,
+            sentences_per_group=1,
+            percentile=0.0,
+            min_length=0,
+            max_length=10000,
+            min_distance_for_two_groups=0.5,
+        )
+        splitter.warm_up()
+
+        result = splitter.run(documents=[Document(content=text)])
+        assert len(result["documents"]) == 2
+
     def test_create_splits_from_points(self):
         mock_embedder = Mock()
         splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder)
@@ -379,6 +474,7 @@ def test_to_dict(self):
         assert result["init_parameters"]["percentile"] == 0.9
         assert result["init_parameters"]["min_length"] == 50
         assert result["init_parameters"]["max_length"] == 1000
+        assert result["init_parameters"]["min_distance_for_two_groups"] is None
         assert "document_embedder" in result["init_parameters"]
 
     @pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")