Batch Paperless chunk embeddings during indexing #73
+20
-7
@@ -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]:
|
||||||
|
|||||||
@@ -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