ryan
c80125f7fe
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>
57 lines
1.9 KiB
Python
57 lines
1.9 KiB
Python
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)]
|