diff --git a/CHANGELOG.md b/CHANGELOG.md index 161d6c44..ff05dc4c 100644 --- a/CHANGELOG.md +++ b/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`. diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 454ba9fa..ac4a7b18 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/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: diff --git a/haiku_rag_slim/haiku/rag/store/models/document_item.py b/haiku_rag_slim/haiku/rag/store/models/document_item.py index 0d38e745..1e2955b0 100644 --- a/haiku_rag_slim/haiku/rag/store/models/document_item.py +++ b/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, ) ) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py index f13e2d58..7dde27b5 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py +++ b/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 ] diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py b/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py index faa8c05f..4af3655d 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py +++ b/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) diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_48_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_48_0.py new file mode 100644 index 00000000..328b7eed --- /dev/null +++ b/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", +) diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index b05f4350..eb419e11 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/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" } diff --git a/pyproject.toml b/pyproject.toml index 642917e6..f650e0b3 100644 --- a/pyproject.toml +++ b/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"] -cross-encoder = ["haiku.rag-slim[cross-encoder]==0.47.0"] +s3 = ["haiku.rag-slim[s3]==0.48.0"] +cross-encoder = ["haiku.rag-slim[cross-encoder]==0.48.0"] [build-system] requires = ["hatchling"] diff --git a/tests/store/test_document_items.py b/tests/store/test_document_items.py index ffe4bd3c..b6741834 100644 --- a/tests/store/test_document_items.py +++ b/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) diff --git a/tests/store/test_v0_48_0_migration.py b/tests/store/test_v0_48_0_migration.py new file mode 100644 index 00000000..bbd9a2d2 --- /dev/null +++ b/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 diff --git a/uv.lock b/uv.lock index 7197a09b..fdc587e7 100644 --- a/uv.lock +++ b/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" },