From 15afb97a6e2a73f44213eb481342e196269be813 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 4 Sep 2026 12:16:00 +0300 Subject: [PATCH] Navigate documents by outline and section from the MCP server 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 --- CHANGELOG.md | 4 + docs/mcp.md | 22 ++- haiku_rag_slim/haiku/rag/context.py | 77 ++++++++ haiku_rag_slim/haiku/rag/mcp.py | 103 ++++++++++- haiku_rag_slim/haiku/rag/sandbox/sandbox.py | 78 +------- haiku_rag_slim/haiku/rag/tools/document.py | 19 ++ tests/test_mcp.py | 191 ++++++++++++++++++-- 7 files changed, 391 insertions(+), 103 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7940d86c..9f6fb5c6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/docs/mcp.md b/docs/mcp.md index d63f5880..c4815434 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -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. diff --git a/haiku_rag_slim/haiku/rag/context.py b/haiku_rag_slim/haiku/rag/context.py index b4214da5..17fb44c5 100644 --- a/haiku_rag_slim/haiku/rag/context.py +++ b/haiku_rag_slim/haiku/rag/context.py @@ -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 diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 5c285952..ec97476f 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -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 ) diff --git a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py index 4102d961..cddbdd46 100644 --- a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py +++ b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py @@ -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, ) diff --git a/haiku_rag_slim/haiku/rag/tools/document.py b/haiku_rag_slim/haiku/rag/tools/document.py index 960d5194..89648127 100644 --- a/haiku_rag_slim/haiku/rag/tools/document.py +++ b/haiku_rag_slim/haiku/rag/tools/document.py @@ -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.""" diff --git a/tests/test_mcp.py b/tests/test_mcp.py index b21b24cd..b908a190 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -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