danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #57 from ggozad/feat/rollback-on-fail

Browse source 
Rollback version if a failure occurs when creating/updating documents
This commit is contained in:
Yiorgis Gozadinos 2025-09-15 11:55:33 +03:00 committed by GitHub
commit 119f72cd1a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 168 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
docs/cli.md
View file

@ -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

8
docs/python.md
View file

@ -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:

15
src/haiku/rag/store/engine.py
View file

@ -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."""

74
src/haiku/rag/store/repositories/document.py
View file

@ -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

94
tests/test_versioning.py Normal file
View file

@ -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