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

Remove _create_and_chunk & _update_and_rechunk, keep storage independent of chunking

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-12-05 12:52:48 +02:00
parent 0ab43d62cd
commit 69b1afa534
No known key found for this signature in database
3 changed files with 115 additions and 274 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

88
haiku_rag_slim/haiku/rag/client.py
View file

@ -231,7 +231,7 @@ class HaikuRAG:
self.store.restore_table_versions(versions) self.store.restore_table_versions(versions)
raise raise
async def _update_document_and_rechunk( async def _update_document_with_chunks(
self, self,
document: Document, document: Document,
chunks: list[Chunk], chunks: list[Chunk],
@ -279,28 +279,6 @@ class HaikuRAG:
self.store.restore_table_versions(versions) self.store.restore_table_versions(versions)
raise raise
async def _create_document_with_docling(
self,
docling_document,
uri: str | None = None,
title: str | None = None,
metadata: dict | None = None,
chunks: list[Chunk] | None = None,
) -> Document:
"""Create a new document from DoclingDocument."""
content = docling_document.export_to_markdown()
document = Document(
content=content,
uri=uri,
title=title,
metadata=metadata or {},
docling_document_json=docling_document.model_dump_json(),
docling_version=docling_document.version,
)
return await self.document_repository._create_and_chunk(
document, docling_document, chunks
)
async def create_document( async def create_document(
self, self,
content: str, content: str,
@ -409,7 +387,7 @@ class HaikuRAG:
docling_version=docling_version, docling_version=docling_version,
) )
return await self.document_repository._create_and_chunk(document, None, chunks) return await self._store_document_with_chunks(document, chunks)
async def create_document_from_source( async def create_document_from_source(
self, source: str | Path, title: str | None = None, metadata: dict | None = None self, source: str | Path, title: str | None = None, metadata: dict | None = None
@ -542,7 +520,7 @@ class HaikuRAG:
existing_doc.docling_version = docling_document.version existing_doc.docling_version = docling_document.version
if title is not None: if title is not None:
existing_doc.title = title existing_doc.title = title
return await self._update_document_and_rechunk( return await self._update_document_with_chunks(
existing_doc, embedded_chunks existing_doc, embedded_chunks
) )
else: else:
@ -649,7 +627,7 @@ class HaikuRAG:
existing_doc.docling_version = docling_document.version existing_doc.docling_version = docling_document.version
if title is not None: if title is not None:
existing_doc.title = title existing_doc.title = title
return await self._update_document_and_rechunk( return await self._update_document_with_chunks(
existing_doc, embedded_chunks existing_doc, embedded_chunks
) )
else: else:
@ -718,18 +696,28 @@ class HaikuRAG:
return await self.document_repository.get_by_uri(uri) return await self.document_repository.get_by_uri(uri)
async def update_document(self, document: Document) -> Document: async def update_document(self, document: Document) -> Document:
"""Update an existing document.""" """Update an existing document.
# Convert content to DoclingDocument
converter = get_converter(self._config) Reconverts content, rechunks, and regenerates embeddings.
docling_document = await converter.convert_text(document.content)
Args:
document: The document to update (must have ID set).
Returns:
The updated Document instance.
"""
from haiku.rag.embeddings import embed_chunks
# Convert → Chunk → Embed using primitives
docling_document = await self.convert(document.content)
chunks = await self.chunk(docling_document)
embedded_chunks = await embed_chunks(chunks, self._config)
# Store DoclingDocument JSON # Store DoclingDocument JSON
document.docling_document_json = docling_document.model_dump_json() document.docling_document_json = docling_document.model_dump_json()
document.docling_version = docling_document.version document.docling_version = docling_document.version
return await self.document_repository._update_and_rechunk( return await self._update_document_with_chunks(document, embedded_chunks)
document, docling_document
)
async def update_document_fields( async def update_document_fields(
self, self,
@ -762,6 +750,8 @@ class HaikuRAG:
""" """
from docling_core.types.doc.document import DoclingDocument from docling_core.types.doc.document import DoclingDocument
from haiku.rag.embeddings import embed_chunks
# Validate: content and docling_document_json are mutually exclusive # Validate: content and docling_document_json are mutually exclusive
if content is not None and docling_document_json is not None: if content is not None and docling_document_json is not None:
raise ValueError( raise ValueError(
@ -800,7 +790,7 @@ class HaikuRAG:
if content is None and chunks is None and docling_document is None: if content is None and chunks is None and docling_document is None:
return await self.document_repository.update(existing_doc) return await self.document_repository.update(existing_doc)
# Custom chunks provided - use them as-is # Custom chunks provided - use them as-is (pre-embedded)
if chunks is not None: if chunks is not None:
# Update content field if provided # Update content field if provided
if content is not None: if content is not None:
@ -814,37 +804,29 @@ class HaikuRAG:
if content is None: if content is None:
existing_doc.content = docling_document.export_to_markdown() existing_doc.content = docling_document.export_to_markdown()
# Delete existing chunks and use custom ones return await self._update_document_with_chunks(existing_doc, chunks)
await self.chunk_repository.delete_by_document_id(document_id)
await self.document_repository.update(existing_doc)
for order, chunk in enumerate(chunks): # DoclingDocument provided without chunks - chunk and embed using primitives
chunk.document_id = document_id
chunk.order = order
await self.chunk_repository.create(chunks)
return existing_doc
# DoclingDocument provided without chunks - extract content and rechunk
if docling_document is not None: if docling_document is not None:
existing_doc.content = docling_document.export_to_markdown() existing_doc.content = docling_document.export_to_markdown()
existing_doc.docling_document_json = docling_document_json existing_doc.docling_document_json = docling_document_json
existing_doc.docling_version = docling_version existing_doc.docling_version = docling_version
return await self.document_repository._update_and_rechunk( new_chunks = await self.chunk(docling_document)
existing_doc, docling_document embedded_chunks = await embed_chunks(new_chunks, self._config)
return await self._update_document_with_chunks(
existing_doc, embedded_chunks
) )
# Content provided without chunks - convert and rechunk # Content provided without chunks - convert, chunk, and embed using primitives
existing_doc.content = content # type: ignore[assignment] existing_doc.content = content # type: ignore[assignment]
converter = get_converter(self._config) converted_docling = await self.convert(existing_doc.content)
converted_docling = await converter.convert_text(existing_doc.content)
existing_doc.docling_document_json = converted_docling.model_dump_json() existing_doc.docling_document_json = converted_docling.model_dump_json()
existing_doc.docling_version = converted_docling.version existing_doc.docling_version = converted_docling.version
return await self.document_repository._update_and_rechunk( new_chunks = await self.chunk(converted_docling)
existing_doc, converted_docling embedded_chunks = await embed_chunks(new_chunks, self._config)
) return await self._update_document_with_chunks(existing_doc, embedded_chunks)
async def delete_document(self, document_id: str) -> bool: async def delete_document(self, document_id: str) -> bool:
"""Delete a document by its ID.""" """Delete a document by its ID."""

87
haiku_rag_slim/haiku/rag/store/repositories/document.py
View file

@ -1,17 +1,10 @@
import asyncio
import json import json
from datetime import datetime from datetime import datetime
from typing import TYPE_CHECKING
from uuid import uuid4 from uuid import uuid4
from haiku.rag.store.engine import DocumentRecord, Store from haiku.rag.store.engine import DocumentRecord, Store
from haiku.rag.store.models.document import Document from haiku.rag.store.models.document import Document
if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.store.models.chunk import Chunk
class DocumentRepository: class DocumentRepository:
"""Repository for Document operations.""" """Repository for Document operations."""
@ -197,83 +190,3 @@ class DocumentRepository:
self.store.documents_table = self.store.db.create_table( self.store.documents_table = self.store.db.create_table(
"documents", schema=DocumentRecord "documents", schema=DocumentRecord
) )
async def _create_and_chunk(
self,
entity: Document,
docling_document: "DoclingDocument | None",
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)
# Attempt to create chunks; on failure, prefer version rollback
try:
# Create chunks if not provided
if chunks is None:
assert docling_document is not None, (
"docling_document is required when chunks are not provided"
)
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"
)
# Set document_id and order for all chunks
for order, chunk in enumerate(chunks):
chunk.document_id = created_doc.id
chunk.order = order
# Batch create all chunks in a single operation
await self.chunk_repository.create(chunks)
# Vacuum old versions in background (non-blocking)
asyncio.create_task(self.store.vacuum())
return created_doc
except Exception:
# Roll back to the captured versions and re-raise
self.store.restore_table_versions(versions)
raise
async def _update_and_rechunk(
self, entity: Document, docling_document: "DoclingDocument"
) -> Document:
"""Update a document and regenerate its 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)
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
)
# Vacuum old versions in background (non-blocking)
asyncio.create_task(self.store.vacuum())
return updated_doc
except Exception:
# Roll back to the captured versions and re-raise
self.store.restore_table_versions(versions)
raise

214
tests/test_versioning.py
View file

@ -1,100 +1,62 @@
import pytest import pytest
from haiku.rag.config import Config from haiku.rag.client import HaikuRAG
from haiku.rag.converters import get_converter
from haiku.rag.store.engine import Store 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
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_version_rollback_on_create_failure(temp_db_path): async def test_version_rollback_on_create_failure(temp_db_path):
store = Store(temp_db_path, create=True) async with HaikuRAG(db_path=temp_db_path, create=True) as client:
repo = DocumentRepository(store) # Patch chunk_repository.create to succeed then fail, triggering rollback
orig_create = client.chunk_repository.create
# Ensure chunk repository is instantiated and stub embeddings to avoid network async def succeed_then_fail(chunks):
dim = repo.chunk_repository.embedder._vector_dim await orig_create(chunks)
raise RuntimeError("boom")
async def fake_embed(x): # type: ignore[no-redef] client.chunk_repository.create = succeed_then_fail # type: ignore[method-assign]
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] # Attempt to create document; expect failure and rollback
content = "Hello, rollback!"
# Patch create_chunks_for_document to succeed then fail, triggering rollback with pytest.raises(RuntimeError):
orig = repo.chunk_repository.create_chunks_for_document await client.create_document(content=content)
async def succeed_then_fail(document_id, dl_doc): # noqa: ARG001 # State should be restored (no documents/chunks)
await orig(document_id, dl_doc) docs = await client.list_documents()
raise RuntimeError("boom") assert len(docs) == 0
all_chunks = await client.chunk_repository.list_all()
repo.chunk_repository.create_chunks_for_document = succeed_then_fail # type: ignore[assignment] assert len(all_chunks) == 0
# Attempt to create document with chunks; expect failure and rollback
content = "Hello, rollback!"
doc = Document(content=content)
converter = get_converter(Config)
dl_doc = await converter.convert_text(content, name="test.md")
with pytest.raises(RuntimeError):
await repo._create_and_chunk(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 @pytest.mark.asyncio
async def test_version_rollback_on_update_failure(temp_db_path): async def test_version_rollback_on_update_failure(temp_db_path):
store = Store(temp_db_path, create=True) async with HaikuRAG(db_path=temp_db_path, create=True) as client:
repo = DocumentRepository(store) # Create a valid document first
base_content = "Base content"
created = await client.create_document(content=base_content)
# Stub embeddings to avoid network # Patch chunk_repository.create to succeed then fail during update
dim = repo.chunk_repository.embedder._vector_dim orig_create = client.chunk_repository.create
async def fake_embed(x): # type: ignore[no-redef] async def succeed_then_fail(chunks):
if isinstance(x, list): await orig_create(chunks)
return [[0.0] * dim for _ in x] raise RuntimeError("update fail")
return [0.0] * dim
repo.chunk_repository.embedder.embed = fake_embed # type: ignore[assignment] client.chunk_repository.create = succeed_then_fail # type: ignore[method-assign]
# Create a valid document first (with real chunking and stubbed embeddings) # Attempt update
base_content = "Base content" created.content = "Updated content"
base_doc = Document(content=base_content)
converter = get_converter(Config)
base_dl = await converter.convert_text(base_content, name="base.md")
created = await repo._create_and_chunk(base_doc, base_dl)
# Force new chunk creation to fail during update after writing with pytest.raises(RuntimeError):
orig = repo.chunk_repository.create_chunks_for_document await client.update_document(created)
async def succeed_then_fail(document_id, dl_doc): # noqa: ARG001 # Content and chunks should remain the original
await orig(document_id, dl_doc) persisted = await client.get_document_by_id(created.id) # type: ignore[arg-type]
raise RuntimeError("update fail") assert persisted is not None
assert persisted.content == base_content
repo.chunk_repository.create_chunks_for_document = succeed_then_fail # type: ignore[assignment] original_chunks = await client.chunk_repository.get_by_document_id(created.id) # type: ignore[arg-type]
assert len(original_chunks) > 0
# Attempt update
updated_content = "Updated content"
created.content = updated_content
updated_dl = await converter.convert_text(updated_content, name="updated.md")
with pytest.raises(RuntimeError):
await repo._update_and_rechunk(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
def test_new_database_does_not_run_upgrades(monkeypatch, temp_db_path): def test_new_database_does_not_run_upgrades(monkeypatch, temp_db_path):
@ -130,81 +92,65 @@ def test_existing_database_runs_upgrades(monkeypatch, temp_db_path):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_vacuum_with_retention_threshold(temp_db_path): async def test_vacuum_with_retention_threshold(temp_db_path):
store = Store(temp_db_path, create=True) async with HaikuRAG(db_path=temp_db_path, create=True) as client:
repo = DocumentRepository(store) # Create first document
await client.create_document(content="First document")
# Stub embeddings to avoid network # Create second document
dim = repo.chunk_repository.embedder._vector_dim await client.create_document(content="Second document")
async def fake_embed(x): # type: ignore[no-redef] store = client.store
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] # Get initial version counts (should have multiple versions from creates)
initial_doc_versions = len(list(store.documents_table.list_versions()))
initial_chunk_versions = len(list(store.chunks_table.list_versions()))
# Create first document assert initial_doc_versions > 1, "Should have multiple document table versions"
converter = get_converter(Config) assert initial_chunk_versions > 1, "Should have multiple chunk table versions"
doc1 = Document(content="First document")
dl_doc1 = await converter.convert_text("First document", name="doc1.md")
await repo._create_and_chunk(doc1, dl_doc1)
# Create second document # Vacuum with default threshold (60 seconds) - should keep recent versions
doc2 = Document(content="Second document") # Note: vacuum may create new versions even when not cleaning up old ones
dl_doc2 = await converter.convert_text("Second document", name="doc2.md") await store.vacuum()
await repo._create_and_chunk(doc2, dl_doc2)
# Get initial version counts (should have multiple versions from creates) after_default_doc_versions = len(list(store.documents_table.list_versions()))
initial_doc_versions = len(list(store.documents_table.list_versions())) after_default_chunk_versions = len(list(store.chunks_table.list_versions()))
initial_chunk_versions = len(list(store.chunks_table.list_versions()))
assert initial_doc_versions > 1, "Should have multiple document table versions" # After vacuum with retention, version count should stay the same or increase
assert initial_chunk_versions > 1, "Should have multiple chunk table versions" # (optimize may create new versions) but not decrease
assert after_default_doc_versions >= initial_doc_versions, (
"Default vacuum should not remove recent versions"
)
assert after_default_chunk_versions >= initial_chunk_versions, (
"Default vacuum should not remove recent versions"
)
# Vacuum with default threshold (60 seconds) - should keep recent versions # Vacuum with 0 threshold - should significantly reduce versions
# Note: vacuum may create new versions even when not cleaning up old ones await store.vacuum(retention_seconds=0)
await store.vacuum()
after_default_doc_versions = len(list(store.documents_table.list_versions())) after_zero_doc_versions = len(list(store.documents_table.list_versions()))
after_default_chunk_versions = len(list(store.chunks_table.list_versions())) after_zero_chunk_versions = len(list(store.chunks_table.list_versions()))
# After vacuum with retention, version count should stay the same or increase # After aggressive vacuum, should have minimal versions (1-2)
# (optimize may create new versions) but not decrease # Note: optimize operation may create a version after cleanup
assert after_default_doc_versions >= initial_doc_versions, ( assert after_zero_doc_versions <= 2, (
"Default vacuum should not remove recent versions" f"Should have minimal document versions after vacuum(0), got {after_zero_doc_versions}"
) )
assert after_default_chunk_versions >= initial_chunk_versions, ( assert after_zero_chunk_versions <= 2, (
"Default vacuum should not remove recent versions" f"Should have minimal chunk versions after vacuum(0), got {after_zero_chunk_versions}"
) )
# Vacuum with 0 threshold - should significantly reduce versions # And it should be significantly fewer than before
await store.vacuum(retention_seconds=0) assert after_zero_doc_versions < initial_doc_versions, (
"Should have fewer versions after vacuum(0)"
after_zero_doc_versions = len(list(store.documents_table.list_versions())) )
after_zero_chunk_versions = len(list(store.chunks_table.list_versions())) assert after_zero_chunk_versions < initial_chunk_versions, (
"Should have fewer versions after vacuum(0)"
# After aggressive vacuum, should have minimal versions (1-2) )
# Note: optimize operation may create a version after cleanup
assert after_zero_doc_versions <= 2, (
f"Should have minimal document versions after vacuum(0), got {after_zero_doc_versions}"
)
assert after_zero_chunk_versions <= 2, (
f"Should have minimal chunk versions after vacuum(0), got {after_zero_chunk_versions}"
)
# And it should be significantly fewer than before
assert after_zero_doc_versions < initial_doc_versions, (
"Should have fewer versions after vacuum(0)"
)
assert after_zero_chunk_versions < initial_chunk_versions, (
"Should have fewer versions after vacuum(0)"
)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_vacuum_completes_before_context_exit(temp_db_path, monkeypatch): async def test_vacuum_completes_before_context_exit(temp_db_path, monkeypatch):
"""Test that background vacuum completes when context manager exits.""" """Test that background vacuum completes when context manager exits."""
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config from haiku.rag.config import Config
# Set aggressive vacuum retention for this test # Set aggressive vacuum retention for this test