Merge pull request 'Batch Paperless chunk embeddings during indexing' (#73) from perf/batch-paperless-embeddings into main

blueprints/rag/logic.py

@@ -67,6 +67,9 @@ def _get_engine():
     return _get_engine._engine
 
 
+# Chunks per embedding request when indexing Paperless documents.
+EMBED_BATCH_SIZE = 32
+
 text_splitter = RecursiveCharacterTextSplitter(
     chunk_size=1000,  # chunk size (characters)
     chunk_overlap=200,  # chunk overlap (characters)
@@ -200,15 +203,25 @@ async def index_documents():
     splits = _sanitize_documents(splits)
     logger.info(f"Indexing {len(splits)} chunks from {len(documents)} documents")
     vector_store = _get_vector_store()
-    for i, split in enumerate(splits):
+    for start in range(0, len(splits), EMBED_BATCH_SIZE):
+        batch = splits[start : start + EMBED_BATCH_SIZE]
         try:
-            await vector_store.aadd_documents(documents=[split])
-        except Exception as e:
-            logger.error(
-                f"Failed to embed chunk {i} from {split.metadata.get('filename', 'unknown')}: {e}"
+            await vector_store.aadd_documents(documents=batch)
+        except Exception as batch_error:
+            # Embedding fails before anything is inserted; retry the batch one
+            # chunk at a time to name the chunk the embedding server rejects.
+            logger.warning(
+                f"Batch of {len(batch)} chunks failed ({batch_error}); retrying one at a time"
             )
-            logger.debug(f"Chunk content preview: {split.page_content[:200]!r}")
-            raise
+            for i, split in enumerate(batch, start):
+                try:
+                    await vector_store.aadd_documents(documents=[split])
+                except Exception as e:
+                    logger.error(
+                        f"Failed to embed chunk {i} from {split.metadata.get('filename', 'unknown')}: {e}"
+                    )
+                    logger.debug(f"Chunk content preview: {split.page_content[:200]!r}")
+                    raise
 
 
 async def fetch_obsidian_documents() -> list[Document]:

tests/unit/test_paperless_indexing.py

@@ -0,0 +1,56 @@
+from unittest.mock import AsyncMock
+
+import pytest
+from langchain_core.documents import Document
+
+from blueprints.rag import logic
+
+
+class FakeVectorStore:
+    def __init__(self, bad_text=None, fail_batches=False):
+        self.calls = []
+        self.bad_text = bad_text
+        self.fail_batches = fail_batches
+
+    async def aadd_documents(self, documents):
+        self.calls.append([d.page_content for d in documents])
+        if self.fail_batches and len(documents) > 1:
+            raise RuntimeError("batch too large")
+        if any(d.page_content == self.bad_text for d in documents):
+            raise RuntimeError("tokenizer error")
+
+
+def _setup(monkeypatch, store, count):
+    docs = [
+        Document(page_content=f"chunk {i}", metadata={"filename": f"doc{i}.pdf"})
+        for i in range(count)
+    ]
+    monkeypatch.setattr(
+        logic, "fetch_documents_from_paperless_ngx", AsyncMock(return_value=docs)
+    )
+    monkeypatch.setattr(logic.text_splitter, "split_documents", lambda d: d)
+    monkeypatch.setattr(logic, "_get_vector_store", lambda: store)
+    monkeypatch.setattr(logic, "EMBED_BATCH_SIZE", 4)
+
+
+async def test_chunks_are_embedded_in_batches(monkeypatch):
+    store = FakeVectorStore()
+    _setup(monkeypatch, store, 10)
+    await logic.index_documents()
+    assert [len(c) for c in store.calls] == [4, 4, 2]
+
+
+async def test_failed_batch_names_the_rejected_chunk(monkeypatch, caplog):
+    store = FakeVectorStore(bad_text="chunk 6")
+    _setup(monkeypatch, store, 10)
+    with pytest.raises(RuntimeError, match="tokenizer error"):
+        await logic.index_documents()
+    assert "Failed to embed chunk 6 from doc6.pdf" in caplog.text
+
+
+async def test_batch_failure_recovers_when_every_chunk_succeeds(monkeypatch):
+    store = FakeVectorStore(fail_batches=True)
+    _setup(monkeypatch, store, 5)
+    await logic.index_documents()
+    singles = [c[0] for c in store.calls if len(c) == 1]
+    assert singles == [f"chunk {i}" for i in range(5)]