Batch Paperless chunk embeddings during indexing
index_documents sent one embedding request and one insert per chunk. Send chunks in batches of 32 and, only when a batch fails, retry it one chunk at a time so the rejected chunk is still named in the logs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
@@ -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)]
|
||||
Reference in New Issue
Block a user