items.jsonl exposes chunk_ids per row; drops position and tree_depth

haiku_rag_slim/haiku/rag/sandbox/sandbox.py

@@ -50,8 +50,9 @@ def _build_toc(items: list["DocumentItem"]) -> list[dict[str, Any]]:
     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.
+    skipped. When all section_headers carry the same level the output is a
+    flat sibling list (see docling-project/docling#2121 for an upstream case
+    where every PDF section_header is emitted at level=1).
     """
     headers: list[DocumentItem] = [
         i for i in items if i.label == "section_header" and i.heading_level > 0
@@ -214,33 +215,44 @@ class Sandbox:
         doc_titles = {doc.id: doc.title for doc in docs if doc.id}
 
         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)``.
+            """Bulk-fetch document items + chunk index once and build both
+            items.jsonl and toc.json views. Returns ``(items_cache, toc_cache)``.
             """
 
-            async def _fetch() -> dict[str, list[DocumentItem]]:
+            async def _fetch() -> tuple[
+                dict[str, list[DocumentItem]], dict[str, dict[str, list[str]]]
+            ]:
                 from haiku.rag.client import HaikuRAG
 
                 async with HaikuRAG(db_path, config=config, read_only=True) as rag:
-                    return await rag.document_item_repository.get_all_items_grouped(
-                        doc_ids
+                    items_grouped = (
+                        await rag.document_item_repository.get_all_items_grouped(
+                            doc_ids
+                        )
                     )
+                    chunk_index = (
+                        await rag.chunk_repository.get_chunk_ids_by_self_ref_grouped(
+                            doc_ids
+                        )
+                    )
+                return items_grouped, chunk_index
 
-            grouped = _run_async(_fetch())
+            grouped, chunk_index = _run_async(_fetch())
 
             items_cache: dict[str, str] = {}
             toc_cache: dict[str, str] = {}
             for did, items in grouped.items():
+                doc_chunk_index = chunk_index.get(did, {})
+
                 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,
+                            "chunk_ids": doc_chunk_index.get(item.self_ref, []),
                         },
                         ensure_ascii=False,
                     )

haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md

@@ -28,9 +28,15 @@ Not supported: class definitions, generators/yield, match statements, decorators
 Search the knowledge base directly (outside code execution). Each result has a `Type:` (paragraph, table, code, list_item, picture). When the Type is `picture`, the corresponding figure may also be attached to the tool response as an image alongside the text — use it directly to answer questions about figures, diagrams, charts, screenshots.
 
 ### cite
-Register the chunk IDs that ground your answer. Call this BEFORE writing your final answer, with the `chunk_id` values from search results (from either the `search` tool or `await search(...)` inside `execute_code`) that support each claim. Every answer that uses search results must be backed by `cite`.
+Register the chunk IDs that ground your answer. Call this BEFORE writing your final answer.
 
-Use chunk_ids exactly as they appear in the search response — copy the full UUID verbatim. Do not abbreviate, paraphrase, or reconstruct chunk_ids from memory; the tool matches them as opaque strings.
+Chunk IDs come from two places:
+- The `chunk_id` field on `search` / `await search(...)` results
+- The `chunk_ids` field on `items.jsonl` rows (when you ground via direct file reads)
+
+Do NOT cite `self_ref` (`#/texts/N` style refs), `position`, or any other identifier-shaped field. They are not chunk IDs and the tool will reject them. Copy chunk IDs verbatim — they are opaque UUIDs.
+
+Every answer that uses search or file-read evidence must be backed by `cite`.
 
 ## Document Filesystem (inside execute_code)
 
@@ -44,7 +50,7 @@ All documents are mounted as a virtual filesystem at `/documents/`:
     toc.json         # Section tree derived from heading_level
 ```
 
-`{document_id}` is an internal identifier, not the user-facing `uri` (filename, URL, etc.). When you only know a document by its URI or title, use `await list_documents()` to enumerate ids, or read `metadata.json` from each directory and match against `uri` / `title`.
+`{document_id}` is an internal identifier, not the user-facing `uri` (filename, URL, etc.). When you only know a document by its URI or title, use `await list_documents()` to enumerate ids and match against `uri` / `title` — that's a single call to the host. Iterating `/documents/` and reading every `metadata.json` works too but is much slower on portal-scale corpora.
 
 ### Reading files
 Always use `Path.read_text()` — do NOT use `open()` or `with` statements (they are not supported).
@@ -75,20 +81,21 @@ Document metadata: `id`, `title`, `uri`, `created_at`.
 Full text content. Use for regex or keyword search across a whole document.
 
 ### items.jsonl
-Structured document items. Each line is a JSON object with:
-- `position`: sequential position in the document
-- `self_ref`: item reference (e.g. "#/texts/5", "#/tables/0")
-- `label`: item type — "section_header", "text", "table", "list_item", "caption", "formula", "picture", "code", "footnote"
+Structured document items. One JSON object per line. The row's **line index** is the item's position — `item_range` values in `toc.json` are line-slice bounds into this file.
+
+Each row carries:
+- `self_ref`: item reference (e.g. `"#/texts/5"`, `"#/tables/0"`) — used to cross-reference with `doc_item_refs` from search results
+- `label`: item type — one of `"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.
+- `chunk_ids`: chunks that contain this item — pass to `cite()` to ground an answer that read this item directly
+- `heading_level`: H-level for `section_header` rows; `0` on non-header rows
 
 ### 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.
+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}`. `item_range` is a line slice into `items.jsonl` — `items[start:end]`. `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. 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.
+Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) that correspond to `self_ref` values in `items.jsonl`. To find which section a hit lives in: locate the item by `self_ref`, take its line index, and walk `toc.json` to find the deepest node whose `item_range` contains that index.
 
 ## Strategy
 

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

@@ -351,6 +351,38 @@ class ChunkRepository:
         chunks.sort(key=lambda c: c.order)
         return chunks
 
+    async def get_chunk_ids_by_self_ref_grouped(
+        self, document_ids: list[str]
+    ) -> dict[str, dict[str, list[str]]]:
+        """For each document, build a self_ref → [chunk_id, ...] index.
+
+        One query across all requested documents. The map lets items.jsonl
+        rows expose which chunks contain them, so callers can bridge from an
+        item to a `cite`-acceptable chunk_id without a separate search.
+        """
+        from haiku.rag.utils import escape_sql_string
+
+        if not document_ids:
+            return {}
+
+        safe_ids = ", ".join(f"'{escape_sql_string(did)}'" for did in document_ids)
+        rows = await (
+            self.store.chunks_table.query()
+            .select(["id", "document_id", "metadata"])
+            .where(f"document_id IN ({safe_ids})")
+            .to_list()
+        )
+
+        index: dict[str, dict[str, list[str]]] = {}
+        for row in rows:
+            did = row["document_id"]
+            md = json.loads(row.get("metadata") or "{}")
+            refs = md.get("doc_item_refs") or []
+            doc_index = index.setdefault(did, {})
+            for ref in refs:
+                doc_index.setdefault(ref, []).append(row["id"])
+        return index
+
     async def count_by_document_id(self, document_id: str) -> int:
         """Count the number of chunks for a specific document."""
         df = await (

tests/sandbox/test_sandbox.py

@@ -372,11 +372,11 @@ class TestSandboxVFS:
                 "lines = text.strip().split('\\n')\n"
                 "print(len(lines) > 0)\n"
                 "item = json.loads(lines[0])\n"
-                "print('position' in item)\n"
                 "print('self_ref' in item)\n"
                 "print('label' in item)\n"
                 "print('text' in item)\n"
-                "print('page_numbers' in item)"
+                "print('page_numbers' in item)\n"
+                "print('chunk_ids' in item)"
             )
             assert result.success
             assert result.stdout.count("True") == 6

tests/sandbox/test_sandbox_toc.py

@@ -226,9 +226,11 @@ class TestTocCaching:
 
 @pytest.mark.asyncio
 class TestItemsJsonlSurfacesNewFields:
-    """items.jsonl rows expose heading_level and tree_depth."""
+    """items.jsonl row shape: heading_level is always present (0 on non-headers);
+    chunk_ids surfaces each item's containing chunks; position and tree_depth
+    are not exposed."""
 
-    async def test_jsonl_contains_heading_level_and_tree_depth(self, temp_db_path):
+    async def test_jsonl_row_shape(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 = [
@@ -243,10 +245,17 @@ class TestItemsJsonlSurfacesNewFields:
 
         assert len(rows) == 3
         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
         for r in rows:
-            assert {"position", "self_ref", "label", "text", "page_numbers"} <= set(r)
+            expected = {
+                "self_ref",
+                "label",
+                "text",
+                "page_numbers",
+                "heading_level",
+                "chunk_ids",
+            }
+            assert expected <= set(r)
+            assert "position" not in r
+            assert "tree_depth" not in r