Batch-capable document and document-item repositories

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:
+    def _to_record(self, entity: Document, doc_id: str, now: str) -> DocumentRecord:
-        """Create a document in the database."""
+        return DocumentRecord(
-        self.store._assert_writable()
-        # Generate new UUID
-        doc_id = str(uuid4())
-
-        # Create timestamp
-        now = datetime.now().isoformat()
-
-        # Create document record
-        doc_record = DocumentRecord(
             id=doc_id,
             content=entity.content,
             uri=entity.uri,
@@ -84,13 +76,46 @@ class DocumentRepository:
             updated_at=now,
         )
 
-        # Add to table
+    @overload
-        await self.store.documents_table.add([doc_record])
+    async def create(self, entity: Document) -> Document: ...
 
-        entity.id = doc_id
+    @overload
-        entity.created_at = datetime.fromisoformat(now)
+    async def create(self, entity: list[Document]) -> list[Document]: ...
-        entity.updated_at = datetime.fromisoformat(now)
+
-        return entity
+    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."""

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 = [
+        records = [self._to_record(document_id, item) for item in items]
-            DocumentItemRecord(
+        await self.store.document_items_table.add(records)
-                document_id=document_id,
+
-                position=item.position,
+    async def create_all(self, items: list[DocumentItem]) -> None:
-                self_ref=item.self_ref,
+        """Bulk insert items spanning any number of documents in a single
-                label=item.label,
+        table version, keyed by each item's own ``document_id``."""
-                text=item.text,
+        if not items:
-                page_numbers=json.dumps(item.page_numbers),
+            return
-                picture_data=item.picture_data,
+
-                heading_level=item.heading_level,
+        self.store._assert_writable()
-                tree_depth=item.tree_depth,
+        records = [self._to_record(item.document_id, item) for item in items]
-            )
-            for item in items
-        ]
         await self.store.document_items_table.add(records)
 
     async def get_all_items(self, document_id: str) -> list[DocumentItem]:

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 = {