persist docling heading hierarchy on document_items

CHANGELOG.md

@@ -1,6 +1,10 @@
 # Changelog
 ## [Unreleased]
 
+### Added
+
+- **`heading_level` and `tree_depth` on `DocumentItem`.** `extract_items` now captures docling's `SectionHeaderItem.level` (H1–H6 for headers, `0` elsewhere) and the traversal depth from `iterate_items()` for every item, persisting both in the `document_items` table. Foundations for tree-based document navigation in the analysis sandbox. The 0.48.0 migration adds the columns to existing DBs and backfills them from each doc's docling blob.
+
 ### Changed
 
 - Bump `docling>=2.93.0` and `docling-core>=2.75.0`.

haiku_rag_slim/haiku/rag/store/engine.py

@@ -141,6 +141,8 @@ class DocumentItemRecord(LanceModel):
     text: str = Field(default="")
     page_numbers: str = Field(default="[]")
     picture_data: bytes | None = None
+    heading_level: int = Field(default=0)
+    tree_depth: int = Field(default=0)
 
 
 def get_document_items_arrow_schema() -> pa.Schema:

haiku_rag_slim/haiku/rag/store/models/document_item.py

@@ -15,6 +15,8 @@ class DocumentItem(BaseModel):
     text: str = ""
     page_numbers: list[int] = []
     picture_data: bytes | None = None
+    heading_level: int = 0
+    tree_depth: int = 0
 
 
 def _picture_description_text(item: "PictureItem") -> str | None:
@@ -110,12 +112,12 @@ def extract_items(
     fall-back lookup against ``existing_picture_data`` (keyed by ``self_ref``)
     preserves the bytes that were captured at original ingest time.
     """
-    from docling_core.types.doc.document import PictureItem
+    from docling_core.types.doc.document import PictureItem, SectionHeaderItem
 
     existing = existing_picture_data or {}
     items: list[DocumentItem] = []
 
-    for position, (item, _level) in enumerate(docling_doc.iterate_items()):
+    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 "")
 
@@ -134,6 +136,8 @@ def extract_items(
             if picture_data is None:
                 picture_data = existing.get(item.self_ref)
 
+        heading_level = item.level if isinstance(item, SectionHeaderItem) else 0
+
         items.append(
             DocumentItem(
                 document_id=document_id,
@@ -143,6 +147,8 @@ def extract_items(
                 text=text,
                 page_numbers=sorted(page_numbers),
                 picture_data=picture_data,
+                heading_level=heading_level,
+                tree_depth=level,
             )
         )
 

haiku_rag_slim/haiku/rag/store/repositories/document_item.py

@@ -14,6 +14,8 @@ _METADATA_COLUMNS = [
     "label",
     "text",
     "page_numbers",
+    "heading_level",
+    "tree_depth",
 ]
 
 
@@ -31,6 +33,8 @@ class DocumentItemRepository:
             label=row.get("label", ""),
             text=row.get("text", ""),
             page_numbers=json.loads(row.get("page_numbers", "[]")),
+            heading_level=row.get("heading_level", 0) or 0,
+            tree_depth=row.get("tree_depth", 0) or 0,
         )
 
     async def create_items(self, document_id: str, items: list[DocumentItem]) -> None:
@@ -48,6 +52,8 @@ class DocumentItemRepository:
                 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
         ]

haiku_rag_slim/haiku/rag/store/upgrades/__init__.py

@@ -87,6 +87,9 @@ from haiku.rag.store.upgrades.v0_40_0 import (
 from haiku.rag.store.upgrades.v0_45_0 import (
     upgrade_extract_picture_bytes as upgrade_0_45_0_extract_picture_bytes,
 )
+from haiku.rag.store.upgrades.v0_48_0 import (
+    upgrade_backfill_heading_hierarchy as upgrade_0_48_0_heading_hierarchy,
+)
 
 upgrades.append(upgrade_0_20_0_docling)
 upgrades.append(upgrade_0_23_1_contextualize)
@@ -94,3 +97,4 @@ upgrades.append(upgrade_0_25_0_compress)
 upgrades.append(upgrade_0_38_0_split_pages)
 upgrades.append(upgrade_0_40_0_document_items)
 upgrades.append(upgrade_0_45_0_extract_picture_bytes)
+upgrades.append(upgrade_0_48_0_heading_hierarchy)

haiku_rag_slim/haiku/rag/store/upgrades/v0_48_0.py

@@ -0,0 +1,151 @@
+import logging
+
+import pyarrow as pa
+
+from haiku.rag.store.engine import DocumentItemRecord, Store
+from haiku.rag.store.upgrades import Upgrade
+from haiku.rag.utils import escape_sql_string
+
+logger = logging.getLogger(__name__)
+
+PROGRESS_INTERVAL = 10
+
+
+async def _ensure_columns(store: Store) -> None:
+    """Add heading_level + tree_depth (int64) columns if missing. Idempotent."""
+    arrow_schema = await store.document_items_table.schema()
+    existing = {f.name for f in arrow_schema}
+    new_fields = []
+    if "heading_level" not in existing:
+        new_fields.append(pa.field("heading_level", pa.int64()))
+    if "tree_depth" not in existing:
+        new_fields.append(pa.field("tree_depth", pa.int64()))
+    if new_fields:
+        await store.document_items_table.add_columns(pa.schema(new_fields))
+
+
+async def _apply_backfill_heading_hierarchy(store: Store) -> None:
+    """Add heading_level + tree_depth columns and backfill from docling structure.
+
+    For each document with a docling_document blob, decompress it, re-run
+    extract_items, and update the existing items rows with the two new ints.
+    Documents without docling (plain-text adds) are skipped; their rows keep
+    the column-add default (NULL, materialised as 0 by the repository).
+
+    Idempotent: re-running on migrated rows yields identical values.
+    """
+    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
+
+    await _ensure_columns(store)
+
+    ids = (await store.documents_table.query().select(["id"]).to_arrow()).to_pylist()
+    ids = [row["id"] for row in ids]
+    total = len(ids)
+    if not total:
+        return
+
+    logger.info("Backfilling heading_level + tree_depth across %d documents", total)
+    backfilled = 0
+    skipped = 0
+
+    for idx, doc_id in enumerate(ids, 1):
+        safe_id = escape_sql_string(doc_id)
+        rows = await (
+            store.documents_table.query()
+            .select(["id", "docling_document"])
+            .where(f"id = '{safe_id}'")
+            .limit(1)
+            .to_list()
+        )
+        if not rows:
+            skipped += 1
+            continue
+        blob = rows[0].get("docling_document")
+        if not blob or not isinstance(blob, bytes):
+            skipped += 1
+            continue
+
+        try:
+            docling_doc = DoclingDocument.model_validate_json(decompress_json(blob))
+            fresh_items = extract_items(doc_id, docling_doc)
+        except Exception:
+            logger.warning("Failed to re-extract items for %s; skipping", doc_id)
+            skipped += 1
+            continue
+
+        if not fresh_items:
+            skipped += 1
+            continue
+
+        existing_rows = await (
+            store.document_items_table.query()
+            .where(f"document_id = '{safe_id}'")
+            .to_list()
+        )
+        existing_by_ref = {r["self_ref"]: r for r in existing_rows}
+
+        # If the stored rows and a fresh extract disagree on count, the
+        # docling parser version has drifted. Skip and let `rebuild` reconcile.
+        if len(existing_rows) != len(fresh_items):
+            logger.warning(
+                "Item count drift for %s (stored=%d, fresh=%d); skipping",
+                doc_id,
+                len(existing_rows),
+                len(fresh_items),
+            )
+            skipped += 1
+            continue
+
+        records: list[DocumentItemRecord] = []
+        for item in fresh_items:
+            row = existing_by_ref.get(item.self_ref)
+            if row is None:
+                continue
+            records.append(
+                DocumentItemRecord(
+                    document_id=doc_id,
+                    position=row["position"],
+                    self_ref=item.self_ref,
+                    label=row.get("label", ""),
+                    text=row.get("text", ""),
+                    page_numbers=row.get("page_numbers", "[]"),
+                    picture_data=row.get("picture_data"),
+                    heading_level=item.heading_level,
+                    tree_depth=item.tree_depth,
+                )
+            )
+
+        if records:
+            await (
+                store.document_items_table.merge_insert(["document_id", "self_ref"])
+                .when_matched_update_all()
+                .when_not_matched_insert_all()
+                .execute(records)
+            )
+            backfilled += 1
+
+        if idx % PROGRESS_INTERVAL == 0 or idx == total:
+            logger.info(
+                "Progress: %d/%d (%d backfilled, %d skipped)",
+                idx,
+                total,
+                backfilled,
+                skipped,
+            )
+
+    logger.info(
+        "Backfill complete: %d backfilled, %d skipped of %d",
+        backfilled,
+        skipped,
+        total,
+    )
+
+
+upgrade_backfill_heading_hierarchy = Upgrade(
+    version="0.48.0",
+    apply=_apply_backfill_heading_hierarchy,
+    description="Backfill heading_level + tree_depth on document_items",
+)

haiku_rag_slim/pyproject.toml

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

pyproject.toml

@@ -2,7 +2,7 @@
 
 name = "haiku.rag"
 description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling"
-version = "0.47.0"
+version = "0.48.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,cross-encoder]==0.47.0",
+    "haiku.rag-slim[docling,voyageai,mxbai,cohere,zeroentropy,tui,cross-encoder]==0.48.0",
 ]
 
 [project.scripts]
@@ -38,8 +38,8 @@ haiku-rag = "haiku.rag.cli:cli"
 
 [project.optional-dependencies]
 tui = ["textual>=8.2.4"]
-s3 = ["haiku.rag-slim[s3]==0.47.0"]
+s3 = ["haiku.rag-slim[s3]==0.48.0"]
-cross-encoder = ["haiku.rag-slim[cross-encoder]==0.47.0"]
+cross-encoder = ["haiku.rag-slim[cross-encoder]==0.48.0"]
 
 [build-system]
 requires = ["hatchling"]

tests/store/test_document_items.py

@@ -29,6 +29,23 @@ def _make_docling_doc():
     return doc
 
 
+def _make_docling_doc_with_levels():
+    """DoclingDocument with explicit multi-level headings for hierarchy tests."""
+    from docling_core.types.doc.document import DoclingDocument
+    from docling_core.types.doc.labels import DocItemLabel
+
+    doc = DoclingDocument(name="leveled")
+    doc.add_heading(text="Introduction", level=1)
+    doc.add_text(label=DocItemLabel.PARAGRAPH, text="Opening paragraph.")
+    doc.add_heading(text="Background", level=2)
+    doc.add_text(label=DocItemLabel.PARAGRAPH, text="Background paragraph.")
+    doc.add_heading(text="Prior Work", level=3)
+    doc.add_text(label=DocItemLabel.PARAGRAPH, text="Prior-work paragraph.")
+    doc.add_heading(text="Methods", level=1)
+    doc.add_text(label=DocItemLabel.PARAGRAPH, text="Methods paragraph.")
+    return doc
+
+
 class TestExtractItems:
     def test_extracts_all_items(self):
         doc = _make_docling_doc()
@@ -71,6 +88,39 @@ class TestExtractItems:
         assert isinstance(table_item.text, str)
 
 
+class TestExtractItemsHierarchy:
+    """heading_level and tree_depth populated from docling structure."""
+
+    def test_heading_level_on_section_headers(self):
+        doc = _make_docling_doc_with_levels()
+        items = extract_items("doc-1", doc)
+        headers = [i for i in items if i.label == "section_header"]
+        assert [h.heading_level for h in headers] == [1, 2, 3, 1]
+
+    def test_heading_level_zero_on_non_headers(self):
+        doc = _make_docling_doc_with_levels()
+        items = extract_items("doc-1", doc)
+        non_headers = [i for i in items if i.label != "section_header"]
+        assert non_headers
+        assert all(i.heading_level == 0 for i in non_headers)
+
+    def test_tree_depth_set_for_all_items(self):
+        doc = _make_docling_doc_with_levels()
+        items = extract_items("doc-1", doc)
+        assert all(i.tree_depth > 0 for i in items)
+
+    def test_plain_doc_has_zero_heading_level(self):
+        from docling_core.types.doc.document import DoclingDocument
+        from docling_core.types.doc.labels import DocItemLabel
+
+        doc = DoclingDocument(name="plain")
+        doc.add_text(label=DocItemLabel.PARAGRAPH, text="One.")
+        doc.add_text(label=DocItemLabel.PARAGRAPH, text="Two.")
+        items = extract_items("doc-1", doc)
+        assert items
+        assert all(i.heading_level == 0 for i in items)
+
+
 class TestExtractItemText:
     def test_text_item(self):
         from docling_core.types.doc.document import DoclingDocument
@@ -177,6 +227,63 @@ class TestDocumentItemRepository:
             assert await repo.get_item_count("doc-1") == 0
             assert await repo.get_item_count("doc-2") == 5
 
+    async def test_round_trip_preserves_heading_level_and_tree_depth(
+        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=0,
+                    self_ref="#/texts/0",
+                    label="section_header",
+                    text="Intro",
+                    heading_level=1,
+                    tree_depth=1,
+                ),
+                DocumentItem(
+                    document_id="doc-1",
+                    position=1,
+                    self_ref="#/texts/1",
+                    label="section_header",
+                    text="Background",
+                    heading_level=2,
+                    tree_depth=2,
+                ),
+                DocumentItem(
+                    document_id="doc-1",
+                    position=2,
+                    self_ref="#/texts/2",
+                    label="paragraph",
+                    text="A paragraph.",
+                    heading_level=0,
+                    tree_depth=2,
+                ),
+            ]
+            await repo.create_items("doc-1", items)
+
+            got = await repo.get_all_items("doc-1")
+            assert [(i.heading_level, i.tree_depth) for i in got] == [
+                (1, 1),
+                (2, 2),
+                (0, 2),
+            ]
+
+            in_range = await repo.get_items_in_range("doc-1", 0, 2)
+            assert [(i.heading_level, i.tree_depth) for i in in_range] == [
+                (1, 1),
+                (2, 2),
+                (0, 2),
+            ]
+
+            grouped = await repo.get_all_items_grouped(["doc-1"])
+            assert [(i.heading_level, i.tree_depth) for i in grouped["doc-1"]] == [
+                (1, 1),
+                (2, 2),
+                (0, 2),
+            ]
+
     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)

tests/store/test_v0_48_0_migration.py

@@ -0,0 +1,177 @@
+import json
+
+import pytest
+
+from haiku.rag.store.compression import compress_docling_split
+from haiku.rag.store.engine import DocumentItemRecord, DocumentRecord, Store
+
+
+def _docling_with_levels():
+    from docling_core.types.doc.document import DoclingDocument
+    from docling_core.types.doc.labels import DocItemLabel
+
+    doc = DoclingDocument(name="leveled")
+    doc.add_heading(text="Introduction", level=1)
+    doc.add_text(label=DocItemLabel.PARAGRAPH, text="Opening.")
+    doc.add_heading(text="Background", level=2)
+    doc.add_text(label=DocItemLabel.PARAGRAPH, text="Bg paragraph.")
+    doc.add_heading(text="Methods", level=1)
+    return doc
+
+
+@pytest.mark.asyncio
+class TestV0_48_0Migration:
+    """v0.48.0 backfills heading_level + tree_depth on existing items rows."""
+
+    async def test_backfill_populates_levels(self, temp_db_path):
+        docling_doc = _docling_with_levels()
+        structure, pages = compress_docling_split(docling_doc.model_dump_json())
+
+        async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
+            await store.set_haiku_version("0.45.0")
+            await store.documents_table.add(
+                [
+                    DocumentRecord(
+                        id="doc-1",
+                        content="legacy",
+                        docling_document=structure,
+                        docling_pages=pages,
+                        docling_version=docling_doc.version,
+                    )
+                ]
+            )
+            # Pre-migration items: heading_level / tree_depth default to 0.
+            # Match what v0.40.0 produced (label, text, page_numbers, no levels).
+            items = list(docling_doc.iterate_items())
+            await store.document_items_table.add(
+                [
+                    DocumentItemRecord(
+                        document_id="doc-1",
+                        position=pos,
+                        self_ref=item.self_ref,
+                        label=str(getattr(item.label, "value", item.label) or ""),
+                        text=getattr(item, "text", "") or "",
+                        page_numbers="[]",
+                    )
+                    for pos, (item, _depth) in enumerate(items)
+                ]
+            )
+
+        async with Store(temp_db_path, skip_migration_check=True) as store:
+            applied = await store.migrate()
+            assert any("0.48.0" in d for d in applied)
+
+            rows = await (
+                store.document_items_table.query()
+                .where("document_id = 'doc-1'")
+                .to_list()
+            )
+            rows.sort(key=lambda r: r["position"])
+            headers = [r for r in rows if r["label"] == "section_header"]
+            assert [r["heading_level"] for r in headers] == [1, 2, 1]
+            non_headers = [r for r in rows if r["label"] != "section_header"]
+            assert non_headers
+            assert all(r["heading_level"] == 0 for r in non_headers)
+            assert all(r["tree_depth"] > 0 for r in rows)
+
+    async def test_backfill_idempotent(self, temp_db_path):
+        docling_doc = _docling_with_levels()
+        structure, pages = compress_docling_split(docling_doc.model_dump_json())
+
+        async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
+            await store.set_haiku_version("0.45.0")
+            await store.documents_table.add(
+                [
+                    DocumentRecord(
+                        id="doc-1",
+                        content="legacy",
+                        docling_document=structure,
+                        docling_pages=pages,
+                        docling_version=docling_doc.version,
+                    )
+                ]
+            )
+            items = list(docling_doc.iterate_items())
+            await store.document_items_table.add(
+                [
+                    DocumentItemRecord(
+                        document_id="doc-1",
+                        position=pos,
+                        self_ref=item.self_ref,
+                        label=str(getattr(item.label, "value", item.label) or ""),
+                        text=getattr(item, "text", "") or "",
+                        page_numbers="[]",
+                    )
+                    for pos, (item, _depth) in enumerate(items)
+                ]
+            )
+
+        async with Store(temp_db_path, skip_migration_check=True) as store:
+            await store.migrate()
+
+            from haiku.rag.store.upgrades.v0_48_0 import (
+                _apply_backfill_heading_hierarchy,
+            )
+
+            await _apply_backfill_heading_hierarchy(store)
+
+            rows = await (
+                store.document_items_table.query()
+                .where("document_id = 'doc-1'")
+                .to_list()
+            )
+            rows.sort(key=lambda r: r["position"])
+            headers = [r for r in rows if r["label"] == "section_header"]
+            assert [r["heading_level"] for r in headers] == [1, 2, 1]
+
+    async def test_skips_documents_without_docling(self, temp_db_path):
+        async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
+            await store.set_haiku_version("0.45.0")
+            await store.documents_table.add(
+                [DocumentRecord(id="plain", content="no docling")]
+            )
+            await store.document_items_table.add(
+                [
+                    DocumentItemRecord(
+                        document_id="plain",
+                        position=0,
+                        self_ref="#/texts/0",
+                        label="paragraph",
+                        text="Hi",
+                        page_numbers="[]",
+                    )
+                ]
+            )
+
+        async with Store(temp_db_path, skip_migration_check=True) as store:
+            await store.migrate()
+            rows = await (
+                store.document_items_table.query()
+                .where("document_id = 'plain'")
+                .to_list()
+            )
+            # Plain doc — no levels to backfill; defaults stay zero.
+            assert rows[0]["heading_level"] == 0
+            assert rows[0]["tree_depth"] == 0
+
+
+@pytest.mark.asyncio
+class TestV0_48_0FreshSchema:
+    """A newly-created DB has the new columns from the start."""
+
+    async def test_fresh_db_has_heading_level_and_tree_depth(self, temp_db_path):
+        from haiku.rag.client import HaikuRAG
+
+        async with HaikuRAG(temp_db_path, create=True) as rag:
+            schema = await rag.store.document_items_table.schema()
+            names = {f.name for f in schema}
+            assert "heading_level" in names
+            assert "tree_depth" in names
+
+
+# Sanity: ensure module imports without side effects.
+def test_module_imports():
+    from haiku.rag.store.upgrades import v0_48_0  # noqa: F401
+
+    assert hasattr(v0_48_0, "upgrade_backfill_heading_hierarchy")
+    assert json.dumps  # silence unused-import flake on json

uv.lock

@@ -1467,7 +1467,7 @@ wheels = [
 
 [[package]]
 name = "haiku-rag"
-version = "0.47.0"
+version = "0.48.0"
 source = { editable = "." }
 dependencies = [
     { name = "haiku-rag-slim", extra = ["cohere", "cross-encoder", "docling", "mxbai", "tui", "voyageai", "zeroentropy"] },
@@ -1559,7 +1559,7 @@ requires-dist = [
 
 [[package]]
 name = "haiku-rag-slim"
-version = "0.47.0"
+version = "0.48.0"
 source = { editable = "haiku_rag_slim" }
 dependencies = [
     { name = "docling-core" },