From 0026142e7d033da067113cd9042397eb71a9afeb Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 4 Sep 2026 15:50:24 +0300 Subject: [PATCH] Give MCP search results one channel and tidy two messages Search results carry text and image blocks only. Claude Code and the Agent SDK do not forward text blocks when structuredContent is present and Desktop forwards both, so sending both either hid the rendering or doubled it. The invalid-filter error keeps the engine's diagnosis and lists our columns instead of lance's internals. format_citations no longer repeats the URI of an untitled document. Refs #599 --- CHANGELOG.md | 12 ++-- docs/mcp.md | 9 ++- haiku_rag_slim/haiku/rag/mcp.py | 21 +++--- haiku_rag_slim/haiku/rag/utils.py | 7 +- tests/test_mcp.py | 109 +++++++++++++++++------------- tests/test_utils.py | 17 +++++ 6 files changed, 107 insertions(+), 68 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9716736d..64343190 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,10 +35,10 @@ - MCP `search_documents` and `search_documents_by_image` expand results to their section (`HaikuRAG.expand_context`) and return the agent rendering as text (rank, `Document ID`, `Collection` over several databases, title, - headings, the matched chunk's metadata, passage), pictures as - `ImageContent` blocks, and the `SearchResult` list without `image_data` as - structured content. `SearchResult.format_for_agent(include_document_id=, - include_chunk_meta=)`; `collect_pictures` in `haiku.rag.tools.search`. + headings, the matched chunk's metadata, passage) and pictures as + `ImageContent` blocks, with no structured content. + `SearchResult.format_for_agent(include_document_id=, include_chunk_meta=)`; + `collect_pictures` in `haiku.rag.tools.search`. - MCP tools raise on failure; an empty result no longer doubles as an error. Unknown document, unknown collection, invalid filter and invalid base64 carry a message; `ask_question` and `analyze` failures name the exception @@ -50,6 +50,10 @@ `DocumentInfo.source`; citations name their database when the server covers several. `format_citations(citations, include_source=False)`. +### Fixed + +- MCP citations no longer repeat the URI of an untitled document. + ### Removed - `cite` on the MCP `ask_question` tool; citations are always appended. diff --git a/docs/mcp.md b/docs/mcp.md index ab3b8aa7..df7b4448 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -113,8 +113,9 @@ title, section headings, the matched chunk's metadata when it has any, and the passage expanded to its section the way the agents get it (`search.max_context_chars` caps it). Pictures in the results follow as image blocks, one per distinct picture, each preceded by a line naming its -result; `include_images: false` leaves them out. The structured content is the -`SearchResult` list without picture bytes. Scores are not comparable across +result; `include_images: false` leaves them out. Search results carry no +structured content, so every client shows the model the same text and +images. Scores are not comparable across queries or search types, so rank is the signal. `search_documents_by_image` embeds the query image and searches by vector similarity alone. @@ -128,7 +129,9 @@ which is how a client learns what a filter can match. `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. +documents. Both cost a model call. Claude Code moves a call still running +after about two minutes to a background task, which a slow local model can +trigger; `--no-agents` leaves both tools out. ### Filters diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index c80ea1ec..112a156c 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -1,5 +1,6 @@ import asyncio import logging +import re from collections.abc import AsyncIterator from contextlib import AsyncExitStack, asynccontextmanager from importlib import metadata @@ -89,7 +90,12 @@ async def _check_filter( try: await selected[0].count_documents(filter=filter) except ValueError as e: - raise ToolError(f"Invalid filter {filter!r}: {e}") from e + # The engine lists its own columns too, lance internals among them. + reason = re.sub(r"\s*Valid fields are .*", "", str(e), flags=re.DOTALL) + raise ToolError( + f"Invalid filter {filter!r}: {reason.rstrip('. ')}. " + f"Columns: {_FILTER_COLUMNS}." + ) from e def _instructions(scope: "DatabaseScope", config: AppConfig, agents: bool) -> str: @@ -121,7 +127,8 @@ def _instructions(scope: "DatabaseScope", config: AppConfig, agents: bool) -> st def _search_result(results: list[SearchResult], covers_multiple: bool) -> ToolResult: """Results as the in-process agents read them, plus the matched chunk's metadata, then each distinct picture as an image block labelled with its - result, and the results as structured content without the picture bytes.""" + result. No structured content: a client given both shows the model the + JSON and drops the text, or shows both.""" import base64 total = len(results) @@ -154,15 +161,7 @@ def _search_result(results: list[SearchResult], covers_multiple: bool) -> ToolRe mime_type="image/png", ) ) - return ToolResult( - content=content, - structured_content={ - "result": [ - result.model_dump(mode="json", exclude={"image_data"}) - for result in results - ] - }, - ) + return ToolResult(content=content) def _node(toc: "dict[str, Any]") -> OutlineNode: diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index b96bf74b..2d71b47c 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -421,11 +421,12 @@ def format_citations(citations: "list[Citation]", include_source: bool = False) if section: location_parts.append(f"Section: {section}") - source = c.document_uri + # The URI is the header when there is no title; do not repeat it. + line = f"{header} {c.document_uri}" if c.document_title else header if location_parts: - source += f" - {', '.join(location_parts)}" + line += f" - {', '.join(location_parts)}" - lines.append(f"{header} {source}") + lines.append(line) for ref in c.picture_refs: lines.append(f"[Figure: {ref}]") lines.append(c.content) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index ddc46fdb..0fada9a3 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1,4 +1,5 @@ import logging +import re from pathlib import Path from types import SimpleNamespace @@ -98,9 +99,25 @@ async def _call(mcp, name, **kwargs): return await client.call_tool(name, kwargs, raise_on_error=False) -def _results(search_result) -> list[dict]: - """The search results a tool returned, as the client sees them.""" - return search_result.structured_content["result"] +def _results(tool_result) -> list[dict]: + """A tool's structured result list, as the client sees it.""" + return tool_result.structured_content["result"] + + +_HEADER = re.compile(r"^\[[^\]]+\] \[rank \d+ of \d+\]$", re.MULTILINE) + + +def _rendered(search_result) -> list[str]: + """The result blocks of a search, split from the text the model reads.""" + text = search_result.content[0].text + starts = [match.start() for match in _HEADER.finditer(text)] + return [text[a:b].strip() for a, b in zip(starts, starts[1:] + [len(text)])] + + +def _line(block: str, name: str) -> str | None: + """The value of a `Name: value` line in a rendered result, if present.""" + match = re.search(rf"^{re.escape(name)}: (.+)$", block, re.MULTILINE) + return match.group(1) if match else None def _png_b64() -> str: @@ -120,17 +137,17 @@ class TestMCPReadTools: mcp = create_mcp_server(mcp_db) search = await _get_tool(mcp, "search_documents") - results = _results(await search(query="artificial intelligence")) - assert len(results) > 0 - assert all(r["chunk_id"] and r["content"] for r in results) + blocks = _rendered(await search(query="artificial intelligence")) + assert blocks + assert all("Content:" in block for block in blocks) @pytest.mark.asyncio async def test_search_documents_with_limit(self, mcp_db): mcp = create_mcp_server(mcp_db) search = await _get_tool(mcp, "search_documents") - results = _results(await search(query="artificial intelligence", limit=1)) - assert len(results) == 1 + blocks = _rendered(await search(query="artificial intelligence", limit=1)) + assert len(blocks) == 1 @pytest.mark.asyncio @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") @@ -143,18 +160,15 @@ class TestMCPReadTools: {"query": "artificial intelligence", "filter": "title = 'ML Basics'"}, ) - results = result.structured_content["result"] - assert results - assert {r["document_title"] for r in results} == {"ML Basics"} + blocks = _rendered(result) + assert blocks + assert all('"ML Basics"' in block for block in blocks) @pytest.mark.asyncio @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") - async def test_search_documents_preserves_chunk_meta_through_serialization( - self, mcp_db - ): - """Chunk_meta must survive FastMCP's actual wire serialization. - - Calling the tool function directly bypasses that serialization step entirely.""" + async def test_search_documents_carries_the_matched_chunks_metadata(self, mcp_db): + """The chunk's own metadata reaches the text the model reads, over the + wire, without haiku.rag's structural keys.""" from fastmcp import Client async with HaikuRAG(mcp_db, create=True) as rag: @@ -176,14 +190,11 @@ class TestMCPReadTools: "search_documents", {"query": "artificial intelligence"} ) - results = result.structured_content["result"] - assert results - assert any( - r["chunk_meta"] == {"fake-metadata-for-testing": "42"} for r in results - ) text = result.content[0].text assert "fake-metadata-for-testing" in text assert "42" in text + assert "doc_item_refs" not in text + assert result.structured_content is None @pytest.mark.asyncio @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") @@ -196,9 +207,8 @@ class TestMCPReadTools: result = await _call(_covering_all(config), "search_documents", query="gardens") - [hit] = _results(result) - assert all(sentence in hit["content"] for sentence in sentences) - assert all(sentence in result.content[0].text for sentence in sentences) + [hit] = _rendered(result) + assert all(sentence in hit for sentence in sentences) @pytest.mark.asyncio async def test_get_document(self, mcp_db): @@ -531,9 +541,12 @@ class TestMCPSearchResultShape: assert [ label for label in labels if "[c3]" in label and "#/pictures/3" in label ] - structured = _results(result) - assert [r["chunk_id"] for r in structured] == ["c1", "c2", "c3"] - assert all("image_data" not in r for r in structured) + assert [block.split("]")[0] for block in _rendered(result)] == [ + "[c1", + "[c2", + "[c3", + ] + assert result.structured_content is None @pytest.mark.asyncio async def test_an_undecodable_picture_yields_no_image(self, mcp_db, monkeypatch): @@ -566,7 +579,7 @@ class TestMCPSearchResultShape: result = await _call(create_mcp_server(mcp_db), "search_documents", query="q") assert [block.text for block in result.content] == ["No results found."] - assert _results(result) == [] + assert result.structured_content is None @pytest.mark.asyncio async def test_search_text_alone_drives_the_document_tools(self, two_dbs): @@ -600,7 +613,7 @@ class TestMCPSearchResultShape: r"Document ID: (\S+)\nCollection: (\S+)", search.content[0].text ) - assert len(pairs) == len(_results(search)) == 2 + assert len(pairs) == len(_rendered(search)) == 2 assert {source for _, source in pairs} == {"alpha", "beta"} for document_id, source in pairs: outline = await _call( @@ -747,19 +760,19 @@ class TestMCPCoversTheConfiguredSet: mcp = _covering_all(two_dbs) search = await _get_tool(mcp, "search_documents") - results = _results(await search(query="cats")) + blocks = _rendered(await search(query="cats")) - assert {r["source"] for r in results} == {"alpha", "beta"} + assert {_line(block, "Collection") for block in blocks} == {"alpha", "beta"} @pytest.mark.asyncio async def test_sources_narrows_the_search(self, two_dbs): mcp = _covering_all(two_dbs) search = await _get_tool(mcp, "search_documents") - results = _results(await search(query="cats", sources=["beta"])) + blocks = _rendered(await search(query="cats", sources=["beta"])) - assert results - assert {r["source"] for r in results} == {"beta"} + assert blocks + assert {_line(block, "Collection") for block in blocks} == {"beta"} @pytest.mark.asyncio @pytest.mark.parametrize( @@ -793,13 +806,13 @@ class TestMCPCoversTheConfiguredSet: mcp = _covering_all(two_dbs) search = await _get_tool(mcp, "search_documents") - results = _results( + blocks = _rendered( await search(query="cats", filter="uri LIKE '%beta%'", sources=["beta"]) ) - assert results - assert {r["source"] for r in results} == {"beta"} + assert blocks + assert {_line(block, "Collection") for block in blocks} == {"beta"} none = await search(query="cats", filter="uri LIKE '%beta%'", sources=[]) - assert _results(none) == [] + assert _rendered(none) == [] @pytest.mark.asyncio async def test_the_listing_covers_every_database(self, two_dbs): @@ -828,9 +841,9 @@ class TestMCPCoversTheConfiguredSet: mcp = create_mcp_server(config=two_dbs) search = await _get_tool(mcp, "search_documents") - results = _results(await search(query="cats")) + blocks = _rendered(await search(query="cats")) - assert {r["source"] for r in results} == {"alpha", "beta"} + assert {_line(block, "Collection") for block in blocks} == {"alpha", "beta"} @pytest.mark.asyncio async def test_ask_question_names_each_citations_database( @@ -924,7 +937,7 @@ class TestMCPImageQuery: sources=[], ) - assert _results(results) == [] + assert _rendered(results) == [] assert seen["query"] == png assert seen["filter"] == "uri LIKE 'x%'" assert seen["sources"] == [] @@ -1036,6 +1049,8 @@ class TestMCPErrorContract: assert result.is_error assert "no_such_column = 1" in result.content[0].text + assert "created_at" in result.content[0].text + assert "_rowid" not in result.content[0].text @pytest.mark.asyncio @pytest.mark.parametrize("filter", [None, "title = 'AI Overview'"]) @@ -1281,12 +1296,12 @@ class TestMCPClientLifetime: mcp = _mcp_covering(scope, config) async with mcp._lifespan_manager(): search = await _get_tool(mcp, "search_documents") - results = _results(await search(query="artificial intelligence")) + blocks = _rendered(await search(query="artificial intelligence")) listing = await _get_tool(mcp, "list_documents") documents = await listing() - assert results - assert {r["source"] for r in results} == {"alpha"} + assert blocks + assert {_line(block, "Collection") for block in blocks} == {None} titles = {d.title for d in documents} assert "AI Overview" in titles assert "Zebras" not in titles @@ -1371,9 +1386,9 @@ class TestMCPClientLifetime: assert opens == 1 async with mcp._lifespan_manager(): - results = _results(await search(query="artificial intelligence")) + blocks = _rendered(await search(query="artificial intelligence")) assert opens == 2 - assert len(results) > 0 + assert blocks @pytest.mark.asyncio async def test_same_dim_drift_starts(self, mcp_db): diff --git a/tests/test_utils.py b/tests/test_utils.py index 5284cb04..9d4cba45 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -769,6 +769,23 @@ def test_format_citations_names_the_source_when_asked(): assert "papers" not in format_citations([citation]) +def test_format_citations_names_an_untitled_document_once(): + from haiku.rag.store.models.citation import Citation + from haiku.rag.utils import format_citations + + citation = Citation( + document_id="doc1", + chunk_id="chunk1", + document_uri="test://doc", + content="Content", + page_numbers=[3], + ) + result = format_citations([citation]) + + assert result.count("test://doc") == 1 + assert "[1] test://doc - p. 3" in result + + # --- format_citations tests (pictures) ---