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) ---