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
This commit is contained in:
Yiorgis Gozadinos 2026-09-04 15:50:24 +03:00
parent 2582f2c05a
commit 0026142e7d
No known key found for this signature in database
6 changed files with 107 additions and 68 deletions

View file

@ -35,10 +35,10 @@
- MCP `search_documents` and `search_documents_by_image` expand results to - MCP `search_documents` and `search_documents_by_image` expand results to
their section (`HaikuRAG.expand_context`) and return the agent rendering their section (`HaikuRAG.expand_context`) and return the agent rendering
as text (rank, `Document ID`, `Collection` over several databases, title, as text (rank, `Document ID`, `Collection` over several databases, title,
headings, the matched chunk's metadata, passage), pictures as headings, the matched chunk's metadata, passage) and pictures as
`ImageContent` blocks, and the `SearchResult` list without `image_data` as `ImageContent` blocks, with no structured content.
structured content. `SearchResult.format_for_agent(include_document_id=, `SearchResult.format_for_agent(include_document_id=, include_chunk_meta=)`;
include_chunk_meta=)`; `collect_pictures` in `haiku.rag.tools.search`. `collect_pictures` in `haiku.rag.tools.search`.
- MCP tools raise on failure; an empty result no longer doubles as an error. - MCP tools raise on failure; an empty result no longer doubles as an error.
Unknown document, unknown collection, invalid filter and invalid base64 Unknown document, unknown collection, invalid filter and invalid base64
carry a message; `ask_question` and `analyze` failures name the exception carry a message; `ask_question` and `analyze` failures name the exception
@ -50,6 +50,10 @@
`DocumentInfo.source`; citations name their database when the server `DocumentInfo.source`; citations name their database when the server
covers several. `format_citations(citations, include_source=False)`. covers several. `format_citations(citations, include_source=False)`.
### Fixed
- MCP citations no longer repeat the URI of an untitled document.
### Removed ### Removed
- `cite` on the MCP `ask_question` tool; citations are always appended. - `cite` on the MCP `ask_question` tool; citations are always appended.

View file

@ -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 passage expanded to its section the way the agents get it
(`search.max_context_chars` caps it). Pictures in the results follow as (`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 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 result; `include_images: false` leaves them out. Search results carry no
`SearchResult` list without picture bytes. Scores are not comparable across 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` queries or search types, so rank is the signal. `search_documents_by_image`
embeds the query image and searches by vector similarity alone. 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 `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 followed by its citations. `analyze` writes and runs Python in a sandbox
over the documents, for counting, aggregation and computation across 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 ### Filters

View file

@ -1,5 +1,6 @@
import asyncio import asyncio
import logging import logging
import re
from collections.abc import AsyncIterator from collections.abc import AsyncIterator
from contextlib import AsyncExitStack, asynccontextmanager from contextlib import AsyncExitStack, asynccontextmanager
from importlib import metadata from importlib import metadata
@ -89,7 +90,12 @@ async def _check_filter(
try: try:
await selected[0].count_documents(filter=filter) await selected[0].count_documents(filter=filter)
except ValueError as e: 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: 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: def _search_result(results: list[SearchResult], covers_multiple: bool) -> ToolResult:
"""Results as the in-process agents read them, plus the matched chunk's """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 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 import base64
total = len(results) total = len(results)
@ -154,15 +161,7 @@ def _search_result(results: list[SearchResult], covers_multiple: bool) -> ToolRe
mime_type="image/png", mime_type="image/png",
) )
) )
return ToolResult( return ToolResult(content=content)
content=content,
structured_content={
"result": [
result.model_dump(mode="json", exclude={"image_data"})
for result in results
]
},
)
def _node(toc: "dict[str, Any]") -> OutlineNode: def _node(toc: "dict[str, Any]") -> OutlineNode:

View file

@ -421,11 +421,12 @@ def format_citations(citations: "list[Citation]", include_source: bool = False)
if section: if section:
location_parts.append(f"Section: {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: 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: for ref in c.picture_refs:
lines.append(f"[Figure: {ref}]") lines.append(f"[Figure: {ref}]")
lines.append(c.content) lines.append(c.content)

View file

@ -1,4 +1,5 @@
import logging import logging
import re
from pathlib import Path from pathlib import Path
from types import SimpleNamespace 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) return await client.call_tool(name, kwargs, raise_on_error=False)
def _results(search_result) -> list[dict]: def _results(tool_result) -> list[dict]:
"""The search results a tool returned, as the client sees them.""" """A tool's structured result list, as the client sees it."""
return search_result.structured_content["result"] 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: def _png_b64() -> str:
@ -120,17 +137,17 @@ class TestMCPReadTools:
mcp = create_mcp_server(mcp_db) mcp = create_mcp_server(mcp_db)
search = await _get_tool(mcp, "search_documents") search = await _get_tool(mcp, "search_documents")
results = _results(await search(query="artificial intelligence")) blocks = _rendered(await search(query="artificial intelligence"))
assert len(results) > 0 assert blocks
assert all(r["chunk_id"] and r["content"] for r in results) assert all("Content:" in block for block in blocks)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_search_documents_with_limit(self, mcp_db): async def test_search_documents_with_limit(self, mcp_db):
mcp = create_mcp_server(mcp_db) mcp = create_mcp_server(mcp_db)
search = await _get_tool(mcp, "search_documents") search = await _get_tool(mcp, "search_documents")
results = _results(await search(query="artificial intelligence", limit=1)) blocks = _rendered(await search(query="artificial intelligence", limit=1))
assert len(results) == 1 assert len(blocks) == 1
@pytest.mark.asyncio @pytest.mark.asyncio
@pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning")
@ -143,18 +160,15 @@ class TestMCPReadTools:
{"query": "artificial intelligence", "filter": "title = 'ML Basics'"}, {"query": "artificial intelligence", "filter": "title = 'ML Basics'"},
) )
results = result.structured_content["result"] blocks = _rendered(result)
assert results assert blocks
assert {r["document_title"] for r in results} == {"ML Basics"} assert all('"ML Basics"' in block for block in blocks)
@pytest.mark.asyncio @pytest.mark.asyncio
@pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning")
async def test_search_documents_preserves_chunk_meta_through_serialization( async def test_search_documents_carries_the_matched_chunks_metadata(self, mcp_db):
self, mcp_db """The chunk's own metadata reaches the text the model reads, over the
): wire, without haiku.rag's structural keys."""
"""Chunk_meta must survive FastMCP's actual wire serialization.
Calling the tool function directly bypasses that serialization step entirely."""
from fastmcp import Client from fastmcp import Client
async with HaikuRAG(mcp_db, create=True) as rag: async with HaikuRAG(mcp_db, create=True) as rag:
@ -176,14 +190,11 @@ class TestMCPReadTools:
"search_documents", {"query": "artificial intelligence"} "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 text = result.content[0].text
assert "fake-metadata-for-testing" in text assert "fake-metadata-for-testing" in text
assert "42" in text assert "42" in text
assert "doc_item_refs" not in text
assert result.structured_content is None
@pytest.mark.asyncio @pytest.mark.asyncio
@pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") @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") result = await _call(_covering_all(config), "search_documents", query="gardens")
[hit] = _results(result) [hit] = _rendered(result)
assert all(sentence in hit["content"] for sentence in sentences) assert all(sentence in hit for sentence in sentences)
assert all(sentence in result.content[0].text for sentence in sentences)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_get_document(self, mcp_db): async def test_get_document(self, mcp_db):
@ -531,9 +541,12 @@ class TestMCPSearchResultShape:
assert [ assert [
label for label in labels if "[c3]" in label and "#/pictures/3" in label label for label in labels if "[c3]" in label and "#/pictures/3" in label
] ]
structured = _results(result) assert [block.split("]")[0] for block in _rendered(result)] == [
assert [r["chunk_id"] for r in structured] == ["c1", "c2", "c3"] "[c1",
assert all("image_data" not in r for r in structured) "[c2",
"[c3",
]
assert result.structured_content is None
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_an_undecodable_picture_yields_no_image(self, mcp_db, monkeypatch): 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") result = await _call(create_mcp_server(mcp_db), "search_documents", query="q")
assert [block.text for block in result.content] == ["No results found."] assert [block.text for block in result.content] == ["No results found."]
assert _results(result) == [] assert result.structured_content is None
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_search_text_alone_drives_the_document_tools(self, two_dbs): 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 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"} assert {source for _, source in pairs} == {"alpha", "beta"}
for document_id, source in pairs: for document_id, source in pairs:
outline = await _call( outline = await _call(
@ -747,19 +760,19 @@ class TestMCPCoversTheConfiguredSet:
mcp = _covering_all(two_dbs) mcp = _covering_all(two_dbs)
search = await _get_tool(mcp, "search_documents") 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 @pytest.mark.asyncio
async def test_sources_narrows_the_search(self, two_dbs): async def test_sources_narrows_the_search(self, two_dbs):
mcp = _covering_all(two_dbs) mcp = _covering_all(two_dbs)
search = await _get_tool(mcp, "search_documents") 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 blocks
assert {r["source"] for r in results} == {"beta"} assert {_line(block, "Collection") for block in blocks} == {"beta"}
@pytest.mark.asyncio @pytest.mark.asyncio
@pytest.mark.parametrize( @pytest.mark.parametrize(
@ -793,13 +806,13 @@ class TestMCPCoversTheConfiguredSet:
mcp = _covering_all(two_dbs) mcp = _covering_all(two_dbs)
search = await _get_tool(mcp, "search_documents") search = await _get_tool(mcp, "search_documents")
results = _results( blocks = _rendered(
await search(query="cats", filter="uri LIKE '%beta%'", sources=["beta"]) await search(query="cats", filter="uri LIKE '%beta%'", sources=["beta"])
) )
assert results assert blocks
assert {r["source"] for r in results} == {"beta"} assert {_line(block, "Collection") for block in blocks} == {"beta"}
none = await search(query="cats", filter="uri LIKE '%beta%'", sources=[]) none = await search(query="cats", filter="uri LIKE '%beta%'", sources=[])
assert _results(none) == [] assert _rendered(none) == []
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_the_listing_covers_every_database(self, two_dbs): async def test_the_listing_covers_every_database(self, two_dbs):
@ -828,9 +841,9 @@ class TestMCPCoversTheConfiguredSet:
mcp = create_mcp_server(config=two_dbs) mcp = create_mcp_server(config=two_dbs)
search = await _get_tool(mcp, "search_documents") 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 @pytest.mark.asyncio
async def test_ask_question_names_each_citations_database( async def test_ask_question_names_each_citations_database(
@ -924,7 +937,7 @@ class TestMCPImageQuery:
sources=[], sources=[],
) )
assert _results(results) == [] assert _rendered(results) == []
assert seen["query"] == png assert seen["query"] == png
assert seen["filter"] == "uri LIKE 'x%'" assert seen["filter"] == "uri LIKE 'x%'"
assert seen["sources"] == [] assert seen["sources"] == []
@ -1036,6 +1049,8 @@ class TestMCPErrorContract:
assert result.is_error assert result.is_error
assert "no_such_column = 1" in result.content[0].text 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.asyncio
@pytest.mark.parametrize("filter", [None, "title = 'AI Overview'"]) @pytest.mark.parametrize("filter", [None, "title = 'AI Overview'"])
@ -1281,12 +1296,12 @@ class TestMCPClientLifetime:
mcp = _mcp_covering(scope, config) mcp = _mcp_covering(scope, config)
async with mcp._lifespan_manager(): async with mcp._lifespan_manager():
search = await _get_tool(mcp, "search_documents") 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") listing = await _get_tool(mcp, "list_documents")
documents = await listing() documents = await listing()
assert results assert blocks
assert {r["source"] for r in results} == {"alpha"} assert {_line(block, "Collection") for block in blocks} == {None}
titles = {d.title for d in documents} titles = {d.title for d in documents}
assert "AI Overview" in titles assert "AI Overview" in titles
assert "Zebras" not in titles assert "Zebras" not in titles
@ -1371,9 +1386,9 @@ class TestMCPClientLifetime:
assert opens == 1 assert opens == 1
async with mcp._lifespan_manager(): async with mcp._lifespan_manager():
results = _results(await search(query="artificial intelligence")) blocks = _rendered(await search(query="artificial intelligence"))
assert opens == 2 assert opens == 2
assert len(results) > 0 assert blocks
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_same_dim_drift_starts(self, mcp_db): async def test_same_dim_drift_starts(self, mcp_db):

View file

@ -769,6 +769,23 @@ def test_format_citations_names_the_source_when_asked():
assert "papers" not in format_citations([citation]) 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) --- # --- format_citations tests (pictures) ---