Add HaikuRAG.import_documents for batch document import

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.

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

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,

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,

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)."""