From 68388032be7a3d43b74e629805e6bb5382f3437c Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 18 May 2026 11:49:48 +0300 Subject: [PATCH] consolidate #/pictures/ prefix, tighten CHANGELOG, log migration exc_info --- CHANGELOG.md | 12 ++++++------ haiku_rag_slim/haiku/rag/agents/analysis/prompts.py | 3 +-- haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py | 4 ++-- haiku_rag_slim/haiku/rag/agents/research/models.py | 6 +++++- haiku_rag_slim/haiku/rag/client/search.py | 11 +++++++---- .../haiku/rag/store/models/document_item.py | 3 +++ haiku_rag_slim/haiku/rag/store/upgrades/v0_48_0.py | 4 +++- 7 files changed, 27 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7c64cb3c..6e1294a1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,14 +3,14 @@ ### 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. -- **`picture_refs` in sandbox `search()` results.** Each search dict in the analysis sandbox now exposes a `picture_refs` key (subset of `doc_item_refs` whose label is `picture`) so the agent can spot picture chunks without zipping refs + labels manually. -- **Citations carry `picture_refs` for inline picture rendering.** `Citation` gains a `picture_refs: list[str]` field; `resolve_citations` derives it from the originating `SearchResult.doc_item_refs` + `labels`. The chat TUI's `CitationWidget` mounts a `textual_image.widget.Image` per picture self_ref inside the existing collapsible — picture citations render the actual figure alongside their text content. The skill's search tool already auto-attaches `BinaryContent` for picture chunks under vision-capable models, so the driving model sees the figure during reasoning; the citation pipeline now also surfaces it in the user's UI. `visualize_chunk` continues to work uniformly because every cited chunk has a `chunk_id`. +- `heading_level` and `tree_depth` on `DocumentItem`, populated by `extract_items` and persisted on `document_items`. 0.48.0 migration backfills existing rows from each doc's docling structure blob. +- `toc.json` in the analysis sandbox VFS at `/documents/{id}/toc.json`. Nested tree on HTML/markdown sources, flat sibling list on PDFs. `items.jsonl` rows now include `heading_level` and `tree_depth`. +- `picture_refs` on sandbox `search()` result dicts and on `Citation` (subset of `doc_item_refs` starting with `#/pictures/`). +- Chat TUI renders picture citations inline via `textual_image.widget.Image` inside the existing `CitationWidget`. -### Changed +### Removed -- **`llm()` removed from the analysis sandbox.** The function was a thin wrapper that spun up an ad-hoc pydantic-ai `Agent` per call. The driving agent already is the LLM — there's no need for a sandbox-internal one. Sandbox external functions are now `search` and `list_documents`. +- `llm()` from the analysis sandbox. Sandbox externals are now `search` and `list_documents` only. ### Changed diff --git a/haiku_rag_slim/haiku/rag/agents/analysis/prompts.py b/haiku_rag_slim/haiku/rag/agents/analysis/prompts.py index 7277d6b2..f962460b 100644 --- a/haiku_rag_slim/haiku/rag/agents/analysis/prompts.py +++ b/haiku_rag_slim/haiku/rag/agents/analysis/prompts.py @@ -147,8 +147,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. +3. **Navigate Structure**: Use `items.jsonl` to find tables, section headers, or specific elements by label and page number (tables are pre-rendered as markdown). 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. **Cite picture chunks for figure-driven questions**: When a question is about a figure or diagram, find the picture chunk (search results with non-empty `picture_refs`) and cite its chunk_id. The driving model already sees figures from search hits; the citation makes the picture visible in the user's UI as well. diff --git a/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py b/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py index 66821d68..89e0e48d 100644 --- a/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py +++ b/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py @@ -13,7 +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 +from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX, DocumentItem if TYPE_CHECKING: from pathlib import PurePosixPath @@ -145,7 +145,7 @@ class Sandbox: out: list[dict[str, Any]] = [] for r in expanded: picture_refs = [ - ref for ref in r.doc_item_refs if ref.startswith("#/pictures/") + ref for ref in r.doc_item_refs if ref.startswith(PICTURE_REF_PREFIX) ] out.append( { diff --git a/haiku_rag_slim/haiku/rag/agents/research/models.py b/haiku_rag_slim/haiku/rag/agents/research/models.py index 0a8e0a9d..ddf5f873 100644 --- a/haiku_rag_slim/haiku/rag/agents/research/models.py +++ b/haiku_rag_slim/haiku/rag/agents/research/models.py @@ -2,6 +2,8 @@ from typing import TYPE_CHECKING from pydantic import BaseModel, Field +from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX + if TYPE_CHECKING: from haiku.rag.store.models import SearchResult @@ -104,7 +106,9 @@ def resolve_citations( r = by_id.get(chunk_id) if not r: continue - picture_refs = [ref for ref in r.doc_item_refs if ref.startswith("#/pictures/")] + picture_refs = [ + ref for ref in r.doc_item_refs if ref.startswith(PICTURE_REF_PREFIX) + ] citations.append( Citation( document_id=r.document_id or "", diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index 0cfaa9b1..c2f0a10a 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING from haiku.rag.reranking import get_reranker from haiku.rag.store.models.chunk import Chunk, SearchResult, SearchType +from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX if TYPE_CHECKING: from PIL import Image as PILImage @@ -93,7 +94,9 @@ def _dedup_picture_chunks(results: list[SearchResult]) -> list[SearchResult]: seen: dict[tuple[str | None, str], int] = {} keep: list[bool] = [True] * len(results) for i, r in enumerate(results): - if len(r.doc_item_refs) == 1 and r.doc_item_refs[0].startswith("#/pictures/"): + if len(r.doc_item_refs) == 1 and r.doc_item_refs[0].startswith( + PICTURE_REF_PREFIX + ): key = (r.document_id, r.doc_item_refs[0]) prior = seen.get(key) if prior is None: @@ -111,13 +114,13 @@ async def _populate_image_data(client: "HaikuRAG", results: list[SearchResult]) Groups results by document_id and batches one picture-bytes lookup per document so a result set spanning N documents costs N reads, not one per - picture. Only refs starting with ``#/pictures/`` are queried. + picture. Only refs starting with ``PICTURE_REF_PREFIX`` are queried. """ by_doc: dict[str, list[SearchResult]] = {} for r in results: if not r.document_id: continue - if not any(ref.startswith("#/pictures/") for ref in r.doc_item_refs): + if not any(ref.startswith(PICTURE_REF_PREFIX) for ref in r.doc_item_refs): continue by_doc.setdefault(r.document_id, []).append(r) @@ -126,7 +129,7 @@ async def _populate_image_data(client: "HaikuRAG", results: list[SearchResult]) seen: set[str] = set() for r in doc_results: for ref in r.doc_item_refs: - if ref.startswith("#/pictures/") and ref not in seen: + if ref.startswith(PICTURE_REF_PREFIX) and ref not in seen: wanted.append(ref) seen.add(ref) if not wanted: 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 1e2955b0..3daf3aed 100644 --- a/haiku_rag_slim/haiku/rag/store/models/document_item.py +++ b/haiku_rag_slim/haiku/rag/store/models/document_item.py @@ -7,6 +7,9 @@ if TYPE_CHECKING: from docling_core.types.doc.document import DoclingDocument, NodeItem, PictureItem +PICTURE_REF_PREFIX = "#/pictures/" + + class DocumentItem(BaseModel): document_id: str position: int 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 index 328b7eed..84eb5a1d 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/v0_48_0.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_48_0.py @@ -72,7 +72,9 @@ async def _apply_backfill_heading_hierarchy(store: Store) -> None: 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) + logger.warning( + "Failed to re-extract items for %s; skipping", doc_id, exc_info=True + ) skipped += 1 continue