Merge pull request #429 from ggozad/feat/batch-document-creation
Batch document import (import_documents)
This commit is contained in:
Yiorgis Gozadinos
GitHub
commit
d45f8f5de8
8 changed files with 368 additions and 31 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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)."""
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
Loading…
Reference in a new issue