feat: optimise embeds

tests/parser/conftest.py

@@ -0,0 +1,23 @@
+import pytest
+
+
+@pytest.fixture(scope="session")
+def _mpnet_tokenizer():
+    """Real WordPiece tokenizer; skipped when the hub is unreachable."""
+    pytest.importorskip("tokenizers")
+    from tokenizers import Tokenizer
+
+    try:
+        tokenizer = Tokenizer.from_pretrained("sentence-transformers/all-mpnet-base-v2")
+    except Exception as exc:  # offline CI
+        pytest.skip(f"tokenizer unavailable: {exc}")
+    tokenizer.no_padding()
+    tokenizer.no_truncation()
+    return tokenizer
+
+
+@pytest.fixture
+def hf_counter(_mpnet_tokenizer):
+    from application.parser.tokenization import HuggingFaceCounter
+
+    return HuggingFaceCounter(_mpnet_tokenizer, "sentence-transformers/all-mpnet-base-v2")

tests/parser/test_chunking.py

@@ -108,18 +108,24 @@ class TestSplitDocument:
 
         result = chunker.split_document(doc)
         assert len(result) > 1
-        # First chunk should contain header
         assert "h1" in result[0].text
+        # Only the first: this is what duplicate_headers=False means.
+        assert all("h1" not in chunk.text for chunk in result[1:])
 
     def test_split_duplicates_header(self):
+        """Every chunk carries the header when the flag is set.
+
+        Asserting only the first chunk passed even while the flag did nothing:
+        the implementation cleared the header after the first iteration, so
+        duplicate_headers was unreachable.
+        """
         chunker = Chunker(max_tokens=50, min_tokens=5, duplicate_headers=True)
         text = "h1\nh2\nh3\n" + "word " * 200
         doc = Document(text=text, doc_id="doc1")
 
         result = chunker.split_document(doc)
         assert len(result) > 1
-        # First chunk should contain header
-        assert "h1" in result[0].text
+        assert all("h1" in chunk.text for chunk in result)
 
     def test_split_preserves_embedding(self):
         chunker = Chunker(max_tokens=50, min_tokens=5)

tests/parser/test_chunking_strategies.py

@@ -15,11 +15,13 @@ from application.parser.chunking_strategies import (
     SemanticChunker,
 )
 from application.parser.schema.base import Document
-from application.utils import get_encoding
+from application.parser.tokenization import get_token_counter
 
 
 def _tok(text: str) -> int:
-    return len(get_encoding().encode(text))
+    # Chunk budgets are enforced in the embedding model's tokenizer, so a
+    # test that measures them in cl100k measures the wrong thing.
+    return get_token_counter().count(text)
 
 
 @pytest.mark.unit

tests/parser/test_tokenization.py

@@ -0,0 +1,107 @@
+"""Chunk sizes must be counted in the embedding model's units, and splitting
+must never rewrite the text it splits."""
+
+import pytest
+
+from application.parser import tokenization
+from application.parser.tokenization import (
+    HuggingFaceCounter,
+    TiktokenCounter,
+    get_token_counter,
+)
+
+SAMPLES = [
+    "Hello World: DocsGPT ANSWERS Questions.",
+    "The quick brown fox jumps over the lazy dog. " * 40,
+    "Comment configurer l'authentification avec une clé API ?",
+    "def embed(text: str) -> list[float]:\n    return model.encode(text)\n",
+    "Ünïcödé — em-dashes, curly “quotes”, and 日本語 text.",
+    "a,b,c\n1,2,3\n4,5,6\n" * 30,
+]
+
+
+class _StubEncoding:
+    """Whitespace tokenizer standing in for tiktoken."""
+
+    def encode_ordinary(self, text):
+        return [ord(c) for c in text]
+
+    def decode(self, ids):
+        return "".join(chr(i) for i in ids)
+
+
+@pytest.fixture(autouse=True)
+def _clear_cache():
+    tokenization.reset_cache()
+    yield
+    tokenization.reset_cache()
+
+
+class TestSplittingPreservesText:
+    """The property that protects every stored document."""
+
+    @pytest.mark.parametrize("text", SAMPLES)
+    def test_tiktoken_split_reassembles_exactly(self, text, monkeypatch):
+        monkeypatch.setattr(tokenization, "get_encoding", lambda: _StubEncoding())
+        counter = TiktokenCounter()
+        pieces = counter.split(text, 7)
+        assert "".join(pieces) == text
+
+    @pytest.mark.parametrize("text", SAMPLES)
+    def test_hf_split_reassembles_exactly(self, text, hf_counter):
+        """WordPiece lowercases on decode, so splitting must slice, not decode."""
+        pieces = hf_counter.split(text, 7)
+        assert "".join(pieces) == text
+
+    def test_hf_split_does_not_lowercase(self, hf_counter):
+        text = "Hello World: DocsGPT ANSWERS Questions."
+        assert "".join(hf_counter.split(text, 3)) == text
+        assert "DocsGPT" in "".join(hf_counter.split(text, 3))
+
+    @pytest.mark.parametrize("text", SAMPLES)
+    def test_every_piece_is_within_budget(self, text, hf_counter):
+        budget = 10
+        for piece in hf_counter.split(text, budget):
+            # The final piece can absorb trailing characters the tokenizer
+            # dropped, so allow a small overshoot there only.
+            assert hf_counter.count(piece) <= budget + 2
+
+    def test_short_text_is_returned_whole(self, hf_counter):
+        assert hf_counter.split("short", 100) == ["short"]
+
+    def test_empty_text_yields_no_pieces(self, hf_counter):
+        assert hf_counter.split("", 10) == []
+
+    def test_zero_budget_is_clamped_not_infinite_loop(self, hf_counter):
+        pieces = hf_counter.split("some words here to split", 0)
+        assert "".join(pieces) == "some words here to split"
+
+
+class TestCounting:
+    def test_counts_differ_between_tokenizers(self, hf_counter, monkeypatch):
+        """The whole point: mpnet and cl100k disagree, so units matter."""
+        monkeypatch.setattr(tokenization, "get_encoding", lambda: _StubEncoding())
+        text = "internationalisation tokenization"
+        assert hf_counter.count(text) != TiktokenCounter().count(text)
+
+    def test_empty_text_counts_zero(self, hf_counter):
+        assert hf_counter.count("") == 0
+
+
+class TestSelection:
+    def test_registered_model_uses_its_own_tokenizer(self):
+        counter = get_token_counter("huggingface_sentence-transformers/all-mpnet-base-v2")
+        assert isinstance(counter, HuggingFaceCounter)
+        assert counter.name == "sentence-transformers/all-mpnet-base-v2"
+
+    def test_openai_model_falls_back_to_cl100k(self):
+        """OpenAI models are served remotely and genuinely count cl100k."""
+        assert isinstance(get_token_counter("openai_text-embedding-ada-002"), TiktokenCounter)
+
+    def test_unreachable_tokenizer_falls_back_rather_than_raising(self, monkeypatch):
+        monkeypatch.setattr(tokenization, "_load_hf_counter", lambda repo: None)
+        assert isinstance(get_token_counter("granite-311m"), TiktokenCounter)
+
+    def test_counter_is_cached_per_model(self):
+        first = get_token_counter("granite-311m")
+        assert get_token_counter("granite-311m") is first