From b43843f8629c705e55038b49c3ca89f0a44cce6c Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Tue, 9 Jun 2026 09:42:47 +0300
Subject: [PATCH 1/2] 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 = {

From a7ed405d75b670a5857b629d454e15b8ed8eb954 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Tue, 9 Jun 2026 09:58:35 +0300
Subject: [PATCH 2/2] Add HaikuRAG.import_documents for batch document import

---
 CHANGELOG.md                                 |  4 +
 docs/python.md                               | 25 ++++++
 haiku_rag_slim/haiku/rag/client/__init__.py  |  9 ++
 haiku_rag_slim/haiku/rag/client/documents.py | 93 ++++++++++++++++++++
 tests/test_client.py                         | 85 ++++++++++++++++++
 5 files changed, 216 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9215ea97..0236c0df 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,10 @@
 # Changelog
 ## [Unreleased]
 
+### Added
+
+- `HaikuRAG.import_documents(imports)` batch-imports prepared documents (`DocumentImport`), writing the `documents`, `chunks`, and `document_items` tables once each regardless of batch size. `DocumentRepository.create` accepts `Document | list[Document]`.
+
 ### Fixed
 
 - Ingester worker circuit breaker is now per-source: a streak of transient failures pauses claims only for the affected source's jobs while healthy sources keep flowing, instead of pausing the whole worker pool. Paused sources are excluded at the claim query.
diff --git a/docs/python.md b/docs/python.md
index f9d6f5e0..c4cc9f0a 100644
--- a/docs/python.md
+++ b/docs/python.md
@@ -394,6 +394,31 @@ doc = await client.import_document(
 
 The `docling_document` provides rich metadata for visual grounding, page numbers, and section headings. Content is automatically extracted from the DoclingDocument.
 
+### Batch Import
+
+Each `create_document*` / `import_document` call writes new versions of the `documents`, `chunks`, and `document_items` tables. Ingesting many documents in a loop therefore creates a table version per document. Use `import_documents()` to write the whole batch in a single version per table:
+
+```python
+from haiku.rag.client import DocumentImport
+
+imports = []
+for path in paths:  # paths: list[Path]
+    docling_doc = await client.convert(path)
+    chunks = await client.chunk(docling_doc)
+    imports.append(
+        DocumentImport(
+            docling_document=docling_doc,
+            chunks=chunks,
+            uri=path.absolute().as_uri(),
+            metadata={"source": "external-pipeline"},
+        )
+    )
+
+docs = await client.import_documents(imports)
+```
+
+Chunks without embeddings are embedded automatically. The import is all-or-nothing: if any document fails, all tables are restored to their pre-batch state.
+
 See [Custom Processing Pipelines](custom-pipelines.md) for building pipelines with `convert()`, `chunk()`, and `embed_chunks()`.
 
 ## Maintenance
diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py
index 61badec5..33d422a5 100644
--- a/haiku_rag_slim/haiku/rag/client/__init__.py
+++ b/haiku_rag_slim/haiku/rag/client/__init__.py
@@ -14,6 +14,7 @@ from urllib.parse import urlparse
 
 import httpx
 
+from haiku.rag.client.documents import DocumentImport
 from haiku.rag.config import AppConfig, Config
 from haiku.rag.converters import get_converter
 from haiku.rag.reranking import get_reranker
@@ -237,6 +238,14 @@ class HaikuRAG:
             self, docling_document, chunks, uri, title, metadata
         )
 
+    async def import_documents(
+        self,
+        imports: "list[DocumentImport]",
+    ) -> list[Document]:
+        from haiku.rag.client.documents import import_documents
+
+        return await import_documents(self, imports)
+
     async def create_document_from_source(
         self,
         source: str | Path,
diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py
index bf395921..cbcb5fa0 100644
--- a/haiku_rag_slim/haiku/rag/client/documents.py
+++ b/haiku_rag_slim/haiku/rag/client/documents.py
@@ -3,6 +3,7 @@ import json
 import logging
 import mimetypes
 import tempfile
+from dataclasses import dataclass, field
 from pathlib import Path
 from typing import TYPE_CHECKING
 from urllib.parse import quote, unquote, urlparse
@@ -32,6 +33,23 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
+
+@dataclass
+class DocumentImport:
+    """A prepared document for batch import via ``import_documents``.
+
+    Carries the same inputs as ``import_document``: a converted
+    ``DoclingDocument``, its chunks (embeddings filled in if missing), and
+    optional uri/title/metadata.
+    """
+
+    docling_document: "DoclingDocument"
+    chunks: list[Chunk]
+    uri: str | None = None
+    title: str | None = None
+    metadata: dict | None = field(default=None)
+
+
 # Maximum length of an attachment chain rooted at a top-level ingest. With
 # value 3, a PDF whose attachments contain PDFs which themselves contain
 # PDFs is fully ingested (3 levels); a fourth nested level logs a warning
@@ -208,6 +226,81 @@ async def import_document(
     return await _store_document_with_chunks(client, document, chunks, docling_document)
 
 
+async def _store_documents_with_chunks(
+    client: "HaikuRAG",
+    prepared: list[tuple[Document, list[Chunk], "DoclingDocument"]],
+) -> list[Document]:
+    """Store many documents with their chunks in a single table version each.
+
+    Embeds any chunks that lack embeddings, then writes the documents, chunks,
+    and document_items tables once apiece. Restores all tables on any failure.
+    """
+    embedded: list[list[Chunk]] = [
+        await ensure_chunks_embedded(client._config, chunks, client.embedder)
+        for _, chunks, _ in prepared
+    ]
+
+    versions = await client.store.current_table_versions()
+
+    created = await client.document_repository.create([doc for doc, _, _ in prepared])
+
+    try:
+        all_chunks: list[Chunk] = []
+        all_items = []
+        for doc, doc_chunks, docling_document in zip(
+            created, embedded, (d for _, _, d in prepared)
+        ):
+            assert doc.id is not None
+            for order, chunk in enumerate(doc_chunks):
+                chunk.document_id = doc.id
+                chunk.order = order
+            all_chunks.extend(doc_chunks)
+            all_items.extend(extract_items(doc.id, docling_document))
+
+        await client.chunk_repository.create(all_chunks)
+        await client.document_item_repository.create_all(all_items)
+
+        if client._config.storage.auto_vacuum:
+            client._schedule_vacuum()
+
+        return created
+    except Exception:
+        await client.store.restore_table_versions(versions)
+        raise
+
+
+async def import_documents(
+    client: "HaikuRAG",
+    imports: list[DocumentImport],
+) -> list[Document]:
+    """Batch-import pre-processed documents with their chunks.
+
+    The batch analog of ``import_document``: writes the documents, chunks, and
+    document_items tables once each regardless of how many documents are
+    imported. Chunks without embeddings are embedded automatically.
+    """
+    if not imports:
+        return []
+
+    prepared: list[tuple[Document, list[Chunk], DoclingDocument]] = []
+    for item in imports:
+        content = item.docling_document.export_to_markdown()
+        title = item.title
+        if title is None:
+            title = await resolve_title(client._config, item.docling_document, content)
+
+        document = Document(
+            content=content,
+            uri=item.uri,
+            title=title,
+            metadata=item.metadata or {},
+        )
+        document.set_docling(item.docling_document)
+        prepared.append((document, item.chunks, item.docling_document))
+
+    return await _store_documents_with_chunks(client, prepared)
+
+
 async def _refresh_doc_metadata(
     client: "HaikuRAG",
     doc: Document,
diff --git a/tests/test_client.py b/tests/test_client.py
index b49cc7f2..571a9c6f 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -7,6 +7,7 @@ import httpx
 import pytest
 
 from haiku.rag.client import HaikuRAG
+from haiku.rag.client.documents import DocumentImport
 from haiku.rag.config import Config
 from haiku.rag.store.compression import decompress_json
 from haiku.rag.store.models.chunk import Chunk
@@ -776,6 +777,90 @@ async def test_client_import_document_with_custom_chunks(temp_db_path):
             )  # Original metadata preserved
 
 
+def _docling_doc(name: str, text: str):
+    from docling_core.types.doc.document import DoclingDocument
+    from docling_core.types.doc.labels import DocItemLabel
+
+    doc = DoclingDocument(name=name)
+    doc.add_text(label=DocItemLabel.TEXT, text=text)
+    return doc
+
+
+def _import(name: str, text: str, **overrides) -> "DocumentImport":
+    """Build a DocumentImport with one pre-embedded chunk (no embedder call)."""
+    dim = Config.embeddings.model.vector_dim
+    chunk = Chunk(content=text, embedding=[0.1] * dim, order=0)
+    return DocumentImport(
+        docling_document=_docling_doc(name, text),
+        chunks=[chunk],
+        **overrides,
+    )
+
+
+async def test_client_import_documents_single_version_per_table(temp_db_path):
+    """import_documents writes documents/chunks/document_items once for the
+    whole batch (issue #287)."""
+    config = Config.model_copy(deep=True)
+    config.storage.auto_vacuum = False
+
+    async with HaikuRAG(temp_db_path, config=config, create=True) as client:
+        imports = [
+            _import("a", "Alpha document body", uri="mem://a", title="Alpha"),
+            _import("b", "Beta document body", uri="mem://b", title="Beta"),
+            _import("c", "Gamma document body", uri="mem://c", title="Gamma"),
+        ]
+
+        before = await client.store.current_table_versions()
+        docs = await client.import_documents(imports)
+        after = await client.store.current_table_versions()
+
+        assert [d.title for d in docs] == ["Alpha", "Beta", "Gamma"]
+        assert all(d.id is not None for d in docs)
+        assert len({d.id for d in docs}) == 3
+
+        for table in ("documents", "chunks", "document_items"):
+            assert after[table] - before[table] == 1, table
+
+        for doc, expected in zip(docs, ("Alpha", "Beta", "Gamma")):
+            assert doc.id is not None
+            stored = await client.get_document_by_id(doc.id)
+            assert stored is not None and stored.title == expected
+            chunks = await client.chunk_repository.get_by_document_id(doc.id)
+            assert len(chunks) == 1 and chunks[0].document_id == doc.id
+            items = await client.document_item_repository.get_all_items(doc.id)
+            assert len(items) >= 1
+            assert all(i.document_id == doc.id for i in items)
+
+
+async def test_client_import_documents_rolls_back_on_failure(temp_db_path):
+    """A failure mid-batch restores all tables: nothing is persisted."""
+    dim = Config.embeddings.model.vector_dim
+
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        good = _import("good", "Good document body", uri="mem://good")
+        bad = DocumentImport(
+            docling_document=_docling_doc("bad", "Bad document body"),
+            chunks=[Chunk(content="bad", embedding=[0.1] * (dim + 1), order=0)],
+            uri="mem://bad",
+        )
+
+        with pytest.raises(Exception):
+            await client.import_documents([good, bad])
+
+        assert await client.count_documents() == 0
+
+
+async def test_client_import_documents_empty(temp_db_path):
+    """import_documents([]) returns [] and bumps no versions."""
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        before = await client.store.current_table_versions()
+        result = await client.import_documents([])
+        after = await client.store.current_table_versions()
+
+        assert result == []
+        assert after == before
+
+
 @pytest.mark.vcr()
 async def test_client_ask(allow_model_requests, temp_db_path):
     """Test asking questions returns answer and citations (VCR recorded)."""