Make toc.json item_range the line slice it documents

build_toc stored item positions, while the sandbox instructions describe
item_range as a slice into items.jsonl; a gap in positions pulled the next
heading into a section. Ranges are now indices into the position-ordered
items, and get_document_section slices by index too. docs/mcp.md names the
tools that take sources.
This commit is contained in:
Yiorgis Gozadinos 2026-09-07 10:27:02 +03:00
parent b374d5eb83
commit 110adf23ee
No known key found for this signature in database
6 changed files with 97 additions and 28 deletions

View file

@ -55,6 +55,11 @@
`search_documents`, `search_documents_by_image` and `execute_code`; `source`
on `get_document`; an unknown name is a tool error. `DocumentInfo.source`.
### Fixed
- `toc.json` `item_range` in the analysis sandbox is a line slice into
`items.jsonl`, as documented; it held item positions.
### Removed
- MCP tools `ask_question` and `analyze`.

View file

@ -32,8 +32,8 @@ The server opens the database read-only. Ingestion goes through the CLI
With several databases in `lancedb.databases`, the server covers all of
them, as `haiku-rag search` does. Results, documents and citations name
theirs in `source`. `sources` on the search and question tools restricts a
call to a subset; `source` on `get_document` names the database holding the
theirs in `source`. `sources` on `search_documents`, `search_documents_by_image`
and `execute_code` restricts a call to a subset; `source` on `get_document` names the database holding the
document. A name the server does not cover is an error.
`haiku-rag --db-name NAME mcp` serves one. See
[Multiple Databases](configuration/storage.md#multiple-databases).

View file

@ -502,10 +502,12 @@ def build_toc(
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.
``item_range = [start, end_exclusive]`` indexes the position-ordered item
list, which is the line numbering of the sandbox's ``items.jsonl``: ``start``
is the header's index and ``end_exclusive`` the index of the next header
whose level is the same or shallower (the next sibling or ancestor that
ends this section), or the item count if no such header exists. Indices,
not positions: positions may have gaps.
``chunk_ids`` aggregates the chunks covered by all items in the section's
``item_range`` (deduped, order preserved). Pass directly to ``cite()`` to
@ -520,33 +522,30 @@ def build_toc(
# but the end_exclusive lookahead below silently miscomputes section
# boundaries if it's not — better to sort once than trust the caller.
items = sorted(items, key=lambda i: i.position)
headers: list[DocumentItem] = [
i for i in items if i.label == "section_header" and i.heading_level > 0
header_indices = [
idx
for idx, i in enumerate(items)
if i.label == "section_header" and i.heading_level > 0
]
if not headers:
if not header_indices:
return []
total = max((i.position for i in items), default=-1) + 1
items_by_position: dict[int, DocumentItem] = {i.position: i for i in items}
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
for n, idx in enumerate(header_indices):
end = len(items)
for later in header_indices[n + 1 :]:
if items[later].heading_level <= items[idx].heading_level:
end = later
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):
for idx, end in zip(header_indices, ends, strict=True):
h = items[idx]
seen: set[str] = set()
chunk_ids: list[str] = []
for pos in range(h.position, end):
item = items_by_position.get(pos)
if item is None:
continue
for item in items[idx:end]:
for cid in chunk_index.get(item.self_ref, []):
if cid not in seen:
seen.add(cid)
@ -556,7 +555,7 @@ def build_toc(
"level": h.heading_level,
"title": h.text,
"page_numbers": list(h.page_numbers),
"item_range": [h.position, end],
"item_range": [idx, end],
"chunk_ids": chunk_ids,
"children": [],
}

View file

@ -409,15 +409,12 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
if node is None:
raise ToolError(f"No section {section_id!r} in document {document_id!r}")
start, end = node["item_range"]
ordered = sorted(items, key=lambda item: item.position)
return DocumentSection(
id=node["self_ref"],
title=node["title"],
page_numbers=node["page_numbers"],
content="\n\n".join(
item.text
for item in items
if start <= item.position < end and item.text
),
content="\n\n".join(item.text for item in ordered[start:end] if item.text),
)
@mcp.tool(annotations=_read_only("List documents"))

View file

@ -467,6 +467,34 @@ class TestVfsReadPaths:
await _read_vfs_text(sandbox, f"/documents/{doc_id}/chunks.jsonl") == first
)
async def test_item_range_is_a_line_slice_into_items_jsonl(self, temp_db_path):
"""`item_range` indexes lines of items.jsonl, as documented, not item
positions: a gap in positions must not pull the next heading into a
section."""
async with HaikuRAG(temp_db_path, create=True) as client:
doc_id = await _empty_doc(client, uri="test://slice", title="Slice")
items = [
_header(doc_id, 0, 1, "Intro"),
_para(doc_id, 1),
_header(doc_id, 3, 1, "Methods"),
_para(doc_id, 4),
]
await client.document_item_repository.create_items(doc_id, items)
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
toc = await _read_toc(sandbox, doc_id)
raw = await _read_vfs_text(sandbox, f"/documents/{doc_id}/items.jsonl")
lines = raw.split("\n")
intro, methods = toc["tree"]
assert intro["item_range"] == [0, 2]
assert methods["item_range"] == [2, 4]
start, end = intro["item_range"]
assert [json.loads(line)["self_ref"] for line in lines[start:end]] == [
"#/texts/0",
"#/texts/1",
]
async def test_toc_skips_gaps_in_item_positions(self, temp_db_path):
"""Positions need not be contiguous — a heading's span may cover
positions that carry no item."""

View file

@ -379,6 +379,46 @@ class TestMCPDocumentNavigation:
assert "para7" in intro.content
assert "Methods" not in intro.content
@pytest.mark.asyncio
async def test_a_section_stops_at_the_next_heading_across_a_position_gap(
self, temp_db_path
):
from haiku.rag.store.models.document import Document as DocumentModel
from haiku.rag.store.models.document_item import DocumentItem
def item(pos, label, text, level=0):
return DocumentItem(
document_id="",
position=pos,
self_ref=f"#/texts/{pos}",
label=label,
text=text,
heading_level=level,
)
async with HaikuRAG(temp_db_path, create=True) as rag:
doc = await rag.document_repository.create(
DocumentModel(content="x", uri="test://gapped", title="Gapped")
)
items = [
item(0, "section_header", "Intro", 1),
item(1, "paragraph", "para1"),
item(3, "section_header", "Methods", 1),
item(4, "paragraph", "para4"),
]
for i in items:
i.document_id = doc.id
await rag.document_item_repository.create_items(doc.id, items)
section = await _call(
create_mcp_server(temp_db_path),
"get_document_section",
document_id=doc.id,
section_id="#/texts/0",
)
assert section.structured_content["content"] == "Intro\n\npara1"
@pytest.mark.asyncio
async def test_an_unknown_section_or_document_is_an_error(self, outlined_db):
db, doc_id = outlined_db