Batch Paperless chunk embeddings during indexing #73

Merged
ryan merged 1 commits from perf/batch-paperless-embeddings into main 2026-10-12 23:06:29 -04:00
2 changed files with 76 additions and 7 deletions
+20 -7
View File
@@ -67,6 +67,9 @@ def _get_engine():
return _get_engine._engine return _get_engine._engine
# Chunks per embedding request when indexing Paperless documents.
EMBED_BATCH_SIZE = 32
text_splitter = RecursiveCharacterTextSplitter( text_splitter = RecursiveCharacterTextSplitter(
chunk_size=1000, # chunk size (characters) chunk_size=1000, # chunk size (characters)
chunk_overlap=200, # chunk overlap (characters) chunk_overlap=200, # chunk overlap (characters)
@@ -200,15 +203,25 @@ async def index_documents():
splits = _sanitize_documents(splits) splits = _sanitize_documents(splits)
logger.info(f"Indexing {len(splits)} chunks from {len(documents)} documents") logger.info(f"Indexing {len(splits)} chunks from {len(documents)} documents")
vector_store = _get_vector_store() 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: try:
await vector_store.aadd_documents(documents=[split]) await vector_store.aadd_documents(documents=batch)
except Exception as e: except Exception as batch_error:
logger.error( # Embedding fails before anything is inserted; retry the batch one
f"Failed to embed chunk {i} from {split.metadata.get('filename', 'unknown')}: {e}" # 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}") for i, split in enumerate(batch, start):
raise 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]: async def fetch_obsidian_documents() -> list[Document]:
+56
View File
@@ -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)]