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

expose toc.json + heading_level/tree_depth in analysis VFS

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-05-15 14:32:39 +03:00
parent 1360303edf
commit 2831030f9b
No known key found for this signature in database
5 changed files with 423 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
CHANGELOG.md
View file

@ -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

45
haiku_rag_slim/haiku/rag/agents/analysis/prompts.py
View file

@ -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.

143
haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py
View file

@ -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)

8
haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md
View file

@ -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

252
tests/agents/analysis/test_sandbox_toc.py Normal file
View file

@ -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)