diff --git a/CHANGELOG.md b/CHANGELOG.md index fc815f0c..81e41e43 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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`. diff --git a/docs/mcp.md b/docs/mcp.md index 9fef2c7d..5079ee26 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -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). diff --git a/haiku_rag_slim/haiku/rag/context.py b/haiku_rag_slim/haiku/rag/context.py index 17fb44c5..82fe6ed7 100644 --- a/haiku_rag_slim/haiku/rag/context.py +++ b/haiku_rag_slim/haiku/rag/context.py @@ -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": [], } diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index ebfc533b..b9eb2bb1 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -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")) diff --git a/tests/sandbox/test_sandbox_toc.py b/tests/sandbox/test_sandbox_toc.py index d08b9340..9ac5c464 100644 --- a/tests/sandbox/test_sandbox_toc.py +++ b/tests/sandbox/test_sandbox_toc.py @@ -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.""" diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 3954625a..9a9faee6 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -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