diff --git a/docs/cli.md b/docs/cli.md index 0d20fa95..5785580c 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -36,8 +36,10 @@ haiku-rag add-src https://example.com/article.html ``` !!! note - As you add documents to `haiku.rag` the database keeps growing. By default, `lanceDB` supports versioning - of your data. You can optimize and compact the database by running the [vaccum](#vacuum-optimize-and-cleanup) command. + As you add documents to `haiku.rag` the database keeps growing. By default, LanceDB supports versioning + of your data. Create/update operations are atomic‑feeling: if anything fails during chunking or embedding, + the database rolls back to the pre‑operation snapshot using LanceDB table versioning. You can optimize and + compact the database by running the [vacuum](#vacuum-optimize-and-cleanup) command. ### Get Document diff --git a/docs/python.md b/docs/python.md index 547de767..925c2d52 100644 --- a/docs/python.md +++ b/docs/python.md @@ -109,6 +109,14 @@ await client.vacuum() This compacts tables and removes historical versions to keep disk usage in check. It’s safe to run anytime, for example after bulk imports or periodically in long‑running apps. +### Atomic Writes and Rollback + +Document create and update operations take a snapshot of table versions before any write and automatically roll back to that snapshot if something fails (for example, during chunking or embedding). This restores both the `documents` and `chunks` tables to their pre‑operation state using LanceDB’s table versioning. + +- Applies to: `create_document(...)`, `create_document_from_source(...)`, `update_document(...)`, and internal rebuild/update flows. +- Scope: Both document rows and all associated chunks are rolled back together. +- Vacuum: Running `vacuum()` later prunes old versions for disk efficiency; rollbacks occur immediately during the failing operation and are not impacted. + ## Searching Documents The search method performs native hybrid search (vector + full-text) using LanceDB with optional reranking for improved relevance: diff --git a/src/haiku/rag/store/engine.py b/src/haiku/rag/store/engine.py index 036729b2..737d90e7 100644 --- a/src/haiku/rag/store/engine.py +++ b/src/haiku/rag/store/engine.py @@ -209,6 +209,21 @@ class Store: # LanceDB connections are automatically managed pass + def current_table_versions(self) -> dict[str, int]: + """Capture current versions of key tables for rollback using LanceDB's API.""" + return { + "documents": int(self.documents_table.version), + "chunks": int(self.chunks_table.version), + "settings": int(self.settings_table.version), + } + + def restore_table_versions(self, versions: dict[str, int]) -> bool: + """Restore tables to the provided versions using LanceDB's API.""" + self.documents_table.restore(int(versions["documents"])) + self.chunks_table.restore(int(versions["chunks"])) + self.settings_table.restore(int(versions["settings"])) + return True + @property def _connection(self): """Compatibility property for repositories expecting _connection.""" diff --git a/src/haiku/rag/store/repositories/document.py b/src/haiku/rag/store/repositories/document.py index 1a218bb9..02d0be6c 100644 --- a/src/haiku/rag/store/repositories/document.py +++ b/src/haiku/rag/store/repositories/document.py @@ -171,44 +171,64 @@ class DocumentRepository: chunks: list["Chunk"] | None = None, ) -> Document: """Create a document with its chunks and embeddings.""" + # Snapshot table versions for versioned rollback (if supported) + versions = self.store.current_table_versions() + # Create the document created_doc = await self.create(entity) - # Create chunks if not provided - if chunks is None: - assert created_doc.id is not None, ( - "Document ID should not be None after creation" - ) - await self.chunk_repository.create_chunks_for_document( - created_doc.id, docling_document - ) - else: - # Use provided chunks, set order from list position - assert created_doc.id is not None, ( - "Document ID should not be None after creation" - ) - for order, chunk in enumerate(chunks): - chunk.document_id = created_doc.id - chunk.metadata["order"] = order - await self.chunk_repository.create(chunk) + # Attempt to create chunks; on failure, prefer version rollback + try: + # Create chunks if not provided + if chunks is None: + assert created_doc.id is not None, ( + "Document ID should not be None after creation" + ) + await self.chunk_repository.create_chunks_for_document( + created_doc.id, docling_document + ) + else: + # Use provided chunks, set order from list position + assert created_doc.id is not None, ( + "Document ID should not be None after creation" + ) + for order, chunk in enumerate(chunks): + chunk.document_id = created_doc.id + chunk.metadata["order"] = order + await self.chunk_repository.create(chunk) - return created_doc + return created_doc + except Exception: + # Roll back to the captured versions and re-raise + self.store.restore_table_versions(versions) + raise async def _update_with_docling( self, entity: Document, docling_document: DoclingDocument ) -> Document: """Update a document and regenerate its chunks.""" - # Delete existing chunks assert entity.id is not None, "Document ID is required for update" + + # Snapshot table versions for versioned rollback + versions = self.store.current_table_versions() + + # Delete existing chunks before writing new ones await self.chunk_repository.delete_by_document_id(entity.id) - # Update the document - updated_doc = await self.update(entity) + try: + # Update the document + updated_doc = await self.update(entity) - # Create new chunks - assert updated_doc.id is not None, "Document ID should not be None after update" - await self.chunk_repository.create_chunks_for_document( - updated_doc.id, docling_document - ) + # Create new chunks + assert updated_doc.id is not None, ( + "Document ID should not be None after update" + ) + await self.chunk_repository.create_chunks_for_document( + updated_doc.id, docling_document + ) - return updated_doc + return updated_doc + except Exception: + # Roll back to the captured versions and re-raise + self.store.restore_table_versions(versions) + raise diff --git a/tests/test_versioning.py b/tests/test_versioning.py new file mode 100644 index 00000000..fc1faa2b --- /dev/null +++ b/tests/test_versioning.py @@ -0,0 +1,94 @@ +import pytest + +from haiku.rag.store.engine import Store +from haiku.rag.store.models.document import Document +from haiku.rag.store.repositories.chunk import ChunkRepository +from haiku.rag.store.repositories.document import DocumentRepository +from haiku.rag.utils import text_to_docling_document + + +@pytest.mark.asyncio +async def test_version_rollback_on_create_failure(temp_db_path): + store = Store(temp_db_path) + repo = DocumentRepository(store) + + # Ensure chunk repository is instantiated and stub embeddings to avoid network + dim = repo.chunk_repository.embedder._vector_dim + + async def fake_embed(x): # type: ignore[no-redef] + if isinstance(x, list): + return [[0.0] * dim for _ in x] + return [0.0] * dim + + repo.chunk_repository.embedder.embed = fake_embed # type: ignore[assignment] + + # Patch create_chunks_for_document to succeed then fail, triggering rollback + orig = repo.chunk_repository.create_chunks_for_document + + async def succeed_then_fail(document_id, dl_doc): # noqa: ARG001 + await orig(document_id, dl_doc) + raise RuntimeError("boom") + + repo.chunk_repository.create_chunks_for_document = succeed_then_fail # type: ignore[assignment] + + # Attempt to create document with chunks; expect failure and rollback + content = "Hello, rollback!" + doc = Document(content=content) + dl_doc = text_to_docling_document(content, name="test.md") + + with pytest.raises(RuntimeError): + await repo._create_with_docling(doc, dl_doc) + + # State should be restored (no documents/chunks) + docs = await repo.list_all() + assert len(docs) == 0 + chunks_repo = ChunkRepository(store) + all_chunks = await chunks_repo.list_all() + assert len(all_chunks) == 0 + + +@pytest.mark.asyncio +async def test_version_rollback_on_update_failure(temp_db_path): + store = Store(temp_db_path) + repo = DocumentRepository(store) + + # Stub embeddings to avoid network + dim = repo.chunk_repository.embedder._vector_dim + + async def fake_embed(x): # type: ignore[no-redef] + if isinstance(x, list): + return [[0.0] * dim for _ in x] + return [0.0] * dim + + repo.chunk_repository.embedder.embed = fake_embed # type: ignore[assignment] + + # Create a valid document first (with real chunking and stubbed embeddings) + base_content = "Base content" + base_doc = Document(content=base_content) + base_dl = text_to_docling_document(base_content, name="base.md") + created = await repo._create_with_docling(base_doc, base_dl) + + # Force new chunk creation to fail during update after writing + orig = repo.chunk_repository.create_chunks_for_document + + async def succeed_then_fail(document_id, dl_doc): # noqa: ARG001 + await orig(document_id, dl_doc) + raise RuntimeError("update fail") + + repo.chunk_repository.create_chunks_for_document = succeed_then_fail # type: ignore[assignment] + + # Attempt update + updated_content = "Updated content" + created.content = updated_content + updated_dl = text_to_docling_document(updated_content, name="updated.md") + + with pytest.raises(RuntimeError): + await repo._update_with_docling(created, updated_dl) + + # Content and chunks should remain the original + persisted = await repo.get_by_id(created.id) # type: ignore[arg-type] + assert persisted is not None + assert persisted.content == base_content + chunks_repo = ChunkRepository(store) + original_chunks = await chunks_repo.get_by_document_id(created.id) # type: ignore[arg-type] + assert len(original_chunks) > 0