From 2582f2c05acdda5690b4b9c323a46417bc78e769 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 4 Sep 2026 14:12:12 +0300 Subject: [PATCH] Expand MCP search results and render the matched chunk's metadata Both search tools pass their results through HaikuRAG.expand_context, as every other consumer of search results already did, so a client reads the hit in its section rather than the chunk that matched. The rendering gains an opt-in include_chunk_meta that shows the metadata stored with the matched chunk beyond haiku.rag's structural keys, labelled as the matched chunk's because an expanded passage spans several chunks and only the anchor's metadata survives expansion. The capabilities' rendering is unchanged. Refs #599 --- CHANGELOG.md | 13 ++++---- claude-plugin/skills/haiku-rag/SKILL.md | 3 +- docs/mcp.md | 4 ++- haiku_rag_slim/haiku/rag/mcp.py | 20 +++++++----- .../haiku/rag/store/models/chunk.py | 20 ++++++++++-- tests/test_chunk.py | 31 +++++++++++++++++++ tests/test_mcp.py | 20 +++++++++++- 7 files changed, 92 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 848bf844..9716736d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -32,12 +32,13 @@ - `processing.conversion_options.picture_description.model` defaults to `enable_thinking: false`, and the field now reaches the VLM: docling's picture-description request carries `reasoning_effort` in `params`. -- MCP `search_documents` and `search_documents_by_image` return the agent - rendering as text (rank, `Document ID`, `Collection` over several - databases, title, headings, passage), pictures as `ImageContent` blocks, - and the `SearchResult` list without `image_data` as structured content. - `SearchResult.format_for_agent(include_document_id=)`; - `collect_pictures` in `haiku.rag.tools.search`. +- 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`. - 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 diff --git a/claude-plugin/skills/haiku-rag/SKILL.md b/claude-plugin/skills/haiku-rag/SKILL.md index f93631a3..a0a0d57b 100644 --- a/claude-plugin/skills/haiku-rag/SKILL.md +++ b/claude-plugin/skills/haiku-rag/SKILL.md @@ -24,7 +24,8 @@ could be about the user's documents. Say so when it has nothing relevant. ## Find `search_documents` is the first call. Results come best first with the document -title, section headings and the matching passage. `filter` restricts which +title, section headings, the matched chunk's metadata when it has any, and the +passage in its section. `filter` restricts which documents are searched, `limit` how many results come back. If it misses, rephrase once or narrow with a filter before concluding the material is not there. diff --git a/docs/mcp.md b/docs/mcp.md index cbf08999..ab3b8aa7 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -109,7 +109,9 @@ repeating it. `search_documents` runs hybrid search, vector and full-text. Its text content is the rendering the in-process agents read: results best first, each with its rank, `Document ID`, `Collection` when the server covers several, the document -title, section headings and the passage. Pictures in the results follow as +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 diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index b3e8d680..c80ea1ec 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -119,9 +119,9 @@ 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, then each distinct picture - as an image block labelled with its result, and the results as structured - content without the picture bytes.""" + """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.""" import base64 total = len(results) @@ -131,6 +131,7 @@ def _search_result(results: list[SearchResult], covers_multiple: bool) -> ToolRe total=total, include_collection=covers_multiple, include_document_id=True, + include_chunk_meta=True, ) for rank, result in enumerate(results, 1) ) @@ -274,9 +275,10 @@ def _covering( Use this first for any question the documents might answer; it needs no model and is the cheapest call. Results come best first, each with its rank, `Document ID`, `Collection` when the server covers several, - the document title, section headings and the matching passage; pass - the id and collection to the document tools. Pictures in the results - follow as images, each labelled with its result. Ranks, not scores, + the document title, section headings, the matched chunk's metadata + when it has any, and the matching passage expanded to its section; + pass the id and collection to the document tools. Pictures in the + results follow as images, each labelled with its result. Ranks, not scores, are the signal: scores are not comparable across queries. If nothing relevant comes back, rephrase once or narrow with `filter` before concluding the material is absent. @@ -300,7 +302,7 @@ def _covering( ) except UnknownDatabaseError as e: raise ToolError(str(e)) from e - return _search_result(results, rag.covers_multiple) + return _search_result(await rag.expand_context(results), rag.covers_multiple) # Image-as-query tool, only registered when the configured embedder # supports image embeddings. Probed at server-build time when no Store is @@ -345,7 +347,9 @@ def _covering( ) except UnknownDatabaseError as e: raise ToolError(str(e)) from e - return _search_result(results, rag.covers_multiple) + return _search_result( + await rag.expand_context(results), rag.covers_multiple + ) @mcp.tool(annotations=_read_only("Get document")) async def get_document(document_id: str, source: str | None = None) -> Document: diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index b21231e1..fc48966f 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/models/chunk.py @@ -1,3 +1,4 @@ +import json from typing import TYPE_CHECKING, Literal from pydantic import BaseModel, PrivateAttr @@ -143,8 +144,9 @@ class SearchResult(BaseModel): consumers (UIs). Never part of ``format_for_agent`` output. ``chunk_meta`` is the anchor chunk's unparsed ``Chunk.metadata`` and does not - include the metadata of any other chunks merged with it. Never part of - ``format_for_agent`` output. + include the metadata of any other chunks merged with it. Left out of + ``format_for_agent`` output unless ``include_chunk_meta`` asks for its custom + keys. ``source`` names the database a result came from: the name from ``lancedb.databases`` or a path's stem, never a path or URI, so a location @@ -203,6 +205,7 @@ class SearchResult(BaseModel): *, include_collection: bool = False, include_document_id: bool = False, + include_chunk_meta: bool = False, ) -> str: """Format this search result for inclusion in agent context. @@ -218,6 +221,9 @@ class SearchResult(BaseModel): search spanning one collection has nothing to distinguish, whether or not that collection is named. `include_document_id` is for a reader that will fetch the document by id from the text alone. + `include_chunk_meta` renders the metadata stored with the matched + chunk beyond haiku.rag's own structural keys; on an expanded result it + locates the hit, not the whole passage. """ if rank is not None and total is not None: parts = [f"[{self.chunk_id}] [rank {rank} of {total}]"] @@ -247,6 +253,16 @@ class SearchResult(BaseModel): if primary_label: parts.append(f"Type: {primary_label}") + if include_chunk_meta: + custom = { + key: value + for key, value in self.chunk_meta.items() + if key not in ChunkMetadata.model_fields + } + if custom: + rendered = json.dumps(custom, ensure_ascii=False, sort_keys=True) + parts.append(f"Matched chunk metadata: {rendered}") + # Surface picture captions when present. Order matches the binary # attachments emitted by build_image_content_from_results, so the model # can correlate caption ↔ attached image by position (BinaryContent diff --git a/tests/test_chunk.py b/tests/test_chunk.py index 1bcf6022..a24ff97a 100644 --- a/tests/test_chunk.py +++ b/tests/test_chunk.py @@ -264,6 +264,37 @@ def test_search_result_format_for_agent_omits_chunk_meta(): assert "para_no" not in formatted +def test_search_result_format_for_agent_chunk_meta_is_opt_in(): + """A caller that asks sees the chunk's own metadata, never the structural + keys haiku.rag stores beside it.""" + result = SearchResult( + content="Some content.", + score=0.9, + chunk_id="chunk-1", + chunk_meta={ + "para_no": "12", + "doc_item_refs": ["#/texts/0"], + "page_numbers": [1], + "headings": ["Intro"], + "labels": ["paragraph"], + }, + ) + + opted = result.format_for_agent(rank=1, total=1, include_chunk_meta=True) + + assert "para_no" in opted + assert "12" in opted + assert "doc_item_refs" not in opted + assert "#/texts/0" not in opted + + structural_only = result.model_copy( + update={"chunk_meta": {"doc_item_refs": ["#/texts/0"], "page_numbers": [1]}} + ) + assert structural_only.format_for_agent( + rank=1, total=1, include_chunk_meta=True + ) == structural_only.format_for_agent(rank=1, total=1) + + def test_search_result_format_for_agent_omits_document_meta(): """Document metadata is UI plumbing, never shown to the model.""" result = SearchResult( diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 002723cc..ddc46fdb 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -10,7 +10,7 @@ from haiku.rag.mcp import _covering as _mcp_covering from haiku.rag.mcp import create_mcp_server from haiku.rag.store.models import Chunk, Document, SearchResult from haiku.rag.tools.document import DocumentInfo -from tests.multi_db.helpers import _config, _seed +from tests.multi_db.helpers import _config, _seed, _seed_expandable @pytest.fixture(autouse=True) @@ -181,6 +181,24 @@ class TestMCPReadTools: 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 + + @pytest.mark.asyncio + @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") + async def test_search_results_come_expanded(self, tmp_path): + """The passage is the hit in its section, as the in-process agents read + it, not the chunk that matched.""" + config = _config(tmp_path, ["alpha"]) + sentences = ["Gardens need water.", "Roses need pruning.", "Tulips need sun."] + await _seed_expandable(config, "alpha", sentences) + + 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) @pytest.mark.asyncio async def test_get_document(self, mcp_db):