feat: semantic chunking + pgvector hybrid (BM25+vector) retriever (#2553)

tests/parser/test_chunking_strategies.py

@@ -1,7 +1,9 @@
-"""Tests for the recursive / markdown / parent_child chunking strategies (D1)."""
+"""Tests for the recursive / markdown / parent_child / semantic strategies."""
 
 from __future__ import annotations
 
+from unittest.mock import patch
+
 import pytest
 
 from application.parser.chunking import Chunker
@@ -10,6 +12,7 @@ from application.parser.chunking_strategies import (
     MarkdownChunker,
     ParentChildChunker,
     RecursiveChunker,
+    SemanticChunker,
 )
 from application.parser.schema.base import Document
 from application.utils import get_encoding
@@ -28,12 +31,13 @@ class TestRegistration:
             ("recursive", RecursiveChunker),
             ("markdown", MarkdownChunker),
             ("parent_child", ParentChildChunker),
+            ("semantic", SemanticChunker),
         ):
             assert ChunkerCreator.chunkers.get(key) is cls
 
     def test_worker_kwargs_accepted(self):
         # The worker builds every strategy with the classic kwarg set.
-        for strat in ("recursive", "markdown", "parent_child"):
+        for strat in ("recursive", "markdown", "parent_child", "semantic"):
             chunker = ChunkerCreator.create_chunker(
                 strat,
                 chunking_strategy=strat,
@@ -44,11 +48,6 @@ class TestRegistration:
             assert chunker.max_tokens == 200
             assert chunker.min_tokens == 20
 
-    def test_semantic_is_not_registered(self):
-        # semantic is explicitly deferred.
-        with pytest.raises(ValueError, match="No chunker class found"):
-            ChunkerCreator.create_chunker("semantic")
-
 
 @pytest.mark.unit
 class TestRecursive:
@@ -128,6 +127,110 @@ class TestParentChild:
         assert all("parent_text" in c.extra_info for c in out)
 
 
+_EMB_TARGET = "application.vectorstore.base.EmbeddingsSingleton.get_instance"
+
+
+class _FakeEmbeddings:
+    def __init__(self, vectors):
+        self._vectors = vectors
+
+    def embed_documents(self, sentences):
+        return self._vectors
+
+
+@pytest.mark.unit
+class TestSemantic:
+    def test_breakpoint_forces_split(self):
+        # Two topics: sentences 0-1 vs 2-3, orthogonal embeddings between.
+        text = "Alpha one. Alpha two. Beta one. Beta two."
+        vectors = [[1.0, 0.0], [1.0, 0.0], [0.0, 1.0], [0.0, 1.0]]
+        chunker = SemanticChunker(max_tokens=2000, min_tokens=0)
+        with patch(_EMB_TARGET, return_value=_FakeEmbeddings(vectors)):
+            out = chunker.chunk([Document(text=text, doc_id="d")])
+        assert len(out) == 2
+        assert "Alpha" in out[0].text and "Beta" not in out[0].text
+        assert "Beta" in out[1].text and "Alpha" not in out[1].text
+
+    def test_no_breakpoint_single_chunk(self):
+        # Identical embeddings -> zero distances -> no split.
+        text = "Same one. Same two. Same three. Same four."
+        vectors = [[1.0, 0.0]] * 4
+        chunker = SemanticChunker(max_tokens=2000, min_tokens=0)
+        with patch(_EMB_TARGET, return_value=_FakeEmbeddings(vectors)):
+            out = chunker.chunk([Document(text=text, doc_id="d")])
+        assert len(out) == 1
+        assert out[0].extra_info["token_count"] == _tok(out[0].text)
+
+    def test_max_tokens_enforced(self):
+        # A single semantic group larger than max_tokens is hard-split.
+        long_sentence = "word " * 300 + "."
+        text = f"{long_sentence} {long_sentence}"
+        vectors = [[1.0, 0.0], [1.0, 0.0]]
+        chunker = SemanticChunker(max_tokens=40, min_tokens=0)
+        with patch(_EMB_TARGET, return_value=_FakeEmbeddings(vectors)):
+            out = chunker.chunk([Document(text=text, doc_id="d")])
+        assert len(out) > 1
+        for c in out:
+            assert _tok(c.text) <= 40
+
+    def test_min_tokens_merges_neighbours(self):
+        # Non-uniform distances yield several breakpoints and tiny groups,
+        # which must merge until they clear min_tokens.
+        text = "A. B. C. D. E. F."
+        vectors = [
+            [1.0, 0.0],
+            [0.0, 1.0],
+            [0.0, 1.0],
+            [1.0, 0.0],
+            [0.0, 1.0],
+            [0.0, 1.0],
+        ]
+        chunker = SemanticChunker(max_tokens=2000, min_tokens=8)
+        with patch(_EMB_TARGET, return_value=_FakeEmbeddings(vectors)):
+            out = chunker.chunk([Document(text=text, doc_id="d")])
+        assert len(out) < 6
+        assert _tok(out[0].text) >= 8
+
+    def test_embeddings_error_falls_back_to_recursive(self):
+        text = "First sentence here. Second sentence here. Third one."
+
+        def _boom(*args, **kwargs):
+            raise RuntimeError("model unavailable")
+
+        chunker = SemanticChunker(max_tokens=2000, min_tokens=0)
+        with patch(_EMB_TARGET, side_effect=_boom):
+            out = chunker.chunk([Document(text=text, doc_id="d")])
+        recursive = RecursiveChunker(max_tokens=2000, min_tokens=0)
+        expected = recursive.chunk([Document(text=text, doc_id="d")])
+        assert [c.text for c in out] == [c.text for c in expected]
+
+    def test_too_few_sentences_falls_back(self):
+        # A single sentence cannot be semantically split.
+        chunker = SemanticChunker(max_tokens=2000, min_tokens=0)
+        with patch(_EMB_TARGET, side_effect=AssertionError("must not embed")):
+            out = chunker.chunk([Document(text="just one sentence", doc_id="d")])
+        assert len(out) == 1
+        assert out[0].text.strip() == "just one sentence"
+
+    def test_source_and_extra_info_preserved(self):
+        text = "Alpha one. Alpha two. Beta one. Beta two."
+        vectors = [[1.0, 0.0], [1.0, 0.0], [0.0, 1.0], [0.0, 1.0]]
+        doc = Document(
+            text=text,
+            doc_id="d",
+            extra_info={"source": "file.md", "title": "T"},
+        )
+        chunker = SemanticChunker(max_tokens=2000, min_tokens=0)
+        with patch(_EMB_TARGET, return_value=_FakeEmbeddings(vectors)):
+            out = chunker.chunk([doc])
+        assert len(out) == 2
+        for c in out:
+            assert c.extra_info["source"] == "file.md"
+            assert c.extra_info["title"] == "T"
+            assert c.extra_info["token_count"] == _tok(c.text)
+            assert c.doc_id.startswith("d-")
+
+
 @pytest.mark.unit
 class TestClassicByteIdentical:
     def test_classic_chunk_unchanged(self):