diff --git a/CHANGELOG.md b/CHANGELOG.md index ff05dc4c..24a1e171 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### 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. +- **`toc.json` in the analysis sandbox VFS.** Each document mounted under `/documents/{id}/` now exposes a `toc.json` view alongside `metadata.json`, `content.txt`, `items.jsonl`. Nodes carry `{self_ref, level, title, position, page_numbers, item_range, children}`; `item_range = [start, end_exclusive]` over the same `position` ints used in `items.jsonl`, so the agent can slice items by range to read a section. HTML/markdown ingests produce a real nested tree; PDF ingests produce a flat sibling list because docling collapses heading levels on PDFs. `tree: []` when the doc has no section headers. `items.jsonl` now surfaces `heading_level` and `tree_depth` on every row. ### Changed diff --git a/haiku_rag_slim/haiku/rag/agents/analysis/prompts.py b/haiku_rag_slim/haiku/rag/agents/analysis/prompts.py index 53147c9b..a7924eef 100644 --- a/haiku_rag_slim/haiku/rag/agents/analysis/prompts.py +++ b/haiku_rag_slim/haiku/rag/agents/analysis/prompts.py @@ -34,6 +34,7 @@ All documents in the knowledge base are available as files under `/documents/`. metadata.json # {"id", "title", "uri", "created_at"} content.txt # Full document text items.jsonl # Structured document items (one JSON object per line) + toc.json # Section tree (nested or flat depending on source) ``` ### metadata.json @@ -59,6 +60,8 @@ Structured document items as JSONL. Each line is a JSON object with: - `label`: item type — "section_header", "text", "table", "list_item", "caption", "formula", "picture", "code", "footnote", etc. - `text`: rendered content (tables are markdown with `|` columns) - `page_numbers`: list of page numbers where the item appears +- `heading_level`: H-level (1–6) on `section_header` items, `0` otherwise. Often `1` for everything when the source is a PDF (docling can't infer heading hierarchy from PDFs) — see `toc.json` for the derived tree. +- `tree_depth`: DOM nesting depth from docling's structure. Useful for HTML where it varies meaningfully (sidebars, captions, nested lists); near-uniform on PDFs. Use items.jsonl to find tables, section headers, or specific structural elements: ```python @@ -70,6 +73,47 @@ for line in items_text.strip().split(chr(10)): print(f"Table on page {item['page_numbers']}: {item['text'][:100]}") ``` +### toc.json +Per-document section tree derived from `heading_level`. Shape: +```json +{"doc_id": "...", "title": "...", "tree": [ + {"self_ref": "#/texts/0", "level": 1, "title": "Intro", + "position": 0, "page_numbers": [1], "item_range": [0, 18], + "children": [ + {"self_ref": "#/texts/8", "level": 2, "title": "Background", + "position": 8, "page_numbers": [2], "item_range": [8, 13], + "children": []} + ]} +]} +``` +- `item_range = [start, end_exclusive]` over the same `position` ints used in items.jsonl. Slice items.jsonl by this range to read a whole section. +- PDF-derived docs typically produce a flat list of level-1 siblings (docling collapses heading levels). HTML/markdown produce a real nested tree. +- `tree: []` when the doc has no section_headers at all. + +```python +import json +toc = json.loads(Path(f'/documents/{doc_id}/toc.json').read_text()) +items = [json.loads(line) for line in Path(f'/documents/{doc_id}/items.jsonl').read_text().strip().split(chr(10))] + +# Read the contents of one section +def items_in(node): + start, end = node['item_range'] + return [it for it in items if start <= it['position'] < end] + +# From a search hit's doc_item_refs, find the deepest TOC node containing it +def find_containing_section(tree, position): + best = None + def walk(nodes): + nonlocal best + for n in nodes: + s, e = n['item_range'] + if s <= position < e: + best = n + walk(n['children']) + walk(tree) + return best +``` + ## Cross-referencing search results with items Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) that correspond to `self_ref` values in items.jsonl. Use this to navigate from a search hit to the surrounding document structure: @@ -108,6 +152,7 @@ Not supported: most imports (only `json`, `re`, `math`, `pathlib` are available) 1. **Search First**: Start with `search()` to find relevant content. Results include expanded context and `doc_item_refs` for cross-referencing. 2. **Discover Documents**: Use `list_documents()` to see what's in the knowledge base. 3. **Use items.jsonl for Structure**: Find tables, section headers, or specific elements by label and page number. Tables are pre-rendered as markdown. +3b. **Use toc.json for Section Navigation**: When a question is scoped to a section, open `toc.json`, find the matching node, and slice `items.jsonl` by its `item_range` instead of streaming `content.txt`. For PDFs where the tree is flat, the sibling list is still useful as a TOC. 4. **Use content.txt for Full Text**: When you need the complete document text (e.g., for regex across the whole document). 5. **Iterate**: Run code, examine results, refine your approach. Don't try to solve everything in one execution. 6. **Use llm() for Reasoning**: When you have content and need classification, summarization, or extraction, use `llm()` rather than writing complex parsing logic. diff --git a/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py b/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py index 21c63d0d..57d3b2db 100644 --- a/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py +++ b/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py @@ -13,6 +13,7 @@ from pydantic_monty import CallbackFile, MemoryFile, MontyRepl, OSAccess from haiku.rag.agents.analysis.dependencies import AnalysisContext from haiku.rag.config.models import AppConfig from haiku.rag.store.models.chunk import SearchResult +from haiku.rag.store.models.document_item import DocumentItem if TYPE_CHECKING: from pathlib import PurePosixPath @@ -36,6 +37,60 @@ def _run_async(coro: Any) -> Any: return _executor.submit(asyncio.run, coro).result() +def _build_toc(items: list["DocumentItem"]) -> list[dict[str, Any]]: + """Build a nested section tree from items in position order. + + Each ``section_header`` with ``heading_level > 0`` becomes a node. Nesting + follows the explicit levels: a header pops the stack until the top is at + a strictly shallower level, then becomes a child of that top (or a root). + + ``item_range = [position, end_exclusive]`` where ``end_exclusive`` is the + position of the next header whose level is the same or shallower (i.e. + the next sibling or ancestor that ends this section), or the total item + count if no such header exists. + + Items without a section_header label (or with ``heading_level == 0``) are + skipped. PDF-derived corpora where docling collapses every header to level + 1 produce a flat list of sibling nodes. + """ + headers: list[DocumentItem] = [ + i for i in items if i.label == "section_header" and i.heading_level > 0 + ] + if not headers: + return [] + + total = max((i.position for i in items), default=-1) + 1 + + # Pre-compute each header's end position: the next header in document order + # whose heading_level <= this header's level. + ends: list[int] = [] + for idx, h in enumerate(headers): + end = total + for j in range(idx + 1, len(headers)): + if headers[j].heading_level <= h.heading_level: + end = headers[j].position + break + ends.append(end) + + roots: list[dict[str, Any]] = [] + stack: list[tuple[int, dict[str, Any]]] = [] + for h, end in zip(headers, ends, strict=True): + node: dict[str, Any] = { + "self_ref": h.self_ref, + "level": h.heading_level, + "title": h.text, + "position": h.position, + "page_numbers": list(h.page_numbers), + "item_range": [h.position, end], + "children": [], + } + while stack and stack[-1][0] >= h.heading_level: + stack.pop() + (stack[-1][1]["children"] if stack else roots).append(node) + stack.append((h.heading_level, node)) + return roots + + class Sandbox: """Execute code in a sandboxed Python interpreter. @@ -57,6 +112,7 @@ class Sandbox: _context: AnalysisContext _search_results: "list[SearchResult]" _items_cache: dict[str, str] | None + _toc_cache: dict[str, str] | None _repl: MontyRepl | None _vfs: OSAccess | None @@ -71,6 +127,7 @@ class Sandbox: self._context = context self._search_results = [] self._items_cache = None + self._toc_cache = None self._repl = None self._vfs = None @@ -155,50 +212,79 @@ class Sandbox: docs = await rag.list_documents(filter=self._context.filter) doc_ids = [doc.id for doc in docs if doc.id] + doc_titles = {doc.id: doc.title for doc in docs if doc.id} - def _load_items_cache() -> dict[str, str]: - """Bulk-fetch all document items in one query, serialize to JSONL.""" + def _load_caches() -> tuple[dict[str, str], dict[str, str]]: + """Bulk-fetch document items once and build both items.jsonl and + toc.json views from the same result. Returns ``(items_cache, toc_cache)``. + """ - async def _fetch() -> dict[str, str]: + async def _fetch() -> dict[str, list[DocumentItem]]: from haiku.rag.client import HaikuRAG async with HaikuRAG(db_path, config=config, read_only=True) as rag: - grouped = await rag.document_item_repository.get_all_items_grouped( + return await rag.document_item_repository.get_all_items_grouped( doc_ids ) - result: dict[str, str] = {} - for did, items in grouped.items(): - lines = [] - for item in items: - lines.append( - json.dumps( - { - "position": item.position, - "self_ref": item.self_ref, - "label": item.label, - "text": item.text, - "page_numbers": item.page_numbers, - }, - ensure_ascii=False, - ) - ) - result[did] = "\n".join(lines) - return result - return _run_async(_fetch()) + grouped = _run_async(_fetch()) + + items_cache: dict[str, str] = {} + toc_cache: dict[str, str] = {} + for did, items in grouped.items(): + items_cache[did] = "\n".join( + json.dumps( + { + "position": item.position, + "self_ref": item.self_ref, + "label": item.label, + "text": item.text, + "page_numbers": item.page_numbers, + "heading_level": item.heading_level, + "tree_depth": item.tree_depth, + }, + ensure_ascii=False, + ) + for item in items + ) + toc_cache[did] = json.dumps( + { + "doc_id": did, + "title": doc_titles.get(did), + "tree": _build_toc(items), + }, + ensure_ascii=False, + ) + return items_cache, toc_cache sandbox = self + def _ensure_caches() -> None: + if sandbox._items_cache is None or sandbox._toc_cache is None: + items_c, toc_c = _load_caches() + sandbox._items_cache = items_c + sandbox._toc_cache = toc_c + def _make_items_reader( did: str, ) -> Callable[["PurePosixPath"], str]: def read_items(_path: "PurePosixPath") -> str: - if sandbox._items_cache is None: - sandbox._items_cache = _load_items_cache() + _ensure_caches() + assert sandbox._items_cache is not None return sandbox._items_cache.get(did, "") return read_items + def _make_toc_reader( + did: str, + ) -> Callable[["PurePosixPath"], str]: + def read_toc(_path: "PurePosixPath") -> str: + _ensure_caches() + assert sandbox._toc_cache is not None + return sandbox._toc_cache.get(did, "") + + return read_toc + for doc in docs: if not doc.id: continue @@ -247,6 +333,13 @@ class Sandbox: write=_deny_write, ) ) + files.append( + CallbackFile( + f"{doc_dir}/toc.json", + read=_make_toc_reader(doc_id), + write=_deny_write, + ) + ) return OSAccess(files) diff --git a/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md b/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md index b23644e1..453480bf 100644 --- a/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md +++ b/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md @@ -43,6 +43,7 @@ All documents are mounted as a virtual filesystem at `/documents/`: metadata.json # {"id", "title", "uri", "created_at"} content.txt # Full document text items.jsonl # Structured items (one JSON object per line) + toc.json # Section tree derived from heading_level ``` ### Reading files @@ -80,9 +81,14 @@ Structured document items. Each line is a JSON object with: - `label`: item type — "section_header", "text", "table", "list_item", "caption", "formula", "picture", "code", "footnote" - `text`: rendered content (tables are markdown with `|` columns) - `page_numbers`: list of page numbers where the item appears +- `heading_level`: H-level (1–6) for `section_header` items, `0` otherwise. PDFs often collapse to `1` for every header. +- `tree_depth`: DOM nesting depth — varies meaningfully on HTML, near-uniform on PDFs. + +### toc.json +Section tree derived from `heading_level`: `{"doc_id", "title", "tree": [...]}` where each node has `{self_ref, level, title, position, page_numbers, item_range: [start, end_exclusive], children}`. Slice `items.jsonl` by `item_range` to read a section's contents. PDFs typically produce a flat sibling list; HTML/markdown produce a real tree. `tree: []` for docs with no headers. ### Cross-referencing search results with items -Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) that correspond to `self_ref` values in items.jsonl. +Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) that correspond to `self_ref` values in items.jsonl. Resolve each ref to a `position`, then walk `toc.json` to find the deepest node whose `item_range` contains it — that's the section the hit lives in. ## Strategy diff --git a/tests/agents/analysis/test_sandbox_toc.py b/tests/agents/analysis/test_sandbox_toc.py new file mode 100644 index 00000000..9590ef99 --- /dev/null +++ b/tests/agents/analysis/test_sandbox_toc.py @@ -0,0 +1,252 @@ +"""Tests for the per-document toc.json view and the heading_level / tree_depth +fields surfaced in items.jsonl. + +The TOC is derived from `DocumentItem.heading_level` (positive only) in +position order. PDF-style corpora (all section_headers at level 1) get a flat +list of siblings; HTML/markdown corpora with real heading hierarchy get a +nested tree. Items with no section_header at all produce `tree: []`. +""" + +import json +from pathlib import PurePosixPath + +import pytest + +from haiku.rag.agents.analysis.dependencies import AnalysisContext +from haiku.rag.agents.analysis.sandbox import Sandbox +from haiku.rag.client import HaikuRAG +from haiku.rag.config.models import AppConfig +from haiku.rag.store.models.document_item import DocumentItem + + +async def _empty_doc(client, *, uri: str, title: str) -> str: + """Create a Document row and drop the auto-extracted items so the test + controls the items table exactly. Returns the document id.""" + doc = await client.create_document(content="x", uri=uri, title=title) + await client.document_item_repository.delete_by_document_id(doc.id) + return doc.id + + +def _para(doc_id: str, pos: int, depth: int = 1) -> DocumentItem: + return DocumentItem( + document_id=doc_id, + position=pos, + self_ref=f"#/texts/{pos}", + label="paragraph", + text=f"para{pos}", + page_numbers=[1], + tree_depth=depth, + ) + + +def _header( + doc_id: str, pos: int, level: int, text: str, depth: int = 1, page: int = 1 +) -> DocumentItem: + return DocumentItem( + document_id=doc_id, + position=pos, + self_ref=f"#/texts/{pos}", + label="section_header", + text=text, + page_numbers=[page], + heading_level=level, + tree_depth=depth, + ) + + +async def _read_toc(sandbox: Sandbox, doc_id: str) -> dict: + vfs = await sandbox._build_vfs() + raw = vfs.path_read_text(PurePosixPath(f"/documents/{doc_id}/toc.json")) + return json.loads(raw) + + +async def _read_items_jsonl(sandbox: Sandbox, doc_id: str) -> list[dict]: + vfs = await sandbox._build_vfs() + raw = vfs.path_read_text(PurePosixPath(f"/documents/{doc_id}/items.jsonl")) + return [json.loads(line) for line in raw.strip().splitlines()] if raw else [] + + +def _flatten(tree: list[dict]) -> list[dict]: + out = [] + for node in tree: + out.append(node) + out.extend(_flatten(node["children"])) + return out + + +@pytest.mark.asyncio +class TestTocShape: + """toc.json builds a section tree from heading_level + position.""" + + async def test_multilevel_tree(self, temp_db_path): + async with HaikuRAG(temp_db_path, create=True) as client: + doc_id = await _empty_doc(client, uri="test://multi", title="TOC Test Doc") + items = [ + _header(doc_id, 0, 1, "Intro"), + _para(doc_id, 1), + _header(doc_id, 2, 2, "Background"), + _para(doc_id, 3), + _header(doc_id, 4, 3, "Prior Work"), + _para(doc_id, 5), + _header(doc_id, 6, 2, "Approach"), + _para(doc_id, 7), + _header(doc_id, 8, 1, "Methods"), + _para(doc_id, 9), + ] + await client.document_item_repository.create_items(doc_id, items) + + sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext()) + toc = await _read_toc(sandbox, doc_id) + + assert toc["doc_id"] == doc_id + assert toc["title"] == "TOC Test Doc" + tree = toc["tree"] + # Two roots: Intro (children: Background>{Prior Work}, Approach) and Methods + assert [n["title"] for n in tree] == ["Intro", "Methods"] + intro = tree[0] + assert intro["level"] == 1 + assert intro["item_range"] == [0, 8] # ends at "Methods" position + assert [c["title"] for c in intro["children"]] == ["Background", "Approach"] + + background = intro["children"][0] + assert background["level"] == 2 + # Background covers positions 2..5; "Approach" begins at 6 (same-level sibling) + assert background["item_range"] == [2, 6] + assert [c["title"] for c in background["children"]] == ["Prior Work"] + + prior = background["children"][0] + assert prior["level"] == 3 + # Prior Work has no descendants and the next same-or-shallower header is + # "Approach" at level 2, position 6. + assert prior["item_range"] == [4, 6] + assert prior["children"] == [] + + approach = intro["children"][1] + assert approach["item_range"] == [6, 8] + + methods = tree[1] + assert methods["item_range"] == [8, 10] # to end of items + assert methods["children"] == [] + + async def test_flat_pdf_style(self, temp_db_path): + """All section_headers at level 1 (PDF reality) -> flat sibling list.""" + async with HaikuRAG(temp_db_path, create=True) as client: + doc_id = await _empty_doc(client, uri="test://pdf-style", title="Flat PDF") + items = [ + _header(doc_id, 0, 1, "Chapter 1"), + _para(doc_id, 1), + _para(doc_id, 2), + _header(doc_id, 3, 1, "Chapter 2"), + _para(doc_id, 4), + _header(doc_id, 5, 1, "Chapter 3"), + _para(doc_id, 6), + ] + await client.document_item_repository.create_items(doc_id, items) + + sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext()) + toc = await _read_toc(sandbox, doc_id) + + tree = toc["tree"] + assert [n["title"] for n in tree] == ["Chapter 1", "Chapter 2", "Chapter 3"] + assert all(n["level"] == 1 and n["children"] == [] for n in tree) + assert tree[0]["item_range"] == [0, 3] + assert tree[1]["item_range"] == [3, 5] + assert tree[2]["item_range"] == [5, 7] + + async def test_no_headers(self, temp_db_path): + async with HaikuRAG(temp_db_path, create=True) as client: + doc_id = await _empty_doc( + client, uri="test://no-headers", title="No Headers" + ) + await client.document_item_repository.create_items( + doc_id, [_para(doc_id, i) for i in range(5)] + ) + + sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext()) + toc = await _read_toc(sandbox, doc_id) + assert toc["tree"] == [] + + async def test_skip_header_with_zero_level(self, temp_db_path): + """A section_header with heading_level=0 (legacy pre-0.46.0 row) is skipped.""" + async with HaikuRAG(temp_db_path, create=True) as client: + doc_id = await _empty_doc(client, uri="test://zero", title="Zero Level") + items = [ + _header(doc_id, 0, 1, "Real H1"), + _para(doc_id, 1), + _header(doc_id, 2, 0, "Pre-migration ghost"), + _para(doc_id, 3), + ] + await client.document_item_repository.create_items(doc_id, items) + + sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext()) + toc = await _read_toc(sandbox, doc_id) + titles = [n["title"] for n in _flatten(toc["tree"])] + assert titles == ["Real H1"] + assert toc["tree"][0]["item_range"] == [0, 4] + + +@pytest.mark.asyncio +class TestTocCaching: + """toc.json is cached across reads — the items query runs once per sandbox.""" + + async def test_cached_across_reads(self, temp_db_path, monkeypatch): + async with HaikuRAG(temp_db_path, create=True) as client: + doc_id = await _empty_doc(client, uri="test://cache", title="Cache") + await client.document_item_repository.create_items( + doc_id, + [_header(doc_id, 0, 1, "Only"), _para(doc_id, 1), _para(doc_id, 2)], + ) + + sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext()) + + # Patch the repository call that backs both items.jsonl and toc.json so we + # can count how many bulk fetches happen. The sandbox opens its own + # HaikuRAG client(s) lazily; patch the class method. + from haiku.rag.store.repositories.document_item import DocumentItemRepository + + call_count = {"n": 0} + original = DocumentItemRepository.get_all_items_grouped + + async def counting(self, document_ids=None): + call_count["n"] += 1 + return await original(self, document_ids) + + monkeypatch.setattr(DocumentItemRepository, "get_all_items_grouped", counting) + + # First read of either file triggers ONE bulk fetch. + _ = await _read_toc(sandbox, doc_id) + _ = await _read_items_jsonl(sandbox, doc_id) + _ = await _read_toc(sandbox, doc_id) + _ = await _read_items_jsonl(sandbox, doc_id) + + assert call_count["n"] == 1 + + +@pytest.mark.asyncio +class TestItemsJsonlSurfacesNewFields: + """items.jsonl rows expose heading_level and tree_depth.""" + + async def test_jsonl_contains_heading_level_and_tree_depth(self, temp_db_path): + async with HaikuRAG(temp_db_path, create=True) as client: + doc_id = await _empty_doc(client, uri="test://jsonl-fields", title="Fields") + items = [ + _header(doc_id, 0, 1, "H1", depth=2), + _para(doc_id, 1, depth=3), + _header(doc_id, 2, 2, "H2", depth=4), + ] + await client.document_item_repository.create_items(doc_id, items) + + sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext()) + rows = await _read_items_jsonl(sandbox, doc_id) + + assert len(rows) == 3 + # Field presence + values + assert rows[0]["heading_level"] == 1 + assert rows[0]["tree_depth"] == 2 + assert rows[1]["heading_level"] == 0 + assert rows[1]["tree_depth"] == 3 + assert rows[2]["heading_level"] == 2 + assert rows[2]["tree_depth"] == 4 + # Existing fields still present and unchanged + for r in rows: + assert {"position", "self_ref", "label", "text", "page_numbers"} <= set(r)