From e8c61ad0c603bede6417dbbc14b6c36bc3a3d5fb Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 19 May 2026 14:20:11 +0300 Subject: [PATCH] items.jsonl exposes chunk_ids per row; drops position and tree_depth --- haiku_rag_slim/haiku/rag/sandbox/sandbox.py | 32 +++++++++++++------ .../haiku/rag/skills/rag-analysis/SKILL.md | 29 ++++++++++------- .../haiku/rag/store/repositories/chunk.py | 32 +++++++++++++++++++ tests/sandbox/test_sandbox.py | 4 +-- tests/sandbox/test_sandbox_toc.py | 21 ++++++++---- 5 files changed, 89 insertions(+), 29 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py index f42753ce..f871bb78 100644 --- a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py +++ b/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, ) 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 13cac3f7..aaba38b0 100644 --- a/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 96a8bbf2..7044066d 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/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 ( diff --git a/tests/sandbox/test_sandbox.py b/tests/sandbox/test_sandbox.py index 2e8cba91..65272d41 100644 --- a/tests/sandbox/test_sandbox.py +++ b/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 diff --git a/tests/sandbox/test_sandbox_toc.py b/tests/sandbox/test_sandbox_toc.py index c68d17e5..a9b0db5f 100644 --- a/tests/sandbox/test_sandbox_toc.py +++ b/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