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
This commit is contained in:
Yiorgis Gozadinos 2026-09-04 14:12:12 +03:00
parent 494294046f
commit 2582f2c05a
No known key found for this signature in database
7 changed files with 92 additions and 19 deletions

View file

@ -32,12 +32,13 @@
- `processing.conversion_options.picture_description.model` defaults to - `processing.conversion_options.picture_description.model` defaults to
`enable_thinking: false`, and the field now reaches the VLM: docling's `enable_thinking: false`, and the field now reaches the VLM: docling's
picture-description request carries `reasoning_effort` in `params`. picture-description request carries `reasoning_effort` in `params`.
- MCP `search_documents` and `search_documents_by_image` return the agent - MCP `search_documents` and `search_documents_by_image` expand results to
rendering as text (rank, `Document ID`, `Collection` over several their section (`HaikuRAG.expand_context`) and return the agent rendering
databases, title, headings, passage), pictures as `ImageContent` blocks, as text (rank, `Document ID`, `Collection` over several databases, title,
and the `SearchResult` list without `image_data` as structured content. headings, the matched chunk's metadata, passage), pictures as
`SearchResult.format_for_agent(include_document_id=)`; `ImageContent` blocks, and the `SearchResult` list without `image_data` as
`collect_pictures` in `haiku.rag.tools.search`. 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. - 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

View file

@ -24,7 +24,8 @@ could be about the user's documents. Say so when it has nothing relevant.
## Find ## Find
`search_documents` is the first call. Results come best first with the document `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, 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 rephrase once or narrow with a filter before concluding the material is not
there. there.

View file

@ -109,7 +109,9 @@ repeating it.
`search_documents` runs hybrid search, vector and full-text. Its text content `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 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 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 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. The structured content is the
`SearchResult` list without picture bytes. Scores are not comparable across `SearchResult` list without picture bytes. Scores are not comparable across

View file

@ -119,9 +119,9 @@ 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, then each distinct picture """Results as the in-process agents read them, plus the matched chunk's
as an image block labelled with its result, and the results as structured metadata, then each distinct picture as an image block labelled with its
content without the picture bytes.""" result, and the results as structured content without the picture bytes."""
import base64 import base64
total = len(results) total = len(results)
@ -131,6 +131,7 @@ def _search_result(results: list[SearchResult], covers_multiple: bool) -> ToolRe
total=total, total=total,
include_collection=covers_multiple, include_collection=covers_multiple,
include_document_id=True, include_document_id=True,
include_chunk_meta=True,
) )
for rank, result in enumerate(results, 1) 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 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 no model and is the cheapest call. Results come best first, each with
its rank, `Document ID`, `Collection` when the server covers several, its rank, `Document ID`, `Collection` when the server covers several,
the document title, section headings and the matching passage; pass the document title, section headings, the matched chunk's metadata
the id and collection to the document tools. Pictures in the results when it has any, and the matching passage expanded to its section;
follow as images, each labelled with its result. Ranks, not scores, 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 are the signal: scores are not comparable across queries. If nothing
relevant comes back, rephrase once or narrow with `filter` before relevant comes back, rephrase once or narrow with `filter` before
concluding the material is absent. concluding the material is absent.
@ -300,7 +302,7 @@ def _covering(
) )
except UnknownDatabaseError as e: except UnknownDatabaseError as e:
raise ToolError(str(e)) from 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 # Image-as-query tool, only registered when the configured embedder
# supports image embeddings. Probed at server-build time when no Store is # supports image embeddings. Probed at server-build time when no Store is
@ -345,7 +347,9 @@ def _covering(
) )
except UnknownDatabaseError as e: except UnknownDatabaseError as e:
raise ToolError(str(e)) from 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")) @mcp.tool(annotations=_read_only("Get document"))
async def get_document(document_id: str, source: str | None = None) -> Document: async def get_document(document_id: str, source: str | None = None) -> Document:

View file

@ -1,3 +1,4 @@
import json
from typing import TYPE_CHECKING, Literal from typing import TYPE_CHECKING, Literal
from pydantic import BaseModel, PrivateAttr from pydantic import BaseModel, PrivateAttr
@ -143,8 +144,9 @@ class SearchResult(BaseModel):
consumers (UIs). Never part of ``format_for_agent`` output. consumers (UIs). Never part of ``format_for_agent`` output.
``chunk_meta`` is the anchor chunk's unparsed ``Chunk.metadata`` and does not ``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 include the metadata of any other chunks merged with it. Left out of
``format_for_agent`` output. ``format_for_agent`` output unless ``include_chunk_meta`` asks for its custom
keys.
``source`` names the database a result came from: the name from ``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 ``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_collection: bool = False,
include_document_id: bool = False, include_document_id: bool = False,
include_chunk_meta: bool = False,
) -> str: ) -> str:
"""Format this search result for inclusion in agent context. """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 search spanning one collection has nothing to distinguish, whether or
not that collection is named. `include_document_id` is for a reader not that collection is named. `include_document_id` is for a reader
that will fetch the document by id from the text alone. 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: if rank is not None and total is not None:
parts = [f"[{self.chunk_id}] [rank {rank} of {total}]"] parts = [f"[{self.chunk_id}] [rank {rank} of {total}]"]
@ -247,6 +253,16 @@ class SearchResult(BaseModel):
if primary_label: if primary_label:
parts.append(f"Type: {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 # Surface picture captions when present. Order matches the binary
# attachments emitted by build_image_content_from_results, so the model # attachments emitted by build_image_content_from_results, so the model
# can correlate caption ↔ attached image by position (BinaryContent # can correlate caption ↔ attached image by position (BinaryContent

View file

@ -264,6 +264,37 @@ def test_search_result_format_for_agent_omits_chunk_meta():
assert "para_no" not in formatted 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(): def test_search_result_format_for_agent_omits_document_meta():
"""Document metadata is UI plumbing, never shown to the model.""" """Document metadata is UI plumbing, never shown to the model."""
result = SearchResult( result = SearchResult(

View file

@ -10,7 +10,7 @@ from haiku.rag.mcp import _covering as _mcp_covering
from haiku.rag.mcp import create_mcp_server from haiku.rag.mcp import create_mcp_server
from haiku.rag.store.models import Chunk, Document, SearchResult from haiku.rag.store.models import Chunk, Document, SearchResult
from haiku.rag.tools.document import DocumentInfo 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) @pytest.fixture(autouse=True)
@ -181,6 +181,24 @@ class TestMCPReadTools:
assert any( assert any(
r["chunk_meta"] == {"fake-metadata-for-testing": "42"} for r in results 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 @pytest.mark.asyncio
async def test_get_document(self, mcp_db): async def test_get_document(self, mcp_db):