feat: ingestion optimisations

tests/parser/file/test_embedding_pipeline.py

@@ -3,8 +3,11 @@ import logging
 from unittest.mock import patch, MagicMock
 
 from application.parser.embedding_pipeline import (
+    DEFAULT_EMBEDDINGS_BATCH_SIZE,
     EmbeddingPipelineError,
+    _resolve_batch_size,
     add_text_to_store_with_retry,
+    add_texts_to_store_with_retry,
     assert_index_complete,
     embed_and_store_documents,
     sanitize_content,
@@ -123,7 +126,7 @@ def test_embed_and_store_documents_progress_band(
     assert currents == sorted(currents)
 
 
-@patch("application.parser.embedding_pipeline.add_text_to_store_with_retry")
+@patch("application.parser.embedding_pipeline.add_texts_to_store_with_retry")
 def test_embed_and_store_documents_partial_failure_raises(
     mock_add_retry, tmp_path, mock_settings, mock_vector_creator, caplog
 ):
@@ -146,9 +149,11 @@ def test_embed_and_store_documents_partial_failure_raises(
     mock_vector_creator.create_vectorstore.return_value = mock_store
 
     # First document succeeds (FAISS init seeds with docs[0]; the loop
-    # picks up at idx=1 and raises on the bad chunk).
-    def side_effect(*args, **kwargs):
-        if "bad" in args[1].page_content:
+    # picks up at idx=1 and raises on the bad chunk). The batch entry point
+    # receives a list, and the per-chunk fallback re-runs it one at a time —
+    # both go through this mock, so "bad" raises either way.
+    def side_effect(store_arg, docs_arg, source_arg):
+        if any("bad" in d.page_content for d in docs_arg):
             raise RuntimeError("Embedding failed")
     mock_add_retry.side_effect = side_effect
 
@@ -165,7 +170,7 @@ def test_embed_and_store_documents_partial_failure_raises(
     mock_store.save_local.assert_called()
 
 
-@patch("application.parser.embedding_pipeline.add_text_to_store_with_retry")
+@patch("application.parser.embedding_pipeline.add_texts_to_store_with_retry")
 def test_embed_and_store_documents_all_chunks_succeed_no_raise(
     mock_add_retry, tmp_path, mock_settings, mock_vector_creator,
 ):
@@ -291,3 +296,143 @@ def test_embed_and_store_documents_save_fails_raises_oserror(
     with pytest.raises(OSError, match="Unable to save vector store"):
         embed_and_store_documents(docs, str(folder_name), source_id, task_status)
 
+
+
+# ── batched embed loop ─────────────────────────────────────────────────────
+
+
+def test_add_texts_to_store_with_retry_sends_one_call_per_batch():
+    """The batch entry point collapses N chunks into a single add_texts."""
+    store = MagicMock()
+    docs = [MagicMock(page_content=f"c{i}", metadata={}) for i in range(3)]
+
+    add_texts_to_store_with_retry(store, docs, "sid")
+
+    store.add_texts.assert_called_once_with(
+        ["c0", "c1", "c2"],
+        metadatas=[{"source_id": "sid"}] * 3,
+    )
+
+
+def test_add_texts_to_store_with_retry_sanitizes_and_skips_empty():
+    store = MagicMock()
+    docs = [MagicMock(page_content="a\x00b", metadata={})]
+
+    add_texts_to_store_with_retry(store, docs, "sid")
+    assert store.add_texts.call_args.args[0] == ["ab"]
+
+    store.reset_mock()
+    add_texts_to_store_with_retry(store, [], "sid")
+    store.add_texts.assert_not_called()
+
+
+def test_resolve_batch_size_falls_back_on_bad_setting(monkeypatch):
+    fake = MagicMock()  # attribute access yields a MagicMock, not an int
+    monkeypatch.setattr("application.parser.embedding_pipeline.settings", fake)
+    assert _resolve_batch_size() == DEFAULT_EMBEDDINGS_BATCH_SIZE
+
+    fake.EMBEDDINGS_BATCH_SIZE = 0
+    assert _resolve_batch_size() == 1  # never below 1
+    fake.EMBEDDINGS_BATCH_SIZE = 32
+    assert _resolve_batch_size() == 32
+
+
+def test_embed_loop_batches_chunks(tmp_path, mock_settings, mock_vector_creator):
+    """70 chunks at batch size 32 => 3 add_texts calls, not 70."""
+    mock_settings.VECTOR_STORE = "chromadb"
+    mock_settings.EMBEDDINGS_BATCH_SIZE = 32
+
+    docs = [MagicMock(page_content=f"d{i}", metadata={}) for i in range(70)]
+    store = MagicMock()
+    mock_vector_creator.create_vectorstore.return_value = store
+
+    with patch("application.parser.embedding_pipeline._record_progress") as rec:
+        embed_and_store_documents(
+            docs, str(tmp_path / "s"), "sid", MagicMock(),
+        )
+
+    assert store.add_texts.call_count == 3
+    assert [len(c.args[0]) for c in store.add_texts.call_args_list] == [32, 32, 6]
+    # One checkpoint per batch, and the final one accounts for every chunk.
+    assert rec.call_count == 3
+    assert rec.call_args.kwargs == {"last_index": 69, "embedded_chunks": 70}
+
+
+def test_embed_loop_batch_size_one_matches_legacy(
+    tmp_path, mock_settings, mock_vector_creator
+):
+    """batch_size=1 restores the pre-batching one-call-per-chunk behaviour."""
+    mock_settings.VECTOR_STORE = "chromadb"
+    mock_settings.EMBEDDINGS_BATCH_SIZE = 1
+
+    docs = [MagicMock(page_content=f"d{i}", metadata={}) for i in range(5)]
+    store = MagicMock()
+    mock_vector_creator.create_vectorstore.return_value = store
+
+    embed_and_store_documents(docs, str(tmp_path / "s"), "sid", MagicMock())
+
+    assert store.add_texts.call_count == 5
+
+
+def test_poison_chunk_isolated_by_per_chunk_fallback(
+    tmp_path, mock_settings, mock_vector_creator
+):
+    """A batch failure re-runs individually: good chunks land, and the
+    reported failure index is the real offender, not the batch head.
+
+    Patches the batch entry point so the ``@retry`` sleeps don't run — the
+    fallback, not the retry, is what's under test here.
+    """
+    mock_settings.VECTOR_STORE = "chromadb"
+    mock_settings.EMBEDDINGS_BATCH_SIZE = 32
+
+    docs = [MagicMock(page_content=f"d{i}", metadata={}) for i in range(10)]
+    docs[6].page_content = "poison"
+    mock_vector_creator.create_vectorstore.return_value = MagicMock()
+
+    def fake_add(store, batch, source_id):
+        if any(d.page_content == "poison" for d in batch):
+            raise RuntimeError("input too large")
+
+    with patch(
+        "application.parser.embedding_pipeline.add_texts_to_store_with_retry",
+        side_effect=fake_add,
+    ):
+        with patch("application.parser.embedding_pipeline._record_progress") as rec:
+            with pytest.raises(EmbeddingPipelineError) as exc:
+                embed_and_store_documents(
+                    docs, str(tmp_path / "s"), "sid", MagicMock(),
+                )
+
+    assert "chunk 6/10" in str(exc.value)
+    assert isinstance(exc.value.__cause__, RuntimeError)
+    # Chunks 0-5 were salvaged one at a time and checkpointed.
+    assert rec.call_args_list[-1].kwargs == {"last_index": 5, "embedded_chunks": 6}
+
+
+def test_batch_only_failure_recovers_via_fallback(
+    tmp_path, mock_settings, mock_vector_creator
+):
+    """When the *batch* is rejected but each chunk is fine on its own (e.g. a
+    request-size limit), the fallback completes the ingest without raising."""
+    mock_settings.VECTOR_STORE = "chromadb"
+    mock_settings.EMBEDDINGS_BATCH_SIZE = 32
+
+    docs = [MagicMock(page_content=f"d{i}", metadata={}) for i in range(4)]
+    mock_vector_creator.create_vectorstore.return_value = MagicMock()
+
+    seen = []
+
+    def fake_add(store, batch, source_id):
+        seen.append(len(batch))
+        if len(batch) > 1:
+            raise RuntimeError("payload too large")
+
+    with patch(
+        "application.parser.embedding_pipeline.add_texts_to_store_with_retry",
+        side_effect=fake_add,
+    ):
+        embed_and_store_documents(docs, str(tmp_path / "s"), "sid", MagicMock())
+
+    # One rejected batch of 4, then four successful singles.
+    assert seen == [4, 1, 1, 1, 1]

tests/parser/remote/test_github_loader.py

@@ -91,9 +91,11 @@ class TestGitHubLoaderLoadData:
         loader = GitHubLoader()
 
         # Stub out network-dependent methods
-        monkeypatch.setattr(loader, "fetch_repo_files", lambda repo, path="": [
-            "README.md", "src/main.py"
-        ])
+        monkeypatch.setattr(loader, "get_default_branch", lambda repo: "main")
+        monkeypatch.setattr(
+            loader, "fetch_repo_tree",
+            lambda repo, branch: ([("README.md", 10), ("src/main.py", 10)], False),
+        )
 
         def fake_fetch_content(repo, file_path):
             return f"content for {file_path}"
@@ -257,8 +259,10 @@ class TestGitHubLoaderFetchFileContentEdgeCases:
 class TestGitHubLoaderLoadDataSkipsNone:
     def test_skips_binary_files(self, monkeypatch):
         loader = GitHubLoader()
+        monkeypatch.setattr(loader, "get_default_branch", lambda repo: "main")
         monkeypatch.setattr(
-            loader, "fetch_repo_files", lambda repo, path="": ["a.py", "b.png"]
+            loader, "fetch_repo_tree",
+            lambda repo, branch: ([("a.py", 10), ("b.png", 10)], False),
         )
 
         def fake_content(repo, fp):
@@ -312,3 +316,199 @@ class TestGitHubLoaderRobustness:
         mock_get.return_value = make_response({"encoding": "base64", "content": "AAA"})
         result = GitHubLoader().fetch_file_content("owner/repo", "bigfile.bin")
         assert result is None
+
+
+class TestGitHubLoaderNormalizeRepo:
+    @pytest.mark.parametrize("raw,expected", [
+        ("https://github.com/owner/repo", "owner/repo"),
+        ("https://github.com/owner/repo.git", "owner/repo"),
+        ("https://github.com/owner/repo/", "owner/repo"),
+        ("http://github.com/owner/repo.git/", "owner/repo"),
+        ("owner/repo", "owner/repo"),
+        ("https://github.com/owner/repo/tree/main/sub", "owner/repo"),
+    ])
+    def test_normalizes(self, raw, expected):
+        assert GitHubLoader.normalize_repo(raw) == expected
+
+    def test_rejects_non_repo_url(self):
+        """Regression: a pasted website URL used to be concatenated straight
+        into the contents path and 404-loop."""
+        loader = GitHubLoader()
+        with pytest.raises(ValueError, match="Not a valid GitHub repository"):
+            loader.load_data("https://suat-handbook.netlify.app/")
+
+
+class TestGitHubLoaderSkipPaths:
+    @pytest.mark.parametrize("path", [
+        "node_modules/left-pad/index.js",
+        "web/dist/app.js",
+        "target/scala-2.13/Foo.class",
+        "vendor/github.com/pkg/errors/errors.go",
+        "package-lock.json",
+        "assets/app.min.js",
+        "build/output.map",
+        ".venv/lib/thing.py",
+        "src/__pycache__/mod.pyc",
+    ])
+    def test_skipped(self, path):
+        assert GitHubLoader().should_skip_path(path) is True
+
+    @pytest.mark.parametrize("path", [
+        "README.md",
+        "src/main/scala/zio/ZIO.scala",
+        ".github/workflows/ci.yml",
+        "docs/guide/setup.md",
+        "Cargo.toml",
+    ])
+    def test_kept(self, path):
+        assert GitHubLoader().should_skip_path(path) is False
+
+
+class TestGitHubLoaderIsTextFileAdditions:
+    @pytest.mark.parametrize("path", [
+        "Cargo.toml", "main.go", "app.vue", "schema.proto",
+        "infra.tf", "Dockerfile", "Makefile", "build.gradle",
+    ])
+    def test_newly_recognised_text(self, path):
+        assert GitHubLoader().is_text_file(path) is True
+
+    def test_env_files_are_not_ingested(self):
+        """.env was in the allowlist, so a committed secrets file would be
+        embedded verbatim into the vector index."""
+        loader = GitHubLoader()
+        assert loader.is_text_file("config/.env") is False
+        assert loader.is_text_file(".env") is False
+
+
+class TestGitHubLoaderSelectFiles:
+    def test_applies_size_cap(self, monkeypatch):
+        loader = GitHubLoader()
+        monkeypatch.setattr(
+            "application.parser.remote.github_loader.settings.GITHUB_INGEST_MAX_FILE_BYTES",
+            100, raising=False,
+        )
+        entries = [("small.py", 50), ("huge.py", 5000), ("ok.md", 99)]
+        assert loader.select_files(entries) == ["small.py", "ok.md"]
+
+    def test_zero_cap_disables_limit(self, monkeypatch):
+        loader = GitHubLoader()
+        monkeypatch.setattr(
+            "application.parser.remote.github_loader.settings.GITHUB_INGEST_MAX_FILE_BYTES",
+            0, raising=False,
+        )
+        assert loader.select_files([("huge.py", 10**9)]) == ["huge.py"]
+
+    def test_filters_binaries_and_vendored(self):
+        entries = [
+            ("README.md", 10), ("logo.png", 10),
+            ("node_modules/x/i.js", 10), ("src/a.py", 10),
+        ]
+        assert GitHubLoader().select_files(entries) == ["README.md", "src/a.py"]
+
+
+class TestGitHubLoaderTree:
+    @patch("application.parser.remote.github_loader.requests.get")
+    def test_single_request_lists_all_blobs(self, mock_get):
+        mock_get.return_value = make_response({
+            "tree": [
+                {"path": "README.md", "type": "blob", "size": 12},
+                {"path": "src", "type": "tree"},
+                {"path": "src/a.py", "type": "blob", "size": 34},
+            ],
+            "truncated": False,
+        })
+        entries, truncated = GitHubLoader().fetch_repo_tree("owner/repo", "main")
+
+        assert entries == [("README.md", 12), ("src/a.py", 34)]
+        assert truncated is False
+        # One call for the whole repo, versus one per directory before.
+        assert mock_get.call_count == 1
+
+    @patch("application.parser.remote.github_loader.requests.get")
+    def test_truncated_tree_falls_back_to_walk(self, mock_get, monkeypatch):
+        loader = GitHubLoader()
+        mock_get.return_value = make_response({"tree": [], "truncated": True})
+        monkeypatch.setattr(
+            loader, "fetch_repo_files", lambda repo, path="": ["README.md"]
+        )
+        assert loader._list_candidate_files("owner/repo", "main") == ["README.md"]
+
+
+class TestGitHubLoaderDefaultBranch:
+    @patch("application.parser.remote.github_loader.requests.get")
+    def test_uses_repo_default_branch(self, mock_get):
+        mock_get.return_value = make_response({"default_branch": "master"})
+        assert GitHubLoader().get_default_branch("owner/repo") == "master"
+
+    @patch("application.parser.remote.github_loader.requests.get")
+    def test_falls_back_to_main(self, mock_get):
+        mock_get.side_effect = requests.ConnectionError("boom")
+        assert GitHubLoader().get_default_branch("owner/repo") == "main"
+
+    def test_citation_url_uses_real_branch(self, monkeypatch):
+        """Regression: blob/main was hard-coded, so every master-default
+        repo produced dead source links."""
+        loader = GitHubLoader()
+        monkeypatch.setattr(loader, "get_default_branch", lambda repo: "master")
+        monkeypatch.setattr(
+            loader, "fetch_repo_tree", lambda repo, branch: ([("a.py", 5)], False)
+        )
+        monkeypatch.setattr(loader, "fetch_file_content", lambda r, p: "code")
+
+        docs = loader.load_data("https://github.com/owner/repo")
+        assert docs[0].extra_info["source"] == (
+            "https://github.com/owner/repo/blob/master/a.py"
+        )
+
+
+class TestGitHubLoaderParallelFetch:
+    def test_fetches_in_parallel_and_preserves_order(self, monkeypatch):
+        loader = GitHubLoader()
+        monkeypatch.setattr(loader, "get_default_branch", lambda repo: "main")
+        paths = [f"f{i}.py" for i in range(10)]
+        monkeypatch.setattr(
+            loader, "fetch_repo_tree",
+            lambda repo, branch: ([(p, 5) for p in paths], False),
+        )
+        monkeypatch.setattr(
+            loader, "fetch_file_content", lambda r, p: f"body {p}"
+        )
+
+        docs = loader.load_data("https://github.com/owner/repo")
+        assert [d.doc_id for d in docs] == paths
+
+    def test_one_bad_file_does_not_sink_the_ingest(self, monkeypatch):
+        loader = GitHubLoader()
+        monkeypatch.setattr(loader, "get_default_branch", lambda repo: "main")
+        monkeypatch.setattr(
+            loader, "fetch_repo_tree",
+            lambda repo, branch: ([("good.py", 5), ("bad.py", 5)], False),
+        )
+
+        def flaky(repo, path):
+            if path == "bad.py":
+                raise requests.HTTPError("500")
+            return "code"
+        monkeypatch.setattr(loader, "fetch_file_content", flaky)
+
+        docs = loader.load_data("https://github.com/owner/repo")
+        assert [d.doc_id for d in docs] == ["good.py"]
+
+
+class TestGitHubLoaderStaleTokenFallback:
+    @patch("application.parser.remote.github_loader.requests.get")
+    def test_401_retries_unauthenticated(self, mock_get):
+        """An expired PAT 401s even public repos; fall back to anonymous
+        rather than failing the ingest outright."""
+        loader = GitHubLoader()
+        loader.access_token = "stale-token"
+        unauthorized = MagicMock(status_code=401)
+        ok = make_response({"ok": True}, 200)
+        mock_get.side_effect = [unauthorized, ok]
+
+        resp = loader._make_request("https://api.github.com/repos/o/r")
+
+        assert resp.status_code == 200
+        assert mock_get.call_count == 2
+        # Second attempt carried no Authorization header.
+        assert "Authorization" not in mock_get.call_args.kwargs["headers"]

tests/test_worker_utils.py

@@ -311,3 +311,79 @@ class TestUploadIndex:
         mock_post.assert_called_once()
         files = mock_post.call_args.kwargs["files"]
         assert "file_faiss" in files and "file_pkl" in files
+
+
+# ── directory_structure accounting ─────────────────────────────────────────
+
+
+class TestCountStructureFiles:
+    """Regression: the remote-ingest completion log used
+    ``len(directory_structure)``, which counts top-level entries only — a
+    1,474-file repo logged "44 files"."""
+
+    def test_counts_nested_leaves_not_top_level_keys(self):
+        from application.worker import count_structure_files
+
+        structure = {
+            "README.md": {"type": "text/markdown", "size_bytes": 10, "token_count": 3},
+            "src": {
+                "main": {
+                    "a.py": {"type": "text/x-python", "size_bytes": 5, "token_count": 2},
+                    "b.py": {"type": "text/x-python", "size_bytes": 5, "token_count": 2},
+                },
+                "c.py": {"type": "text/x-python", "size_bytes": 5, "token_count": 2},
+            },
+        }
+        assert len(structure) == 2  # the old, wrong number
+        assert count_structure_files(structure) == 4
+
+    def test_empty_and_non_dict_are_zero(self):
+        from application.worker import count_structure_files
+
+        assert count_structure_files({}) == 0
+        assert count_structure_files(None) == 0
+        assert count_structure_files({"empty_dir": {}}) == 0
+
+
+class TestDirectoryStructureAccumulatesChunks:
+    """Regression: ``directory_structure`` was built from *chunks* but
+    assigned per path, so a multi-chunk file kept only its last fragment's
+    size and token count."""
+
+    def test_multi_chunk_file_sums_its_chunks(self):
+        from application.worker import add_file_to_structure, count_structure_files
+
+        structure: dict = {}
+        for tokens, size in [(900, 3600), (900, 3600), (120, 480)]:
+            add_file_to_structure(
+                structure, "docs/guide.md", "text/markdown",
+                size_bytes=size, token_count=tokens,
+            )
+
+        leaf = structure["docs"]["guide.md"]
+        assert leaf["token_count"] == 1920   # not 120, the last chunk
+        assert leaf["size_bytes"] == 7680    # not 480
+        assert leaf["type"] == "text/markdown"
+        assert count_structure_files(structure) == 1
+
+    def test_distinct_files_stay_separate(self):
+        from application.worker import add_file_to_structure, count_structure_files
+
+        structure: dict = {}
+        for path in ["a.md", "src/b.py", "src/deep/c.py"]:
+            add_file_to_structure(
+                structure, path, "text/plain", size_bytes=10, token_count=4,
+            )
+
+        assert count_structure_files(structure) == 3
+        assert structure["src"]["deep"]["c.py"]["token_count"] == 4
+        assert structure["a.md"]["token_count"] == 4
+
+    def test_empty_path_is_ignored(self):
+        from application.worker import add_file_to_structure
+
+        structure: dict = {}
+        add_file_to_structure(
+            structure, "", "text/plain", size_bytes=1, token_count=1,
+        )
+        assert structure == {}

tests/worker/test_ingest_checkpoint.py

@@ -89,21 +89,22 @@ class TestEmbedCheckpoint:
         docs = _make_docs(5)
         captured_chunks: list[str] = []
 
-        def _fake_add(store, doc, sid):
-            captured_chunks.append(doc.page_content)
+        def _fake_add(store, batch, sid):
+            captured_chunks.extend(d.page_content for d in batch)
 
         # The retry decorator wraps the real fn. Patch the module-level name
-        # so our loop calls the spy directly.
+        # so our loop calls the spy directly. The loop embeds a batch at a
+        # time, so the spy receives a list and flattens it back out.
         import application.parser.embedding_pipeline as ep_mod
 
-        original = ep_mod.add_text_to_store_with_retry
-        ep_mod.add_text_to_store_with_retry = _fake_add
+        original = ep_mod.add_texts_to_store_with_retry
+        ep_mod.add_texts_to_store_with_retry = _fake_add
         try:
             ep_mod.embed_and_store_documents(
                 docs, str(tmp_path), source_id, MagicMock()
             )
         finally:
-            ep_mod.add_text_to_store_with_retry = original
+            ep_mod.add_texts_to_store_with_retry = original
 
         # Faiss seeds the store with docs[0]; the loop picks up at idx=1.
         assert captured_chunks == ["chunk-1", "chunk-2", "chunk-3", "chunk-4"]
@@ -136,17 +137,17 @@ class TestEmbedCheckpoint:
 
         captured_chunks: list[str] = []
 
-        def _fake_add(store, doc, sid):
-            captured_chunks.append(doc.page_content)
+        def _fake_add(store, batch, sid):
+            captured_chunks.extend(d.page_content for d in batch)
 
-        original = ep_mod.add_text_to_store_with_retry
-        ep_mod.add_text_to_store_with_retry = _fake_add
+        original = ep_mod.add_texts_to_store_with_retry
+        ep_mod.add_texts_to_store_with_retry = _fake_add
         try:
             ep_mod.embed_and_store_documents(
                 docs, str(tmp_path), source_id, MagicMock()
             )
         finally:
-            ep_mod.add_text_to_store_with_retry = original
+            ep_mod.add_texts_to_store_with_retry = original
 
         # On resume the FAISS store is loaded from storage (no docs_init);
         # the loop iterates the un-popped docs list starting at resume_index.
@@ -187,14 +188,14 @@ class TestEmbedCheckpoint:
         docs = _make_docs(6)
         _seed_progress_row(pg_conn, source_id, total=6, last_index=2)
 
-        original = ep_mod.add_text_to_store_with_retry
-        ep_mod.add_text_to_store_with_retry = lambda store, doc, sid: None
+        original = ep_mod.add_texts_to_store_with_retry
+        ep_mod.add_texts_to_store_with_retry = lambda store, batch, sid: None
         try:
             ep_mod.embed_and_store_documents(
                 docs, str(tmp_path), source_id, MagicMock()
             )
         finally:
-            ep_mod.add_text_to_store_with_retry = original
+            ep_mod.add_texts_to_store_with_retry = original
 
         assert len(captured_kwargs) == 1
         # No ``docs_init`` on resume — this is what triggers FaissStore to
@@ -219,14 +220,14 @@ class TestEmbedCheckpoint:
         docs = _make_docs(4)
         _seed_progress_row(pg_conn, source_id, total=4, last_index=1)
 
-        original = ep_mod.add_text_to_store_with_retry
-        ep_mod.add_text_to_store_with_retry = lambda store, doc, sid: None
+        original = ep_mod.add_texts_to_store_with_retry
+        ep_mod.add_texts_to_store_with_retry = lambda store, batch, sid: None
         try:
             ep_mod.embed_and_store_documents(
                 docs, str(tmp_path), source_id, MagicMock()
             )
         finally:
-            ep_mod.add_text_to_store_with_retry = original
+            ep_mod.add_texts_to_store_with_retry = original
 
         fake_store.delete_index.assert_not_called()
 
@@ -245,9 +246,11 @@ class TestEmbedCheckpoint:
         )
 
         captured: list[str] = []
-        original = ep_mod.add_text_to_store_with_retry
-        ep_mod.add_text_to_store_with_retry = (
-            lambda store, doc, sid: captured.append(doc.page_content)
+        original = ep_mod.add_texts_to_store_with_retry
+        ep_mod.add_texts_to_store_with_retry = (
+            lambda store, batch, sid: captured.extend(
+                d.page_content for d in batch
+            )
         )
         try:
             ep_mod.embed_and_store_documents(
@@ -255,7 +258,7 @@ class TestEmbedCheckpoint:
                 attempt_id="att-A",
             )
         finally:
-            ep_mod.add_text_to_store_with_retry = original
+            ep_mod.add_texts_to_store_with_retry = original
 
         # Same attempt → resume past the last persisted index.
         assert captured == ["chunk-3", "chunk-4", "chunk-5"]
@@ -278,9 +281,11 @@ class TestEmbedCheckpoint:
         docs = _make_docs(5)
 
         captured: list[str] = []
-        original = ep_mod.add_text_to_store_with_retry
-        ep_mod.add_text_to_store_with_retry = (
-            lambda store, doc, sid: captured.append(doc.page_content)
+        original = ep_mod.add_texts_to_store_with_retry
+        ep_mod.add_texts_to_store_with_retry = (
+            lambda store, batch, sid: captured.extend(
+                d.page_content for d in batch
+            )
         )
         try:
             ep_mod.embed_and_store_documents(
@@ -288,7 +293,7 @@ class TestEmbedCheckpoint:
                 attempt_id="att-new",
             )
         finally:
-            ep_mod.add_text_to_store_with_retry = original
+            ep_mod.add_texts_to_store_with_retry = original
 
         # Fresh attempt → reset to chunk 0; FAISS branch seeds with
         # docs[0] and the loop picks up at idx=1.
@@ -327,9 +332,11 @@ class TestEmbedCheckpoint:
         docs = _make_docs(5)
 
         captured: list[str] = []
-        original = ep_mod.add_text_to_store_with_retry
-        ep_mod.add_text_to_store_with_retry = (
-            lambda store, doc, sid: captured.append(doc.page_content)
+        original = ep_mod.add_texts_to_store_with_retry
+        ep_mod.add_texts_to_store_with_retry = (
+            lambda store, batch, sid: captured.extend(
+                d.page_content for d in batch
+            )
         )
         try:
             ep_mod.embed_and_store_documents(
@@ -337,7 +344,7 @@ class TestEmbedCheckpoint:
                 attempt_id="sync-2",
             )
         finally:
-            ep_mod.add_text_to_store_with_retry = original
+            ep_mod.add_texts_to_store_with_retry = original
 
         # All non-seed chunks re-embedded under the new attempt.
         assert captured == ["chunk-1", "chunk-2", "chunk-3", "chunk-4"]
@@ -358,16 +365,18 @@ class TestEmbedCheckpoint:
         docs = _make_docs(4)
 
         captured: list[str] = []
-        original = ep_mod.add_text_to_store_with_retry
-        ep_mod.add_text_to_store_with_retry = (
-            lambda store, doc, sid: captured.append(doc.page_content)
+        original = ep_mod.add_texts_to_store_with_retry
+        ep_mod.add_texts_to_store_with_retry = (
+            lambda store, batch, sid: captured.extend(
+                d.page_content for d in batch
+            )
         )
         try:
             ep_mod.embed_and_store_documents(
                 docs, str(tmp_path), source_id, MagicMock(),
             )
         finally:
-            ep_mod.add_text_to_store_with_retry = original
+            ep_mod.add_texts_to_store_with_retry = original
 
         # Resumed past last_index=1.
         assert captured == ["chunk-2", "chunk-3"]
@@ -387,15 +396,15 @@ class TestEmbedCheckpoint:
         source_id = str(uuid.uuid4())
         docs = _make_docs(1)
 
-        original = ep_mod.add_text_to_store_with_retry
-        ep_mod.add_text_to_store_with_retry = lambda store, doc, sid: None
+        original = ep_mod.add_texts_to_store_with_retry
+        ep_mod.add_texts_to_store_with_retry = lambda store, batch, sid: None
         try:
             ep_mod.embed_and_store_documents(
                 docs, str(tmp_path), source_id, MagicMock(),
                 attempt_id="att-single",
             )
         finally:
-            ep_mod.add_text_to_store_with_retry = original
+            ep_mod.add_texts_to_store_with_retry = original
 
         row = pg_conn.execute(
             text(
@@ -426,15 +435,15 @@ class TestEmbedCheckpoint:
         source_id = str(uuid.uuid4())
         docs = _make_docs(4)
 
-        original = ep_mod.add_text_to_store_with_retry
-        ep_mod.add_text_to_store_with_retry = lambda store, doc, sid: None
+        original = ep_mod.add_texts_to_store_with_retry
+        ep_mod.add_texts_to_store_with_retry = lambda store, batch, sid: None
         try:
             ep_mod.embed_and_store_documents(
                 docs, str(tmp_path), source_id, MagicMock(),
                 attempt_id="att-multi",
             )
         finally:
-            ep_mod.add_text_to_store_with_retry = original
+            ep_mod.add_texts_to_store_with_retry = original
 
         row = pg_conn.execute(
             text(