danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

add document_items table for fast context expansion

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-04-14 18:22:34 +03:00
parent 661e0d34d8
commit 364b1bc509
No known key found for this signature in database
15 changed files with 780 additions and 36 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
CHANGELOG.md
View file

@ -1,6 +1,15 @@
# Changelog
## [Unreleased]
### Added
- **Document items table**: Pre-extracted document items stored as individual rows with scalar indexes, enabling context expansion via indexed range queries (~2.5ms) instead of full DoclingDocument deserialization (~8.7s for large documents)
### Changed
- **Database migration required**: Run `haiku-rag migrate` to populate `document_items` table for existing documents
- **Pin docling-core**: Upper bound added (`<2.72`) to prevent uncontrolled schema changes
## [0.39.0] - 2026-04-09
### Added

60
haiku_rag_slim/haiku/rag/client.py
View file

@ -20,11 +20,13 @@ from haiku.rag.reranking import get_reranker
from haiku.rag.store.engine import Store
from haiku.rag.store.models.chunk import Chunk, SearchResult
from haiku.rag.store.models.document import Document
from haiku.rag.store.models.document_item import extract_items
from haiku.rag.store.repositories.chunk import ChunkRepository
from haiku.rag.store.repositories.document import (
DocumentRepository,
_escape_sql_string,
)
from haiku.rag.store.repositories.document_item import DocumentItemRepository
from haiku.rag.store.repositories.settings import SettingsRepository
if TYPE_CHECKING:
@ -96,6 +98,7 @@ class HaikuRAG:
)
self.document_repository = DocumentRepository(self.store)
self.chunk_repository = ChunkRepository(self.store)
self.document_item_repository = DocumentItemRepository(self.store)
@property
def is_read_only(self) -> bool:
@ -354,6 +357,7 @@ class HaikuRAG:
self,
document: Document,
chunks: list[Chunk],
docling_document: "DoclingDocument",
) -> Document:
"""Store a document with chunks, embedding any that lack embeddings.
@ -362,6 +366,7 @@ class HaikuRAG:
Args:
document: The document to store (will be created).
chunks: Chunks to store (will be embedded if lacking embeddings).
docling_document: The DoclingDocument to extract items from.
Returns:
The created Document instance with ID set.
@ -389,6 +394,10 @@ class HaikuRAG:
# Batch create all chunks in a single operation
await self.chunk_repository.create(chunks)
# Extract and store document items for context expansion
items = extract_items(created_doc.id, docling_document)
await self.document_item_repository.create_items(created_doc.id, items)
# Vacuum old versions in background (non-blocking) if auto_vacuum enabled
if self._config.storage.auto_vacuum:
asyncio.create_task(self.store.vacuum())
@ -403,6 +412,7 @@ class HaikuRAG:
self,
document: Document,
chunks: list[Chunk],
docling_document: "DoclingDocument | None" = None,
) -> Document:
"""Update a document and replace its chunks, embedding any that lack embeddings.
@ -411,6 +421,8 @@ class HaikuRAG:
Args:
document: The document to update (must have ID set).
chunks: Chunks to replace existing (will be embedded if lacking embeddings).
docling_document: The DoclingDocument to extract items from.
When None, existing items are preserved.
Returns:
The updated Document instance.
@ -441,6 +453,14 @@ class HaikuRAG:
# Batch create all chunks in a single operation
await self.chunk_repository.create(chunks)
# Replace document items when a new DoclingDocument is provided
if docling_document is not None:
await self.document_item_repository.delete_by_document_id(
updated_doc.id
)
items = extract_items(updated_doc.id, docling_document)
await self.document_item_repository.create_items(updated_doc.id, items)
# Vacuum old versions in background (non-blocking) if auto_vacuum enabled
if self._config.storage.auto_vacuum:
asyncio.create_task(self.store.vacuum())
@ -498,7 +518,9 @@ class HaikuRAG:
document.set_docling(docling_document)
# Store document and chunks
return await self._store_document_with_chunks(document, embedded_chunks)
return await self._store_document_with_chunks(
document, embedded_chunks, docling_document
)
async def import_document(
self,
@ -536,7 +558,9 @@ class HaikuRAG:
)
document.set_docling(docling_document)
return await self._store_document_with_chunks(document, chunks)
return await self._store_document_with_chunks(
document, chunks, docling_document
)
async def create_document_from_source(
self, source: str | Path, title: str | None = None, metadata: dict | None = None
@ -677,7 +701,7 @@ class HaikuRAG:
docling_document, stored_content
)
return await self._update_document_with_chunks(
existing_doc, embedded_chunks
existing_doc, embedded_chunks, docling_document
)
else:
# Create new document
@ -690,7 +714,9 @@ class HaikuRAG:
metadata=metadata,
)
document.set_docling(docling_document)
return await self._store_document_with_chunks(document, embedded_chunks)
return await self._store_document_with_chunks(
document, embedded_chunks, docling_document
)
async def _create_or_update_document_from_url(
self, url: str, title: str | None = None, metadata: dict | None = None
@ -790,7 +816,7 @@ class HaikuRAG:
docling_document, stored_content
)
return await self._update_document_with_chunks(
existing_doc, embedded_chunks
existing_doc, embedded_chunks, docling_document
)
else:
# Create new document
@ -803,7 +829,9 @@ class HaikuRAG:
metadata=metadata,
)
document.set_docling(docling_document)
return await self._store_document_with_chunks(document, embedded_chunks)
return await self._store_document_with_chunks(
document, embedded_chunks, docling_document
)
def _get_extension_from_content_type_or_url(
self, url: str, content_type: str
@ -956,7 +984,9 @@ class HaikuRAG:
elif content is not None:
existing_doc.content = content
return await self._update_document_with_chunks(existing_doc, chunks)
return await self._update_document_with_chunks(
existing_doc, chunks, docling_document
)
# DoclingDocument provided without chunks - chunk and embed using primitives
if docling_document is not None:
@ -966,7 +996,7 @@ class HaikuRAG:
new_chunks = await self.chunk(docling_document)
embedded_chunks = await embed_chunks(new_chunks, self._config)
return await self._update_document_with_chunks(
existing_doc, embedded_chunks
existing_doc, embedded_chunks, docling_document
)
# Content provided without chunks - convert, chunk, and embed using primitives
@ -977,7 +1007,9 @@ class HaikuRAG:
new_chunks = await self.chunk(converted_docling)
embedded_chunks = await embed_chunks(new_chunks, self._config)
return await self._update_document_with_chunks(existing_doc, embedded_chunks)
return await self._update_document_with_chunks(
existing_doc, embedded_chunks, converted_docling
)
async def delete_document(self, document_id: str) -> bool:
"""Delete a document by its ID."""
@ -1758,6 +1790,7 @@ class HaikuRAG:
"""Batch write documents and chunks during rebuild.
This performs two writes: one for all document updates, one for all chunks.
Also repopulates document items from the stored docling document.
Used by RECHUNK and FULL modes after the chunks table has been cleared.
"""
from haiku.rag.store.engine import DocumentRecord
@ -1800,6 +1833,15 @@ class HaikuRAG:
if chunks:
await self.chunk_repository.create(chunks)
# Repopulate document items from stored docling data
for doc in documents:
assert doc.id is not None
docling_doc = doc.get_docling_document()
if docling_doc is not None:
await self.document_item_repository.delete_by_document_id(doc.id)
items = extract_items(doc.id, docling_doc)
await self.document_item_repository.create_items(doc.id, items)
async def _rebuild_rechunk(
self, documents: list[Document]
) -> AsyncGenerator[str, None]:

33
haiku_rag_slim/haiku/rag/store/engine.py
View file

@ -107,6 +107,15 @@ def create_chunk_model(vector_dim: int):
return ChunkRecord
class DocumentItemRecord(LanceModel):
document_id: str
position: int
self_ref: str
label: str = Field(default="")
text: str = Field(default="")
page_numbers: str = Field(default="[]")
class SettingsRecord(LanceModel):
id: str = Field(default="settings")
settings: str = Field(default="{}")
@ -256,6 +265,7 @@ class Store:
for table in [
self.documents_table,
self.chunks_table,
self.document_items_table,
self.settings_table,
]:
table.optimize(cleanup_older_than=retention)
@ -358,7 +368,7 @@ class Store:
def _init_tables(self):
"""Initialize database tables (create if they don't exist)."""
existing_tables = self.db.list_tables().tables
required_tables = {"documents", "chunks", "settings"}
required_tables = {"documents", "chunks", "document_items", "settings"}
missing_tables = required_tables - set(existing_tables)
if missing_tables and self._read_only:
@ -385,6 +395,23 @@ class Store:
"content_fts", replace=True, with_position=True, remove_stop_words=False
)
# Create or open document_items table
if "document_items" in existing_tables:
self.document_items_table = self.db.open_table("document_items")
else:
self.document_items_table = self.db.create_table(
"document_items", schema=DocumentItemRecord
)
self.document_items_table.create_scalar_index(
"document_id", index_type="BTREE", replace=True
)
self.document_items_table.create_scalar_index(
"position", index_type="BTREE", replace=True
)
self.document_items_table.create_scalar_index(
"self_ref", index_type="BTREE", replace=True
)
# Create or open settings table
if "settings" in existing_tables:
self.settings_table = self.db.open_table("settings")
@ -528,6 +555,7 @@ class Store:
return {
"documents": int(self.documents_table.version),
"chunks": int(self.chunks_table.version),
"document_items": int(self.document_items_table.version),
"settings": int(self.settings_table.version),
}
@ -540,6 +568,7 @@ class Store:
self._assert_writable()
self.documents_table.restore(int(versions["documents"]))
self.chunks_table.restore(int(versions["chunks"]))
self.document_items_table.restore(int(versions["document_items"]))
self.settings_table.restore(int(versions["settings"]))
return True
@ -569,6 +598,7 @@ class Store:
tables = [
("documents", self.documents_table),
("chunks", self.chunks_table),
("document_items", self.document_items_table),
("settings", self.settings_table),
]
@ -620,6 +650,7 @@ class Store:
table_map = {
"documents": self.documents_table,
"chunks": self.chunks_table,
"document_items": self.document_items_table,
"settings": self.settings_table,
}
table = table_map.get(table_name)

2
haiku_rag_slim/haiku/rag/store/models/__init__.py
View file

@ -1,10 +1,12 @@
from .chunk import BoundingBox, Chunk, ChunkMetadata, SearchResult
from .document import Document
from .document_item import DocumentItem
__all__ = [
"BoundingBox",
"Chunk",
"ChunkMetadata",
"Document",
"DocumentItem",
"SearchResult",
]

86
haiku_rag_slim/haiku/rag/store/models/document_item.py Normal file
View file

@ -0,0 +1,86 @@
from typing import TYPE_CHECKING
from pydantic import BaseModel
if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument, NodeItem
class DocumentItem(BaseModel):
document_id: str
position: int
self_ref: str
label: str = ""
text: str = ""
page_numbers: list[int] = []
def extract_item_text(item: "NodeItem", docling_doc: "DoclingDocument") -> str | None:
"""Extract text content from a DocItem.
Handles different item types:
- TextItem, SectionHeaderItem, etc.: Use .text attribute
- TableItem: Use export_to_markdown() for table content
- PictureItem: Use export_to_markdown() with PLACEHOLDER mode to avoid base64
"""
from docling_core.types.doc.base import ImageRefMode
from docling_core.types.doc.document import PictureItem, TableItem
if text := getattr(item, "text", None):
return text
if isinstance(item, PictureItem):
return item.export_to_markdown(
docling_doc,
image_mode=ImageRefMode.PLACEHOLDER,
image_placeholder="",
)
if isinstance(item, TableItem):
try:
return item.export_to_markdown(docling_doc)
except Exception:
pass
if caption := getattr(item, "caption", None):
if hasattr(caption, "text"):
return caption.text
return None
def extract_items(
document_id: str, docling_doc: "DoclingDocument"
) -> list[DocumentItem]:
"""Extract document items from a DoclingDocument for the items table.
Runs iterate_items() and extracts the fields needed for context expansion:
self_ref, label, pre-rendered text, and page numbers from provenance.
"""
items: list[DocumentItem] = []
for position, (item, _level) in enumerate(docling_doc.iterate_items()):
label = getattr(item, "label", None)
label_str = str(label.value) if hasattr(label, "value") else str(label or "")
text = extract_item_text(item, docling_doc) or ""
page_numbers: list[int] = []
if prov := getattr(item, "prov", None):
for p in prov:
page_no = getattr(p, "page_no", None)
if page_no is not None and page_no not in page_numbers:
page_numbers.append(page_no)
items.append(
DocumentItem(
document_id=document_id,
position=position,
self_ref=item.self_ref,
label=label_str,
text=text,
page_numbers=sorted(page_numbers),
)
)
return items

2
haiku_rag_slim/haiku/rag/store/repositories/__init__.py
View file

@ -1,9 +1,11 @@
from haiku.rag.store.repositories.chunk import ChunkRepository
from haiku.rag.store.repositories.document import DocumentRepository
from haiku.rag.store.repositories.document_item import DocumentItemRepository
from haiku.rag.store.repositories.settings import SettingsRepository
__all__ = [
"ChunkRepository",
"DocumentItemRepository",
"DocumentRepository",
"SettingsRepository",
]

33
haiku_rag_slim/haiku/rag/store/repositories/document.py
View file

@ -17,6 +17,7 @@ class DocumentRepository:
def __init__(self, store: Store) -> None:
self.store = store
self._chunk_repository = None
self._document_item_repository = None
@property
def chunk_repository(self):
@ -27,6 +28,17 @@ class DocumentRepository:
self._chunk_repository = ChunkRepository(self.store)
return self._chunk_repository
@property
def document_item_repository(self):
"""Lazy-load DocumentItemRepository when needed."""
if self._document_item_repository is None:
from haiku.rag.store.repositories.document_item import (
DocumentItemRepository,
)
self._document_item_repository = DocumentItemRepository(self.store)
return self._document_item_repository
def _record_to_document(self, record: DocumentRecord) -> Document:
"""Convert a DocumentRecord to a Document model."""
return Document(
@ -179,11 +191,11 @@ class DocumentRepository:
if doc is None:
return False
# Invalidate cache before delete
invalidate_docling_document_cache(entity_id)
# Delete associated chunks first
# Delete associated chunks and items first
await self.chunk_repository.delete_by_document_id(entity_id)
await self.document_item_repository.delete_by_document_id(entity_id)
# Delete the document
safe_id = _escape_sql_string(entity_id)
@ -272,8 +284,23 @@ class DocumentRepository:
async def delete_all(self) -> None:
"""Delete all documents from the database."""
self.store._assert_writable()
# Delete all chunks first
from haiku.rag.store.engine import DocumentItemRecord
# Delete all chunks and items first
await self.chunk_repository.delete_all()
self.store.db.drop_table("document_items")
self.store.document_items_table = self.store.db.create_table(
"document_items", schema=DocumentItemRecord
)
self.store.document_items_table.create_scalar_index(
"document_id", index_type="BTREE", replace=True
)
self.store.document_items_table.create_scalar_index(
"position", index_type="BTREE", replace=True
)
self.store.document_items_table.create_scalar_index(
"self_ref", index_type="BTREE", replace=True
)
# Get count before deletion
count = len(

90
haiku_rag_slim/haiku/rag/store/repositories/document_item.py Normal file
View file

@ -0,0 +1,90 @@
import json
from haiku.rag.store.engine import DocumentItemRecord, Store
from haiku.rag.store.models.document_item import DocumentItem
def _escape_sql_string(value: str) -> str:
"""Escape single quotes in SQL string literals."""
return value.replace("'", "''")
class DocumentItemRepository:
"""Repository for DocumentItem operations."""
def __init__(self, store: Store) -> None:
self.store = store
def _record_to_item(self, row: dict) -> DocumentItem:
return DocumentItem(
document_id=row["document_id"],
position=row["position"],
self_ref=row["self_ref"],
label=row.get("label", ""),
text=row.get("text", ""),
page_numbers=json.loads(row.get("page_numbers", "[]")),
)
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),
)
for item in items
]
self.store.document_items_table.add(records)
async def get_items_in_range(
self, document_id: str, start: int, end: int
) -> list[DocumentItem]:
"""Get items for a document within a position range (inclusive)."""
safe_id = _escape_sql_string(document_id)
rows = (
self.store.document_items_table.search()
.where(
f"document_id = '{safe_id}' "
f"AND position >= {start} AND position <= {end}"
)
.to_list()
)
items = [self._record_to_item(row) for row in rows]
items.sort(key=lambda x: x.position)
return items
async def resolve_refs(self, document_id: str, refs: list[str]) -> dict[str, int]:
"""Resolve self_refs to positions. Returns {self_ref: position}."""
if not refs:
return {}
safe_id = _escape_sql_string(document_id)
refs_sql = ", ".join(f"'{_escape_sql_string(r)}'" for r in refs)
rows = (
self.store.document_items_table.search()
.select(["self_ref", "position"])
.where(f"document_id = '{safe_id}' AND self_ref IN ({refs_sql})")
.to_list()
)
return {row["self_ref"]: row["position"] for row in rows}
async def get_item_count(self, document_id: str) -> int:
"""Count items for a document."""
safe_id = _escape_sql_string(document_id)
return self.store.document_items_table.count_rows(
filter=f"document_id = '{safe_id}'"
)
async def delete_by_document_id(self, document_id: str) -> None:
"""Delete all items for a document."""
self.store._assert_writable()
safe_id = _escape_sql_string(document_id)
self.store.document_items_table.delete(f"document_id = '{safe_id}'")

4
haiku_rag_slim/haiku/rag/store/upgrades/__init__.py
View file

@ -81,8 +81,12 @@ from haiku.rag.store.upgrades.v0_25_0 import (
from haiku.rag.store.upgrades.v0_38_0 import (
upgrade_split_pages_zstd as upgrade_0_38_0_split_pages,
)
from haiku.rag.store.upgrades.v0_40_0 import (
upgrade_populate_document_items as upgrade_0_40_0_document_items,
)
upgrades.append(upgrade_0_20_0_docling)
upgrades.append(upgrade_0_23_1_contextualize)
upgrades.append(upgrade_0_25_0_compress)
upgrades.append(upgrade_0_38_0_split_pages)
upgrades.append(upgrade_0_40_0_document_items)

97
haiku_rag_slim/haiku/rag/store/upgrades/v0_40_0.py Normal file
View file

@ -0,0 +1,97 @@
import json
import logging
from haiku.rag.store.engine import DocumentItemRecord, Store
from haiku.rag.store.upgrades import Upgrade
logger = logging.getLogger(__name__)
def _apply_populate_document_items(store: Store) -> None: # pragma: no cover
"""Populate document_items table from existing docling documents."""
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.store.compression import decompress_json
from haiku.rag.store.models.document_item import extract_items
# Get all document IDs that have docling data
ids = [
row["id"]
for row in store.documents_table.search().select(["id"]).to_arrow().to_pylist()
]
if not ids:
logger.info("No documents to migrate")
return
total = len(ids)
logger.info("Populating document_items for %d documents", total)
migrated = 0
skipped = 0
for idx, doc_id in enumerate(ids, 1):
# Load only docling data
safe_id = doc_id.replace("'", "''")
rows = (
store.documents_table.search()
.select(["id", "docling_document"])
.where(f"id = '{safe_id}'")
.limit(1)
.to_list()
)
if not rows:
skipped += 1
continue
row = rows[0]
docling_blob = row.get("docling_document")
if not docling_blob or not isinstance(docling_blob, bytes):
skipped += 1
continue
try:
json_str = decompress_json(docling_blob)
docling_doc = DoclingDocument.model_validate_json(json_str)
items = extract_items(doc_id, docling_doc)
if items:
records = [
DocumentItemRecord(
document_id=item.document_id,
position=item.position,
self_ref=item.self_ref,
label=item.label,
text=item.text,
page_numbers=json.dumps(item.page_numbers),
)
for item in items
]
store.document_items_table.add(records)
migrated += 1
if idx % 10 == 0 or idx == total:
logger.info(
"Progress: %d/%d documents (%d migrated, %d skipped)",
idx,
total,
migrated,
skipped,
)
except Exception:
logger.warning("Failed to extract items for document %s", doc_id)
skipped += 1
logger.info(
"Migration complete: %d migrated, %d skipped out of %d",
migrated,
skipped,
total,
)
upgrade_populate_document_items = Upgrade(
version="0.40.0",
apply=_apply_populate_document_items,
description="Populate document_items table for context expansion",
)

2
haiku_rag_slim/pyproject.toml
View file

@ -2,7 +2,7 @@
name = "haiku.rag-slim"
description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling - Minimal dependencies"
version = "0.39.0"
version = "0.40.0"
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
license = { text = "MIT" }
readme = { file = "README.md", content-type = "text/markdown" }

4
pyproject.toml
View file

@ -2,7 +2,7 @@
name = "haiku.rag"
description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling"
version = "0.39.0"
version = "0.40.0"
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
license = { text = "MIT" }
readme = { file = "README.md", content-type = "text/markdown" }
@ -30,7 +30,7 @@ classifiers = [
]
dependencies = [
"haiku.rag-slim[docling,voyageai,mxbai,cohere,zeroentropy,tui]==0.39.0",
"haiku.rag-slim[docling,voyageai,mxbai,cohere,zeroentropy,tui]==0.40.0",
]
[project.scripts]

351
tests/store/test_document_items.py Normal file
View file

@ -0,0 +1,351 @@
import pytest
from haiku.rag.client import HaikuRAG
from haiku.rag.store.engine import Store
from haiku.rag.store.models.document_item import (
DocumentItem,
extract_item_text,
extract_items,
)
from haiku.rag.store.repositories.document_item import DocumentItemRepository
def _make_docling_doc():
"""Create a DoclingDocument with mixed item types for testing."""
from docling_core.types.doc.document import DoclingDocument, TableData
from docling_core.types.doc.labels import DocItemLabel
doc = DoclingDocument(name="test")
doc.add_text(label=DocItemLabel.SECTION_HEADER, text="Introduction")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="This is the first paragraph.")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="This is the second paragraph.")
doc.add_table(data=TableData(num_rows=2, num_cols=2, table_cells=[]))
doc.add_text(label=DocItemLabel.SECTION_HEADER, text="Conclusion")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Final thoughts here.")
return doc
class TestExtractItems:
def test_extracts_all_items(self):
doc = _make_docling_doc()
items = extract_items("doc-1", doc)
assert len(items) == 6
assert all(item.document_id == "doc-1" for item in items)
assert [item.position for item in items] == [0, 1, 2, 3, 4, 5]
def test_extracts_labels(self):
doc = _make_docling_doc()
items = extract_items("doc-1", doc)
assert items[0].label == "section_header"
assert items[1].label == "paragraph"
assert items[3].label == "table"
assert items[4].label == "section_header"
def test_extracts_text(self):
doc = _make_docling_doc()
items = extract_items("doc-1", doc)
assert items[0].text == "Introduction"
assert items[1].text == "This is the first paragraph."
assert items[5].text == "Final thoughts here."
def test_extracts_self_refs(self):
doc = _make_docling_doc()
items = extract_items("doc-1", doc)
assert all(item.self_ref.startswith("#/") for item in items)
def test_table_gets_markdown_text(self):
doc = _make_docling_doc()
items = extract_items("doc-1", doc)
table_item = items[3]
assert table_item.label == "table"
# Table should have some text from export_to_markdown
assert isinstance(table_item.text, str)
class TestExtractItemText:
def test_text_item(self):
from docling_core.types.doc.document import DoclingDocument
from docling_core.types.doc.labels import DocItemLabel
doc = DoclingDocument(name="test")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Hello world")
item, _ = next(iter(doc.iterate_items()))
assert extract_item_text(item, doc) == "Hello world"
def test_returns_none_for_empty_item(self):
from docling_core.types.doc.document import DoclingDocument
doc = DoclingDocument(name="test")
# An empty doc has no items to extract text from
items = extract_items("doc-1", doc)
assert items == []
@pytest.mark.asyncio
class TestDocumentItemRepository:
async def test_create_and_get_range(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag:
repo = DocumentItemRepository(rag.store)
items = [
DocumentItem(
document_id="doc-1",
position=i,
self_ref=f"#/texts/{i}",
label="paragraph",
text=f"Item {i}",
page_numbers=[1],
)
for i in range(10)
]
await repo.create_items("doc-1", items)
result = await repo.get_items_in_range("doc-1", 3, 7)
assert len(result) == 5
assert result[0].position == 3
assert result[-1].position == 7
assert result[0].text == "Item 3"
async def test_resolve_refs(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag:
repo = DocumentItemRepository(rag.store)
items = [
DocumentItem(
document_id="doc-1",
position=i,
self_ref=f"#/texts/{i}",
label="paragraph",
text=f"Item {i}",
)
for i in range(10)
]
await repo.create_items("doc-1", items)
refs = await repo.resolve_refs(
"doc-1", ["#/texts/2", "#/texts/7", "#/texts/999"]
)
assert refs == {"#/texts/2": 2, "#/texts/7": 7}
async def test_get_item_count(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag:
repo = DocumentItemRepository(rag.store)
items = [
DocumentItem(
document_id="doc-1",
position=i,
self_ref=f"#/texts/{i}",
label="paragraph",
text=f"Item {i}",
)
for i in range(15)
]
await repo.create_items("doc-1", items)
assert await repo.get_item_count("doc-1") == 15
assert await repo.get_item_count("nonexistent") == 0
async def test_delete_by_document_id(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag:
repo = DocumentItemRepository(rag.store)
for doc_id in ["doc-1", "doc-2"]:
items = [
DocumentItem(
document_id=doc_id,
position=i,
self_ref=f"#/texts/{i}",
label="paragraph",
text=f"Item {i}",
)
for i in range(5)
]
await repo.create_items(doc_id, items)
await repo.delete_by_document_id("doc-1")
assert await repo.get_item_count("doc-1") == 0
assert await repo.get_item_count("doc-2") == 5
async def test_empty_refs_returns_empty(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag:
repo = DocumentItemRepository(rag.store)
assert await repo.resolve_refs("doc-1", []) == {}
async def test_items_sorted_by_position(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag:
repo = DocumentItemRepository(rag.store)
# Insert in reverse order
items = [
DocumentItem(
document_id="doc-1",
position=9 - i,
self_ref=f"#/texts/{9 - i}",
label="paragraph",
text=f"Item {9 - i}",
)
for i in range(10)
]
await repo.create_items("doc-1", items)
result = await repo.get_items_in_range("doc-1", 0, 9)
positions = [item.position for item in result]
assert positions == sorted(positions)
@pytest.mark.asyncio
class TestDocumentItemPopulation:
async def test_store_document_populates_items(self, temp_db_path):
"""Test that _store_document_with_chunks populates items when given a docling_document."""
from haiku.rag.store.models.document import Document
docling_doc = _make_docling_doc()
async with HaikuRAG(temp_db_path, create=True) as rag:
document = Document(
content="test content",
uri="test://doc",
)
document.set_docling(docling_doc)
# Use _store_document_with_chunks directly with empty chunks
# to avoid needing embeddings
created = await rag._store_document_with_chunks(document, [], docling_doc)
assert created.id is not None
count = await rag.document_item_repository.get_item_count(created.id)
assert count == 6
items = await rag.document_item_repository.get_items_in_range(
created.id, 0, count
)
assert items[0].label == "section_header"
assert items[0].text == "Introduction"
assert items[1].label == "paragraph"
async def test_update_document_replaces_items(self, temp_db_path):
"""Test that _update_document_with_chunks replaces items."""
from docling_core.types.doc.document import DoclingDocument
from docling_core.types.doc.labels import DocItemLabel
from haiku.rag.store.models.document import Document
docling_doc = _make_docling_doc()
async with HaikuRAG(temp_db_path, create=True) as rag:
document = Document(
content="test content",
uri="test://doc",
)
document.set_docling(docling_doc)
created = await rag._store_document_with_chunks(document, [], docling_doc)
assert created.id is not None
assert await rag.document_item_repository.get_item_count(created.id) == 6
# Update with a simpler document
new_doc = DoclingDocument(name="updated")
new_doc.add_text(label=DocItemLabel.PARAGRAPH, text="Only one item now.")
created.set_docling(new_doc)
await rag._update_document_with_chunks(created, [], new_doc)
assert await rag.document_item_repository.get_item_count(created.id) == 1
async def test_delete_document_cascades_items(self, temp_db_path):
"""Test that deleting a document also deletes its items."""
from haiku.rag.store.models.document import Document
docling_doc = _make_docling_doc()
async with HaikuRAG(temp_db_path, create=True) as rag:
document = Document(
content="test content",
uri="test://doc",
)
document.set_docling(docling_doc)
created = await rag._store_document_with_chunks(document, [], docling_doc)
assert created.id is not None
assert await rag.document_item_repository.get_item_count(created.id) == 6
await rag.delete_document(created.id)
assert await rag.document_item_repository.get_item_count(created.id) == 0
class TestDocumentItemMigration:
def test_migration_populates_items_for_existing_documents(self, temp_db_path):
"""Test that the v0.40.0 migration populates items for pre-existing documents."""
from haiku.rag.store.compression import compress_docling_split
from haiku.rag.store.engine import DocumentRecord
docling_doc = _make_docling_doc()
json_str = docling_doc.model_dump_json()
structure, pages = compress_docling_split(json_str)
# Create a database at a pre-migration version with a document
store = Store(temp_db_path, create=True, skip_migration_check=True)
store.set_haiku_version("0.39.0")
doc_record = DocumentRecord(
id="test-doc-1",
content="test content",
uri="test://doc",
docling_document=structure,
docling_pages=pages,
docling_version=docling_doc.version,
)
store.documents_table.add([doc_record])
# Verify no items exist yet
assert store.document_items_table.count_rows() == 0
store.close()
# Re-open with skip_migration_check and run migration
store = Store(temp_db_path, skip_migration_check=True)
applied = store.migrate()
# Should have applied the v0.40.0 migration
assert any("document_items" in desc for desc in applied)
# Items should now exist
item_count = store.document_items_table.count_rows(
filter="document_id = 'test-doc-1'"
)
assert item_count == 6
# Verify item content
items = (
store.document_items_table.search()
.where("document_id = 'test-doc-1'")
.to_list()
)
labels = {row["label"] for row in items}
assert "section_header" in labels
assert "paragraph" in labels
assert "table" in labels
store.close()
def test_migration_skips_documents_without_docling(self, temp_db_path):
"""Test that migration handles documents without docling data."""
from haiku.rag.store.engine import DocumentRecord
store = Store(temp_db_path, create=True, skip_migration_check=True)
store.set_haiku_version("0.39.0")
doc_record = DocumentRecord(
id="no-docling",
content="plain text document",
)
store.documents_table.add([doc_record])
store.close()
store = Store(temp_db_path, skip_migration_check=True)
store.migrate()
# No items should have been created
assert store.document_items_table.count_rows() == 0
store.close()

3
tests/test_info.py
View file

@ -5,6 +5,7 @@ import pytest
from haiku.rag.app import HaikuRAGApp
from haiku.rag.config.models import AppConfig, LanceDBConfig
from haiku.rag.store.engine import DocumentItemRecord
@pytest.mark.asyncio
@ -33,6 +34,7 @@ async def test_app_info_outputs(temp_db_path, capsys):
settings_tbl = db.create_table("settings", schema=SettingsRecord)
docs_tbl = db.create_table("documents", schema=DocumentRecord)
chunks_tbl = db.create_table("chunks", schema=ChunkRecord)
db.create_table("document_items", schema=DocumentItemRecord)
# Insert one of each - using the new config format
settings_tbl.add(
@ -113,6 +115,7 @@ async def test_app_info_with_vector_index(temp_db_path, capsys):
settings_tbl = db.create_table("settings", schema=SettingsRecord)
docs_tbl = db.create_table("documents", schema=DocumentRecord)
chunks_tbl = db.create_table("chunks", schema=ChunkRecord)
db.create_table("document_items", schema=DocumentItemRecord)
# Insert settings
settings_tbl.add(

40
uv.lock
View file

@ -40,14 +40,14 @@ wheels = [
[[package]]
name = "ag-ui-protocol"
version = "0.1.14"
version = "0.1.15"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "pydantic" },
]
sdist = { url = "https://files.pythonhosted.org/packages/53/ee/319d189343e1dc67b1109c950a0d1091fe32498104b5917fbbd806ff58dd/ag_ui_protocol-0.1.14.tar.gz", hash = "sha256:d8e86b308f86a6cf6a5e18ca7154d7642895de2fe94cd2cece57723cdbba6406", size = 5687, upload-time = "2026-03-18T00:43:13.358Z" }
sdist = { url = "https://files.pythonhosted.org/packages/57/71/96c21ae7e2fb9b610c1a90d38bd2de8b6e5b2900a63001f3882f43e519af/ag_ui_protocol-0.1.15.tar.gz", hash = "sha256:5e23c1042c7d4e364d685e68d2fb74d37c16bc83c66d270102d8eaedce56ad82", size = 6269, upload-time = "2026-04-01T15:44:33.136Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/0c/99/eaa83816924791fc25ebb44ac7987196b687a3fdf597b1e7a62c69306a8d/ag_ui_protocol-0.1.14-py3-none-any.whl", hash = "sha256:ec072e6a45e0d45b8714e6d54919cc9bde3d097fdc36f7e82953b2f21f1cdbef", size = 8069, upload-time = "2026-03-18T00:43:12.1Z" },
{ url = "https://files.pythonhosted.org/packages/e4/a0/a73398d30bb0f9ad70cd70426151a4a19527a7296e48a3a16a50e1d5db05/ag_ui_protocol-0.1.15-py3-none-any.whl", hash = "sha256:85cde077023ccbc37b5ce2ad953537883c262d210320f201fc2ec4e85408b06a", size = 8661, upload-time = "2026-04-01T15:44:32.079Z" },
]
[[package]]
@ -1418,7 +1418,7 @@ wheels = [
[[package]]
name = "haiku-rag"
version = "0.39.0"
version = "0.40.0"
source = { editable = "." }
dependencies = [
{ name = "haiku-rag-slim", extra = ["cohere", "docling", "mxbai", "tui", "voyageai", "zeroentropy"] },
@ -1493,14 +1493,14 @@ requires-dist = [
{ name = "gepa", specifier = ">=0.1.0" },
{ name = "haiku-rag-slim", editable = "haiku_rag_slim" },
{ name = "huggingface-hub", specifier = ">=0.20.0" },
{ name = "pydantic-ai-slim", extras = ["evals", "logfire"], specifier = ">=1.70.0" },
{ name = "pydantic-ai-slim", extras = ["evals", "logfire"], specifier = ">=1.81.0" },
{ name = "python-dotenv", specifier = ">=1.2.2" },
{ name = "typer", specifier = ">=0.21.0,<0.22.0" },
]
[[package]]
name = "haiku-rag-slim"
version = "0.39.0"
version = "0.40.0"
source = { editable = "haiku_rag_slim" }
dependencies = [
{ name = "cachetools" },
@ -1572,7 +1572,7 @@ requires-dist = [
{ name = "cohere", marker = "extra == 'cohere'", specifier = ">=5.21.1" },
{ name = "docling", marker = "extra == 'docling'", specifier = ">=2.84.0" },
{ name = "docling-core", specifier = ">=2.71.0,<2.72" },
{ name = "haiku-skills", specifier = ">=0.13.0" },
{ name = "haiku-skills", specifier = ">=0.13.3" },
{ name = "httpx", specifier = ">=0.28.1" },
{ name = "jinja2", specifier = ">=3.1.0" },
{ name = "jsonpatch", specifier = ">=1.33" },
@ -1586,7 +1586,7 @@ requires-dist = [
{ name = "pydantic-ai-slim", extras = ["google"], marker = "extra == 'google'" },
{ name = "pydantic-ai-slim", extras = ["groq"], marker = "extra == 'groq'" },
{ name = "pydantic-ai-slim", extras = ["mistral"], marker = "extra == 'mistral'" },
{ name = "pydantic-ai-slim", extras = ["openai", "fastmcp", "logfire", "ag-ui"], specifier = ">=1.77.0" },
{ name = "pydantic-ai-slim", extras = ["openai", "fastmcp", "logfire", "ag-ui"], specifier = ">=1.81.0" },
{ name = "pydantic-ai-slim", extras = ["vertexai"], marker = "extra == 'vertexai'" },
{ name = "pydantic-ai-slim", extras = ["voyageai"], marker = "extra == 'voyageai'" },
{ name = "pydantic-monty", specifier = ">=0.0.9" },
@ -1606,7 +1606,7 @@ provides-extras = ["docling", "voyageai", "mxbai", "cohere", "zeroentropy", "jin
[[package]]
name = "haiku-skills"
version = "0.13.0"
version = "0.13.3"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "ag-ui-protocol" },
@ -1616,9 +1616,9 @@ dependencies = [
{ name = "pyyaml" },
{ name = "skills-ref" },
]
sdist = { url = "https://files.pythonhosted.org/packages/5f/61/5c1a8e7978b1eab1a17fb65026f8f837f48f34830f43ac944feb473c71b7/haiku_skills-0.13.0.tar.gz", hash = "sha256:02bca51247254d685c26a142d6f5321279ef10c4aed5aee90e8e4db996e82c26", size = 249960, upload-time = "2026-03-27T15:51:00.074Z" }
sdist = { url = "https://files.pythonhosted.org/packages/9c/65/44140f3289ebf3647c6f3dbf051de2162bc4ade11f020f98676e14d1d8b7/haiku_skills-0.13.3.tar.gz", hash = "sha256:4c918d6c8183c01274a34111159eeb56d1befed5d38a39779d9e511f92d4e088", size = 250458, upload-time = "2026-04-14T09:13:29.975Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/84/85/740e9a778bede480aa5540c7f520803a0226636f425f2b0320bada13a02c/haiku_skills-0.13.0-py3-none-any.whl", hash = "sha256:6f18c84e31de3b67c9476f7a19b1719877314dcaceb7a0b94963ca9d43ccac2f", size = 31451, upload-time = "2026-03-27T15:50:59.072Z" },
{ url = "https://files.pythonhosted.org/packages/b9/c6/e3aaed844b80c94b8e8462b80751b372061b19bf17f57e531f327e0366f7/haiku_skills-0.13.3-py3-none-any.whl", hash = "sha256:2bf0d28fd3e7c65b627f1bec33245bc36dc997d95d6efcea949cb9260aec2d73", size = 31530, upload-time = "2026-04-14T09:13:28.787Z" },
]
[[package]]
@ -3622,7 +3622,7 @@ email = [
[[package]]
name = "pydantic-ai-slim"
version = "1.77.0"
version = "1.81.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "genai-prices" },
@ -3633,9 +3633,9 @@ dependencies = [
{ name = "pydantic-graph" },
{ name = "typing-inspection" },
]
sdist = { url = "https://files.pythonhosted.org/packages/79/a7/ad011e626bed1f275fbaf933181573a50b05c2b9a0be927583d46fb8ff13/pydantic_ai_slim-1.77.0.tar.gz", hash = "sha256:a6e7006a4b048193d45b6ba816d301271e3f5ef1cdc4f9fb340617f382c6ce0d", size = 518781, upload-time = "2026-04-03T02:16:54.524Z" }
sdist = { url = "https://files.pythonhosted.org/packages/85/e9/8fbc609f28cb708bfcc5db7864d40bd4ac84f3e3780321ca50d7739f3c3d/pydantic_ai_slim-1.81.0.tar.gz", hash = "sha256:9895e2d3ae46b8e0342af5b862c987cfb86df87ef887dc8352e372b183db6c06", size = 550997, upload-time = "2026-04-14T01:47:58.757Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/96/c5/5913cc4ae99047901c602f0d8208e3a75a7952b7e57d76169547307d7cea/pydantic_ai_slim-1.77.0-py3-none-any.whl", hash = "sha256:110c516935de384f1beddc36fda04e8df36cdf5bee3a5bfd0da562726182e52b", size = 664494, upload-time = "2026-04-03T02:16:46.668Z" },
{ url = "https://files.pythonhosted.org/packages/e4/81/dc7062fc325ebb448ed9311edfde895c009ddf3cbd3d65c815b8a52c4bf8/pydantic_ai_slim-1.81.0-py3-none-any.whl", hash = "sha256:1d3dd19a53bcdcc9baf75d7b60daee87a80f0647df221776a76b177f4526ed2a", size = 705019, upload-time = "2026-04-14T01:47:51.399Z" },
]
[package.optional-dependencies]
@ -3755,7 +3755,7 @@ wheels = [
[[package]]
name = "pydantic-evals"
version = "1.77.0"
version = "1.81.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "anyio" },
@ -3765,14 +3765,14 @@ dependencies = [
{ name = "pyyaml" },
{ name = "rich" },
]
sdist = { url = "https://files.pythonhosted.org/packages/de/23/8da617d3803362325790e73e6219dec359559c2ef2350b35dbeb9e39c92f/pydantic_evals-1.77.0.tar.gz", hash = "sha256:64a12324c9a3f4fefa34b5a5eb4c2320c0976e425404372fa69f2872f585c3d0", size = 65811, upload-time = "2026-04-03T02:16:55.71Z" }
sdist = { url = "https://files.pythonhosted.org/packages/5e/47/8e1bc88e1fce1a636c4d4bc5fe883f37d2737f9c03efcd7aeb9238ebf07c/pydantic_evals-1.81.0.tar.gz", hash = "sha256:70fd1d9a1e8c17b8e45affd635427f05e2d1e92111d98287b3bd1393d32a65b3", size = 65783, upload-time = "2026-04-14T01:48:00.03Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/e6/07/b9e6ba4afcce41f70a0f149cd7bc299a0babaf9c62a57d7cb291679a6cc7/pydantic_evals-1.77.0-py3-none-any.whl", hash = "sha256:2b536081e36d70826da216a3a6df8d84e3ac1982fc3b8078c123fd539ce6bfb6", size = 77740, upload-time = "2026-04-03T02:16:48.681Z" },
{ url = "https://files.pythonhosted.org/packages/8e/8e/9bee73a780b769c1cb0fa4ad88f22f0dc27d2231358e0ad727b0a9c2f036/pydantic_evals-1.81.0-py3-none-any.whl", hash = "sha256:49516ddadd8064365684cdb4bdfda2e5f1d9b786bebae239d84445727d1e3f26", size = 77710, upload-time = "2026-04-14T01:47:53.166Z" },
]
[[package]]
name = "pydantic-graph"
version = "1.77.0"
version = "1.81.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "httpx" },
@ -3780,9 +3780,9 @@ dependencies = [
{ name = "pydantic" },
{ name = "typing-inspection" },
]
sdist = { url = "https://files.pythonhosted.org/packages/60/40/a8b8e256bb90e4e284b35cc1c5e1a8e2724fa88ad89b7eac958fbf85852b/pydantic_graph-1.77.0.tar.gz", hash = "sha256:ba75dbdf221cd7e366e5c5d250f4d9f3138e05400ea52d3f36330772d989deee", size = 58689, upload-time = "2026-04-03T02:16:56.625Z" }
sdist = { url = "https://files.pythonhosted.org/packages/bd/69/c38ea1c4c8b9789ce1777b408b9fe0f889638415cd97f53a9f95ffbbb044/pydantic_graph-1.81.0.tar.gz", hash = "sha256:721b33324dc25b2ce5956fed8a362e8c558163b45d39e9f83e8ea7b5d44743c0", size = 59240, upload-time = "2026-04-14T01:48:00.99Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/2d/09/3c9c3aba8031adbd21d1833e8e4edd749697e50a88fa9bdab641874abe4f/pydantic_graph-1.77.0-py3-none-any.whl", hash = "sha256:063803e87aec901919c2073ccf3fdd6e4fff84e8b05dbfbe8a6c1af63dd12c05", size = 72503, upload-time = "2026-04-03T02:16:49.97Z" },
{ url = "https://files.pythonhosted.org/packages/4c/bd/9b0561bea26a9918c819b391c3de08cdaa9ef99bab9b4fba8f557df70503/pydantic_graph-1.81.0-py3-none-any.whl", hash = "sha256:9f6256612323d9708b2a3a140db3a3e8ee2fe30c3f1befa897ea473e82cc0faa", size = 73063, upload-time = "2026-04-14T01:47:54.631Z" },
]
[[package]]