danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Navigate documents by outline and section from the MCP server

Browse source 
build_toc moves from the sandbox into haiku.rag.context; the sandbox
keeps its toc.json unchanged. get_document_outline returns the heading
tree with page numbers and get_document_section one section's text,
subsections included, both resolved in the database holding the
document. Chunk ids never leave the server. ask_question drops `cite`
and always appends its citations.

Refs #599
This commit is contained in:
Yiorgis Gozadinos 2026-09-04 12:16:00 +03:00
parent 2a6d72171d
commit 15afb97a6e
No known key found for this signature in database
7 changed files with 391 additions and 103 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
CHANGELOG.md
View file

@ -4,6 +4,9 @@
### Added
- MCP tools `get_document_outline` (heading tree with page numbers) and
`get_document_section` (one section's text, subsections included), built
on `document_items`. `build_toc` in `haiku.rag.context`.
- MCP server `instructions`, `version`, and read-only `ToolAnnotations` on
every tool; every parameter carries a description. `filter` on
`search_documents` and `search_documents_by_image`. `DocumentInfo.metadata`.
@ -35,6 +38,7 @@
### Removed
- `cite` on the MCP `ask_question` tool; citations are always appended.
- MCP write tools `add_document_from_file`, `add_document_from_url`,
`add_document_from_text` and `delete_document`. The server opens the
database read-only; ingest with `haiku-rag add`, `add-src`, `delete` or

22
docs/mcp.md
View file

@ -78,8 +78,10 @@ repeating it.
| `search_documents` | always | `query`, `limit`, `include_images`, `filter`, `sources` |
| `search_documents_by_image` | multimodal embedder only | `image_base64`, `limit`, `include_images`, `filter`, `sources` |
| `get_document` | always | `document_id`, `source` |
| `get_document_outline` | always | `document_id`, `source` |
| `get_document_section` | always | `document_id`, `section_id`, `source` |
| `list_documents` | always | `limit`, `offset`, `filter` |
| `ask_question` | always | `question`, `cite`, `images_base64`, `sources` |
| `ask_question` | always | `question`, `images_base64`, `sources` |
| `analyze` | always | `question`, `filter`, `images_base64`, `sources` |
`search_documents` runs hybrid search, vector and full-text, and returns
@ -88,12 +90,15 @@ Rank is the signal. `include_images` attaches picture bytes as base64 PNG under
`image_data`. `search_documents_by_image` embeds the query image and searches
by vector similarity alone.
`get_document` returns a document whole, in reading order. `list_documents`
returns titles, URIs and metadata, which is how a client learns what a filter
can match.
`get_document` returns a document whole, in reading order. For a long one,
`get_document_outline` returns the heading tree with page numbers and
`get_document_section` the text of one section, subsections included; a
node's `id` in the outline is the `section_id`. A document without headings
has an empty outline. `list_documents` returns titles, URIs and metadata,
which is how a client learns what a filter can match.
`ask_question` runs the RAG agent on the server and returns an answer, with
citations when `cite` is set. `analyze` writes and runs Python in a sandbox
`ask_question` runs the RAG agent on the server and returns an answer
followed by its citations. `analyze` writes and runs Python in a sandbox
over the documents, for counting, aggregation and computation across
documents. Both cost a model call.
@ -112,8 +117,9 @@ title = 'Q3 report'
### Errors
A failure is an MCP error, never an empty result. Expected failures carry a
message: a document id that matches nothing, a collection the server does not
cover, a filter the query engine rejects (with its message), invalid base64,
message: a document or section id that matches nothing, a collection the
server does not cover, a filter the query engine rejects (with its message),
invalid base64,
and an `ask_question` or `analyze` failure naming only the exception type.
Anything else reaches the client as `Error calling tool 'name'` and its
traceback goes to the server log.

77
haiku_rag_slim/haiku/rag/context.py
View file

@ -32,6 +32,8 @@ In both cases:
- Results without doc_item_refs pass through unexpanded
"""
from typing import Any
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.document_item import DocumentItem
@ -488,3 +490,78 @@ def expand_with_items(
final_results.append(built)
return final_results + passthrough
def build_toc(
items: list["DocumentItem"],
chunk_index: dict[str, list[str]],
) -> list[dict[str, Any]]:
"""Build a nested section tree from items in position order.
Each ``section_header`` with ``heading_level > 0`` becomes a node. Nesting
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.
``chunk_ids`` aggregates the chunks covered by all items in the section's
``item_range`` (deduped, order preserved). Pass directly to ``cite()`` to
ground a section-scoped answer without a corpus-wide ``search()`` call.
Items without a section_header label (or with ``heading_level == 0``) are
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).
"""
# Defensive: every consumer is supposed to pass items in position order,
# 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
]
if not headers:
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
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):
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 cid in chunk_index.get(item.self_ref, []):
if cid not in seen:
seen.add(cid)
chunk_ids.append(cid)
node: dict[str, Any] = {
"self_ref": h.self_ref,
"level": h.heading_level,
"title": h.text,
"page_numbers": list(h.page_numbers),
"item_range": [h.position, end],
"chunk_ids": chunk_ids,
"children": [],
}
while stack and stack[-1][0] >= h.heading_level:
stack.pop()
(stack[-1][1]["children"] if stack else roots).append(node)
stack.append((h.heading_level, node))
return roots

103
haiku_rag_slim/haiku/rag/mcp.py
View file

@ -13,13 +13,17 @@ from pydantic import Field
from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig, get_config
from haiku.rag.context import build_toc
from haiku.rag.store.exceptions import UnknownDatabaseError
from haiku.rag.store.models import Document, SearchResult
from haiku.rag.store.models.document_item import DocumentItem
from haiku.rag.store.schema import DocumentMetaRecord
from haiku.rag.tools.document import DocumentInfo
from haiku.rag.tools.document import DocumentInfo, DocumentSection, OutlineNode
from haiku.rag.utils import format_citations
if TYPE_CHECKING:
from typing import Any
from haiku.rag.client.scope import DatabaseScope
logger = logging.getLogger(__name__)
@ -91,8 +95,8 @@ def _instructions(scope: "DatabaseScope", config: AppConfig) -> str:
description from the listing."""
lines = [
"haiku-rag is the user's knowledge base: documents they ingested, "
"searchable by meaning and keyword, readable whole, answered with "
"citations, or computed across documents.",
"searchable by meaning and keyword, readable whole or section by "
"section, answered with citations, or computed across documents.",
"Use it whenever a question could be answered from those documents, "
"before answering from memory, and say when it had nothing relevant.",
]
@ -106,6 +110,26 @@ def _instructions(scope: "DatabaseScope", config: AppConfig) -> str:
return "\n".join(lines)
def _node(toc: "dict[str, Any]") -> OutlineNode:
return OutlineNode(
id=toc["self_ref"],
title=toc["title"],
level=toc["level"],
page_numbers=toc["page_numbers"],
children=[_node(child) for child in toc["children"]],
)
def _find(toc: list["dict[str, Any]"], section_id: str) -> "dict[str, Any] | None":
for node in toc:
if node["self_ref"] == section_id:
return node
found = _find(node["children"], section_id)
if found is not None:
return found
return None
def create_mcp_server(
db_path: Path | None = None,
config: AppConfig | None = None,
@ -279,6 +303,72 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
raise ToolError(f"No document with id {document_id!r}")
return document
async def _items_of(document_id: str, source: str | None) -> list[DocumentItem]:
"""A document's items in reading order, from the database holding it."""
rag = await _client()
try:
document = await rag.get_document_by_id(document_id, source)
if document is None:
raise ToolError(f"No document with id {document_id!r}")
owner = await rag.reader_for(source or document.source)
except UnknownDatabaseError as e:
raise ToolError(str(e)) from e
assert owner is not None, "a stored document names its database"
return await owner.document_item_repository.get_all_items(document_id)
@mcp.tool(annotations=_read_only("Document outline"))
async def get_document_outline(
document_id: str, source: str | None = None
) -> list[OutlineNode]:
"""The heading tree of a document, with page numbers.
Use this on a long document to see its structure before reading, then
pass a node's `id` to `get_document_section`. Returns the headings
nested by level; an empty list means the document has no headings,
so read it with `get_document`.
Args:
document_id: The document's id.
source: The collection holding it. Without one every collection
is asked.
"""
return [
_node(toc) for toc in build_toc(await _items_of(document_id, source), {})
]
@mcp.tool(annotations=_read_only("Document section"))
async def get_document_section(
document_id: str, section_id: str, source: str | None = None
) -> DocumentSection:
"""The text of one section of a document, subsections included.
Use this to read a part of a long document instead of the whole.
`section_id` is a node `id` from `get_document_outline`. Returns the
section's heading, page numbers and text in reading order, up to the
next heading of the same or a higher level.
Args:
document_id: The document's id.
section_id: The `id` of a node in the document's outline.
source: The collection holding it. Without one every collection
is asked.
"""
items = await _items_of(document_id, source)
node = _find(build_toc(items, {}), section_id)
if node is None:
raise ToolError(f"No section {section_id!r} in document {document_id!r}")
start, end = node["item_range"]
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
),
)
@mcp.tool(annotations=_read_only("List documents"))
async def list_documents(
limit: int | None = None,
@ -313,7 +403,6 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
@mcp.tool(annotations=_read_only("Ask a question"))
async def ask_question(
question: str,
cite: bool = False,
images_base64: list[str] | None = None,
sources: Sources = None,
) -> str:
@ -321,12 +410,10 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
Use this when the user wants an answer rather than material to read.
It runs a model on the server and is slower than a search. Returns
the answer, followed by citations to the passages it rests on when
`cite` is set.
the answer, followed by citations to the passages it rests on.
Args:
question: The question, in natural language.
cite: Append citations to the answer.
images_base64: Images to attach to the question, PNG or JPEG
bytes as base64. Needs a vision-capable model on the server.
"""
@ -339,7 +426,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
except Exception as e:
logger.exception("ask_question failed")
raise ToolError(f"ask_question failed: {type(e).__name__}") from e
if cite and citations:
if citations:
answer += "\n\n" + format_citations(
citations, include_source=rag.covers_multiple
)

78
haiku_rag_slim/haiku/rag/sandbox/sandbox.py
View file

@ -17,6 +17,7 @@ from pydantic_monty import (
)
from haiku.rag.config.models import AppConfig
from haiku.rag.context import build_toc
from haiku.rag.sandbox.dependencies import AnalysisContext
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX, DocumentItem
@ -38,81 +39,6 @@ class SandboxResult:
success: bool
def _build_toc(
items: list["DocumentItem"],
chunk_index: dict[str, list[str]],
) -> list[dict[str, Any]]:
"""Build a nested section tree from items in position order.
Each ``section_header`` with ``heading_level > 0`` becomes a node. Nesting
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.
``chunk_ids`` aggregates the chunks covered by all items in the section's
``item_range`` (deduped, order preserved). Pass directly to ``cite()`` to
ground a section-scoped answer without a corpus-wide ``search()`` call.
Items without a section_header label (or with ``heading_level == 0``) are
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).
"""
# Defensive: every consumer is supposed to pass items in position order,
# 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
]
if not headers:
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
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):
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 cid in chunk_index.get(item.self_ref, []):
if cid not in seen:
seen.add(cid)
chunk_ids.append(cid)
node: dict[str, Any] = {
"self_ref": h.self_ref,
"level": h.heading_level,
"title": h.text,
"page_numbers": list(h.page_numbers),
"item_range": [h.position, end],
"chunk_ids": chunk_ids,
"children": [],
}
while stack and stack[-1][0] >= h.heading_level:
stack.pop()
(stack[-1][1]["children"] if stack else roots).append(node)
stack.append((h.heading_level, node))
return roots
class Sandbox:
"""Execute code in a sandboxed Python interpreter.
@ -520,7 +446,7 @@ class Sandbox:
{
"doc_id": did,
"title": doc_titles.get(did),
"tree": _build_toc(items, chunk_index),
"tree": build_toc(items, chunk_index),
},
ensure_ascii=False,
)

19
haiku_rag_slim/haiku/rag/tools/document.py
View file

@ -31,6 +31,25 @@ class DocumentInfo(BaseModel):
metadata: dict = {}
class OutlineNode(BaseModel):
"""A heading in a document's outline. `id` is the heading item's self_ref."""
id: str
title: str
level: int
page_numbers: list[int] = []
children: list["OutlineNode"] = []
class DocumentSection(BaseModel):
"""One section's text in reading order, subsections included."""
id: str
title: str
page_numbers: list[int] = []
content: str
class DocumentListResponse(BaseModel):
"""Response from list_documents tool."""

191
tests/test_mcp.py
View file

@ -236,9 +236,7 @@ class TestMCPReadTools:
assert overview["metadata"] == {"author": "Ada"}
@pytest.mark.asyncio
async def test_ask_question_appends_citations_when_requested(
self, mcp_db, monkeypatch
):
async def test_ask_question_appends_the_citations(self, mcp_db, monkeypatch):
from haiku.rag.store.models.citation import Citation
citation = Citation(
@ -257,13 +255,182 @@ class TestMCPReadTools:
mcp = create_mcp_server(mcp_db)
ask = await _get_tool(mcp, "ask_question")
with_cite = await ask(question="q", cite=True)
assert with_cite.startswith("the answer")
assert "AI Overview" in with_cite
answer = await ask(question="q")
assert answer.startswith("the answer")
assert "AI Overview" in answer
# One database: its name adds nothing.
assert "alpha" not in with_cite
assert "alpha" not in answer
assert await ask(question="q", cite=False) == "the answer"
@pytest.fixture
async def outlined_db(temp_db_path):
"""A database with one document whose items carry a heading hierarchy.
Rows are written through the repositories, so no embedder is involved.
Returns the path and the document id."""
from haiku.rag.store.models.document import Document as DocumentModel
from haiku.rag.store.models.document_item import DocumentItem
def header(pos, level, text):
return DocumentItem(
document_id="",
position=pos,
self_ref=f"#/texts/{pos}",
label="section_header",
text=text,
page_numbers=[pos // 4 + 1],
heading_level=level,
)
def para(pos):
return DocumentItem(
document_id="",
position=pos,
self_ref=f"#/texts/{pos}",
label="paragraph",
text=f"para{pos}",
page_numbers=[pos // 4 + 1],
)
async with HaikuRAG(temp_db_path, create=True) as rag:
doc = await rag.document_repository.create(
DocumentModel(content="x", uri="test://outlined", title="Outlined")
)
items = [
header(0, 1, "Intro"),
para(1),
header(2, 2, "Background"),
para(3),
header(4, 3, "Prior Work"),
para(5),
header(6, 2, "Approach"),
para(7),
header(8, 1, "Methods"),
para(9),
]
for item in items:
item.document_id = doc.id
await rag.document_item_repository.create_items(doc.id, items)
return temp_db_path, doc.id
class TestMCPDocumentNavigation:
@pytest.mark.asyncio
async def test_the_outline_nests_headings_by_level(self, outlined_db):
db, doc_id = outlined_db
outline = await _get_tool(create_mcp_server(db), "get_document_outline")
roots = await outline(document_id=doc_id)
assert [n.title for n in roots] == ["Intro", "Methods"]
intro = roots[0]
assert (intro.id, intro.level, intro.page_numbers) == ("#/texts/0", 1, [1])
assert [c.title for c in intro.children] == ["Background", "Approach"]
assert [c.title for c in intro.children[0].children] == ["Prior Work"]
assert intro.children[0].children[0].level == 3
assert roots[1].children == []
@pytest.mark.asyncio
async def test_a_document_without_headings_has_an_empty_outline(self, mcp_db):
mcp = create_mcp_server(mcp_db)
[doc] = await (await _get_tool(mcp, "list_documents"))(limit=1)
outline = await _get_tool(mcp, "get_document_outline")
assert await outline(document_id=doc.id) == []
@pytest.mark.asyncio
async def test_a_section_covers_its_subsections_and_stops_at_its_sibling(
self, outlined_db
):
db, doc_id = outlined_db
section = await _get_tool(create_mcp_server(db), "get_document_section")
background = await section(document_id=doc_id, section_id="#/texts/2")
assert background.title == "Background"
assert background.content.split("\n\n") == [
"Background",
"para3",
"Prior Work",
"para5",
]
assert background.page_numbers == [1]
intro = await section(document_id=doc_id, section_id="#/texts/0")
assert intro.content.startswith("Intro")
assert "para7" in intro.content
assert "Methods" not in intro.content
@pytest.mark.asyncio
async def test_an_unknown_section_or_document_is_an_error(self, outlined_db):
db, doc_id = outlined_db
mcp = create_mcp_server(db)
section = await _get_tool(mcp, "get_document_section")
outline = await _get_tool(mcp, "get_document_outline")
with pytest.raises(ToolError, match="#/texts/99"):
await section(document_id=doc_id, section_id="#/texts/99")
with pytest.raises(ToolError, match="nonexistent-id"):
await outline(document_id="nonexistent-id")
with pytest.raises(ToolError, match="nonexistent-id"):
await section(document_id="nonexistent-id", section_id="#/texts/0")
@pytest.mark.asyncio
@pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning")
async def test_outline_and_section_serialize_over_the_wire(self, outlined_db):
db, doc_id = outlined_db
mcp = create_mcp_server(db)
outline = await _call(mcp, "get_document_outline", document_id=doc_id)
section = await _call(
mcp, "get_document_section", document_id=doc_id, section_id="#/texts/8"
)
assert not outline.is_error and not section.is_error
[intro, methods] = outline.structured_content["result"]
assert set(intro) == {"id", "title", "level", "page_numbers", "children"}
assert intro["children"][0]["children"][0]["title"] == "Prior Work"
assert set(section.structured_content) == {
"id",
"title",
"page_numbers",
"content",
}
assert section.structured_content["content"] == "Methods\n\npara9"
@pytest.mark.asyncio
async def test_source_routes_to_the_database_holding_the_document(self, two_dbs):
from haiku.rag.store.models.document_item import DocumentItem
async with HaikuRAG(config=two_dbs, sources=["beta"]) as beta:
[doc] = await beta.list_documents()
await beta.document_item_repository.create_items(
doc.id,
[
DocumentItem(
document_id=doc.id,
position=0,
self_ref="#/texts/0",
label="section_header",
text="Only in beta",
heading_level=1,
)
],
)
mcp = _covering_all(two_dbs)
outline = await _get_tool(mcp, "get_document_outline")
section = await _get_tool(mcp, "get_document_section")
named = await outline(document_id=doc.id, source="beta")
found = await outline(document_id=doc.id)
assert [n.title for n in named] == [n.title for n in found] == ["Only in beta"]
assert (
await section(document_id=doc.id, section_id="#/texts/0", source="beta")
).title == "Only in beta"
with pytest.raises(ToolError, match="nope"):
await outline(document_id=doc.id, source="nope")
with pytest.raises(ToolError, match=doc.id):
await outline(document_id=doc.id, source="alpha")
@pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning")
@ -324,7 +491,7 @@ class TestMCPDescribesItself:
async with Client(create_mcp_server(mcp_db)) as client:
tools = await client.list_tools()
assert len(tools) == 6
assert len(tools) == 8
for tool in tools:
assert tool.annotations is not None, tool.name
assert tool.annotations.readOnlyHint is True, tool.name
@ -344,7 +511,7 @@ class TestMCPDescribesItself:
for name, schema in tool.inputSchema.get("properties", {}).items()
if not schema.get("description")
]
assert len(tools) == 6
assert len(tools) == 8
assert undescribed == []
@ -356,6 +523,8 @@ class TestMCPToolSet:
assert {t.name for t in await mcp.list_tools()} == {
"search_documents",
"get_document",
"get_document_outline",
"get_document_section",
"list_documents",
"ask_question",
"analyze",
@ -475,7 +644,7 @@ class TestMCPCoversTheConfiguredSet:
mcp = _covering_all(two_dbs)
ask = await _get_tool(mcp, "ask_question")
answer = await ask(question="q", cite=True)
answer = await ask(question="q")
assert "alpha" in answer
assert "beta" in answer