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