From b43843f8629c705e55038b49c3ca89f0a44cce6c Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 9 Jun 2026 09:42:47 +0300 Subject: [PATCH] Batch-capable document and document-item repositories --- .../haiku/rag/store/repositories/document.py | 59 +++++++++---- .../rag/store/repositories/document_item.py | 38 +++++--- tests/test_document.py | 86 +++++++++++++++++++ 3 files changed, 152 insertions(+), 31 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index d2d0b234..fb06435b 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py @@ -1,5 +1,6 @@ import json from datetime import datetime +from typing import overload from uuid import uuid4 from lancedb.index import BTree @@ -61,17 +62,8 @@ class DocumentRepository: else datetime.now(), ) - async def create(self, entity: Document) -> Document: - """Create a document in the database.""" - self.store._assert_writable() - # Generate new UUID - doc_id = str(uuid4()) - - # Create timestamp - now = datetime.now().isoformat() - - # Create document record - doc_record = DocumentRecord( + def _to_record(self, entity: Document, doc_id: str, now: str) -> DocumentRecord: + return DocumentRecord( id=doc_id, content=entity.content, uri=entity.uri, @@ -84,13 +76,46 @@ class DocumentRepository: updated_at=now, ) - # Add to table - await self.store.documents_table.add([doc_record]) + @overload + async def create(self, entity: Document) -> Document: ... - entity.id = doc_id - entity.created_at = datetime.fromisoformat(now) - entity.updated_at = datetime.fromisoformat(now) - return entity + @overload + async def create(self, entity: list[Document]) -> list[Document]: ... + + async def create( + self, entity: Document | list[Document] + ) -> Document | list[Document]: + """Create one or more documents in the database. + + A list is written in a single table version regardless of length. + """ + self.store._assert_writable() + + if isinstance(entity, Document): + doc_id = str(uuid4()) + now = datetime.now().isoformat() + await self.store.documents_table.add([self._to_record(entity, doc_id, now)]) + entity.id = doc_id + entity.created_at = datetime.fromisoformat(now) + entity.updated_at = datetime.fromisoformat(now) + return entity + + documents = entity + if not documents: + return [] + + now = datetime.now().isoformat() + created_at = datetime.fromisoformat(now) + records = [] + for document in documents: + doc_id = str(uuid4()) + records.append(self._to_record(document, doc_id, now)) + document.id = doc_id + document.created_at = created_at + document.updated_at = created_at + + await self.store.documents_table.add(records) + return documents async def get_by_id(self, entity_id: str) -> Document | None: """Get a document by its ID.""" diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py index dbdbac2b..4e1f5d49 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py @@ -37,26 +37,36 @@ class DocumentItemRepository: tree_depth=row.get("tree_depth", 0) or 0, ) + def _to_record(self, document_id: str, item: DocumentItem) -> DocumentItemRecord: + return DocumentItemRecord( + document_id=document_id, + position=item.position, + self_ref=item.self_ref, + label=item.label, + text=item.text, + page_numbers=json.dumps(item.page_numbers), + picture_data=item.picture_data, + heading_level=item.heading_level, + tree_depth=item.tree_depth, + ) + async def create_items(self, document_id: str, items: list[DocumentItem]) -> None: """Bulk insert items for a document.""" if not items: return self.store._assert_writable() - records = [ - DocumentItemRecord( - document_id=document_id, - position=item.position, - self_ref=item.self_ref, - label=item.label, - text=item.text, - page_numbers=json.dumps(item.page_numbers), - picture_data=item.picture_data, - heading_level=item.heading_level, - tree_depth=item.tree_depth, - ) - for item in items - ] + records = [self._to_record(document_id, item) for item in items] + await self.store.document_items_table.add(records) + + async def create_all(self, items: list[DocumentItem]) -> None: + """Bulk insert items spanning any number of documents in a single + table version, keyed by each item's own ``document_id``.""" + if not items: + return + + self.store._assert_writable() + records = [self._to_record(item.document_id, item) for item in items] await self.store.document_items_table.add(records) async def get_all_items(self, document_id: str) -> list[DocumentItem]: diff --git a/tests/test_document.py b/tests/test_document.py index e8d038cd..eb41a9a4 100644 --- a/tests/test_document.py +++ b/tests/test_document.py @@ -2,7 +2,9 @@ import pytest from haiku.rag.store.engine import Store from haiku.rag.store.models.document import Document +from haiku.rag.store.models.document_item import DocumentItem from haiku.rag.store.repositories.document import DocumentRepository +from haiku.rag.store.repositories.document_item import DocumentItemRepository @pytest.mark.asyncio @@ -93,6 +95,90 @@ async def test_document_list_with_filter(qa_corpus: list[dict[str, str]], temp_d } +@pytest.mark.asyncio +async def test_document_create_batch(qa_corpus: list[dict[str, str]], temp_db_path): + """create accepts a list of documents and writes them in a single version.""" + async with Store(temp_db_path, create=True) as store: + doc_repo = DocumentRepository(store) + + content = qa_corpus[0]["document_extracted"] + doc_a = Document(content=content, uri="https://example.com/a.txt", title="A") + doc_b = Document(content=content, uri="https://example.com/b.txt", title="B") + + before = await store.documents_table.version() + created = await doc_repo.create([doc_a, doc_b]) + after = await store.documents_table.version() + + assert isinstance(created, list) + assert len(created) == 2 + assert created[0].id is not None + assert created[1].id is not None + assert created[0].id != created[1].id + assert after - before == 1 + + round_a = await doc_repo.get_by_id(created[0].id) + round_b = await doc_repo.get_by_id(created[1].id) + assert round_a is not None and round_a.title == "A" + assert round_b is not None and round_b.title == "B" + + +@pytest.mark.asyncio +async def test_document_create_empty_batch(temp_db_path): + """create([]) is a no-op returning an empty list with no version bump.""" + async with Store(temp_db_path, create=True) as store: + doc_repo = DocumentRepository(store) + + before = await store.documents_table.version() + created = await doc_repo.create([]) + after = await store.documents_table.version() + + assert created == [] + assert after == before + + +@pytest.mark.asyncio +async def test_document_item_create_all(temp_db_path): + """create_all writes items spanning multiple documents in a single version.""" + async with Store(temp_db_path, create=True) as store: + item_repo = DocumentItemRepository(store) + + items = [ + DocumentItem( + document_id="doc-1", position=0, self_ref="#/texts/0", text="a" + ), + DocumentItem( + document_id="doc-1", position=1, self_ref="#/texts/1", text="b" + ), + DocumentItem( + document_id="doc-2", position=0, self_ref="#/texts/0", text="c" + ), + ] + + before = await store.document_items_table.version() + await item_repo.create_all(items) + after = await store.document_items_table.version() + + assert after - before == 1 + + doc1_items = await item_repo.get_all_items("doc-1") + doc2_items = await item_repo.get_all_items("doc-2") + assert [i.text for i in doc1_items] == ["a", "b"] + assert [i.text for i in doc2_items] == ["c"] + + +@pytest.mark.asyncio +async def test_document_item_create_all_empty(temp_db_path): + """create_all([]) is a no-op with no version bump.""" + async with Store(temp_db_path, create=True) as store: + item_repo = DocumentItemRepository(store) + + before = await store.document_items_table.version() + await item_repo.create_all([]) + after = await store.document_items_table.version() + + assert after == before + + def test_document_get_docling_document(): """Test parsing stored DoclingDocument JSON.""" doc_json = {