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)
raise
async def _update_document_and_rechunk(
async def _update_document_with_chunks(
self,
document: Document,
chunks: list[Chunk],
@ -279,28 +279,6 @@ class HaikuRAG:
self.store.restore_table_versions(versions)
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(
self,
content: str,
@ -409,7 +387,7 @@ class HaikuRAG:
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(
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
if title is not None:
existing_doc.title = title
return await self._update_document_and_rechunk(
return await self._update_document_with_chunks(
existing_doc, embedded_chunks
)
else:
@ -649,7 +627,7 @@ class HaikuRAG:
existing_doc.docling_version = docling_document.version
if title is not None:
existing_doc.title = title
return await self._update_document_and_rechunk(
return await self._update_document_with_chunks(
existing_doc, embedded_chunks
)
else:
@ -718,18 +696,28 @@ class HaikuRAG:
return await self.document_repository.get_by_uri(uri)
async def update_document(self, document: Document) -> Document:
"""Update an existing document."""
# Convert content to DoclingDocument
converter = get_converter(self._config)
docling_document = await converter.convert_text(document.content)
"""Update an existing document.
Reconverts content, rechunks, and regenerates embeddings.
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
document.docling_document_json = docling_document.model_dump_json()
document.docling_version = docling_document.version
return await self.document_repository._update_and_rechunk(
document, docling_document
)
return await self._update_document_with_chunks(document, embedded_chunks)
async def update_document_fields(
self,
@ -762,6 +750,8 @@ class HaikuRAG:
"""
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.embeddings import embed_chunks
# Validate: content and docling_document_json are mutually exclusive
if content is not None and docling_document_json is not None:
raise ValueError(
@ -800,7 +790,7 @@ class HaikuRAG:
if content is None and chunks is None and docling_document is None:
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:
# Update content field if provided
if content is not None:
@ -814,37 +804,29 @@ class HaikuRAG:
if content is None:
existing_doc.content = docling_document.export_to_markdown()
# Delete existing chunks and use custom ones
await self.chunk_repository.delete_by_document_id(document_id)
await self.document_repository.update(existing_doc)
return await self._update_document_with_chunks(existing_doc, chunks)
for order, chunk in enumerate(chunks):
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
# DoclingDocument provided without chunks - chunk and embed using primitives
if docling_document is not None:
existing_doc.content = docling_document.export_to_markdown()
existing_doc.docling_document_json = docling_document_json
existing_doc.docling_version = docling_version
return await self.document_repository._update_and_rechunk(
existing_doc, docling_document
new_chunks = await self.chunk(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]
converter = get_converter(self._config)
converted_docling = await converter.convert_text(existing_doc.content)
converted_docling = await self.convert(existing_doc.content)
existing_doc.docling_document_json = converted_docling.model_dump_json()
existing_doc.docling_version = converted_docling.version
return await self.document_repository._update_and_rechunk(
existing_doc, converted_docling
)
new_chunks = await self.chunk(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:
"""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
from datetime import datetime
from typing import TYPE_CHECKING
from uuid import uuid4
from haiku.rag.store.engine import DocumentRecord, Store
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:
"""Repository for Document operations."""
@ -197,83 +190,3 @@ class DocumentRepository:
self.store.documents_table = self.store.db.create_table(
"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
from haiku.rag.config import Config
from haiku.rag.converters import get_converter
from haiku.rag.client import HaikuRAG
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
async def test_version_rollback_on_create_failure(temp_db_path):
store = Store(temp_db_path, create=True)
repo = DocumentRepository(store)
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
# 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
dim = repo.chunk_repository.embedder._vector_dim
async def succeed_then_fail(chunks):
await orig_create(chunks)
raise RuntimeError("boom")
async def fake_embed(x): # type: ignore[no-redef]
if isinstance(x, list):
return [[0.0] * dim for _ in x]
return [0.0] * dim
client.chunk_repository.create = succeed_then_fail # type: ignore[method-assign]
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
orig = repo.chunk_repository.create_chunks_for_document
with pytest.raises(RuntimeError):
await client.create_document(content=content)
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)
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
# State should be restored (no documents/chunks)
docs = await client.list_documents()
assert len(docs) == 0
all_chunks = await client.chunk_repository.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, create=True)
repo = DocumentRepository(store)
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
# Create a valid document first
base_content = "Base content"
created = await client.create_document(content=base_content)
# Stub embeddings to avoid network
dim = repo.chunk_repository.embedder._vector_dim
# Patch chunk_repository.create to succeed then fail during update
orig_create = client.chunk_repository.create
async def fake_embed(x): # type: ignore[no-redef]
if isinstance(x, list):
return [[0.0] * dim for _ in x]
return [0.0] * dim
async def succeed_then_fail(chunks):
await orig_create(chunks)
raise RuntimeError("update fail")
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)
base_content = "Base 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)
# Attempt update
created.content = "Updated content"
# Force new chunk creation to fail during update after writing
orig = repo.chunk_repository.create_chunks_for_document
with pytest.raises(RuntimeError):
await client.update_document(created)
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 = 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
# Content and chunks should remain the original
persisted = await client.get_document_by_id(created.id) # type: ignore[arg-type]
assert persisted is not None
assert persisted.content == base_content
original_chunks = await client.chunk_repository.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):
@ -130,81 +92,65 @@ def test_existing_database_runs_upgrades(monkeypatch, temp_db_path):
@pytest.mark.asyncio
async def test_vacuum_with_retention_threshold(temp_db_path):
store = Store(temp_db_path, create=True)
repo = DocumentRepository(store)
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
# Create first document
await client.create_document(content="First document")
# Stub embeddings to avoid network
dim = repo.chunk_repository.embedder._vector_dim
# Create second document
await client.create_document(content="Second document")
async def fake_embed(x): # type: ignore[no-redef]
if isinstance(x, list):
return [[0.0] * dim for _ in x]
return [0.0] * dim
store = client.store
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
converter = get_converter(Config)
doc1 = Document(content="First document")
dl_doc1 = await converter.convert_text("First document", name="doc1.md")
await repo._create_and_chunk(doc1, dl_doc1)
assert initial_doc_versions > 1, "Should have multiple document table versions"
assert initial_chunk_versions > 1, "Should have multiple chunk table versions"
# Create second document
doc2 = Document(content="Second document")
dl_doc2 = await converter.convert_text("Second document", name="doc2.md")
await repo._create_and_chunk(doc2, dl_doc2)
# Vacuum with default threshold (60 seconds) - should keep recent versions
# Note: vacuum may create new versions even when not cleaning up old ones
await store.vacuum()
# 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()))
after_default_doc_versions = len(list(store.documents_table.list_versions()))
after_default_chunk_versions = len(list(store.chunks_table.list_versions()))
assert initial_doc_versions > 1, "Should have multiple document table versions"
assert initial_chunk_versions > 1, "Should have multiple chunk table versions"
# After vacuum with retention, version count should stay the same or increase
# (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
# Note: vacuum may create new versions even when not cleaning up old ones
await store.vacuum()
# Vacuum with 0 threshold - should significantly reduce versions
await store.vacuum(retention_seconds=0)
after_default_doc_versions = len(list(store.documents_table.list_versions()))
after_default_chunk_versions = len(list(store.chunks_table.list_versions()))
after_zero_doc_versions = len(list(store.documents_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
# (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"
)
# 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}"
)
# Vacuum with 0 threshold - should significantly reduce versions
await store.vacuum(retention_seconds=0)
after_zero_doc_versions = len(list(store.documents_table.list_versions()))
after_zero_chunk_versions = len(list(store.chunks_table.list_versions()))
# 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)"
)
# 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
async def test_vacuum_completes_before_context_exit(temp_db_path, monkeypatch):
"""Test that background vacuum completes when context manager exits."""
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
# Set aggressive vacuum retention for this test