diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a90db87..bc98984f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,12 @@ ## [Unreleased] +### Added + +- MCP server `instructions`, `version`, and read-only `ToolAnnotations` on + every tool; every parameter carries a description. `filter` on + `search_documents` and `search_documents_by_image`. `DocumentInfo.metadata`. + ### Changed - Default models are `ollama:qwen3.8`: `ModelConfig`, `qa.model`, diff --git a/docs/mcp.md b/docs/mcp.md index ca07ce4c..fd500d07 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -67,47 +67,54 @@ With a custom database path: After restarting Claude Desktop, you can ask Claude to search your documents or answer questions using your knowledge base. -## Available Tools +## Tools -### Documents +Every tool is read-only and says so in its annotations. Each parameter carries +a description in the tool schema, so the listing below names them without +repeating it. -- **`get_document`** - Retrieve a document by ID - - `document_id` (required): The document ID - - `source` (optional): The database holding it +| Tool | Registered | Parameters | +|---|---|---| +| `search_documents` | always | `query`, `limit`, `include_images`, `filter`, `sources` | +| `search_documents_by_image` | multimodal embedder only | `image_base64`, `limit`, `include_images`, `filter`, `sources` | +| `get_document` | always | `document_id`, `source` | +| `list_documents` | always | `limit`, `offset`, `filter` | +| `ask_question` | always | `question`, `cite`, `images_base64`, `sources` | +| `analyze` | always | `question`, `filter`, `images_base64`, `sources` | -- **`list_documents`** - List documents with pagination and filtering - - `limit` (optional): Maximum number to return - - `offset` (optional): Number to skip - - `filter` (optional): SQL WHERE clause for filtering +`search_documents` runs hybrid search, vector and full-text, and returns +results best first. Scores are not comparable across queries or search types. +Rank is the signal. `include_images` attaches picture bytes as base64 PNG under +`image_data`. `search_documents_by_image` embeds the query image and searches +by vector similarity alone. -### Search +`get_document` returns a document whole, in reading order. `list_documents` +returns titles, URIs and metadata, which is how a client learns what a filter +can match. -- **`search_documents`** - Search using hybrid search (vector + full-text) - - `query` (required): Search query - - `limit` (optional): Maximum results (uses config default if not specified) - - `include_images` (optional, default `true`): Attach base64-encoded picture bytes to picture-labeled results - - `sources` (optional): The databases to search +`ask_question` runs the RAG agent on the server and returns an answer, with +citations when `cite` is set. `analyze` writes and runs Python in a sandbox +over the documents, for counting, aggregation and computation across +documents. Both cost a model call. -- **`search_documents_by_image`** - Search using an image as the query (registered only when the configured embedder supports images) - - `image_base64` (required): Base64-encoded image (PNG/JPEG bytes) - - `limit` (optional): Maximum results - - `include_images` (optional, default `true`) - - `sources` (optional): The databases to search +### Filters -### Question Answering +`filter` is a SQL WHERE clause over the document columns `id`, `uri`, `title`, +`metadata`, `created_at`, `updated_at`. `metadata` is a JSON string, so match +its keys with LIKE: -- **`ask_question`** - Ask questions about your documents - - `question` (required): The question to ask - - `cite` (optional): Include source citations (default: false) - - `images_base64` (optional): Base64-encoded images attached to the question (requires a vision-capable QA model) - - `sources` (optional): The databases to answer from +```sql +metadata LIKE '%"author": "Smith"%' +uri LIKE '%.pdf' +title = 'Q3 report' +``` -- **`analyze`** - Answer complex analytical questions via code execution - - `question` (required): The question to answer - - `filter` (optional): SQL WHERE clause to restrict document access - - `images_base64` (optional): Base64-encoded images attached to the question (requires a vision-capable analysis model) - - `sources` (optional): The databases to analyze - - Best for aggregation, computation, and multi-document analysis +### Instructions + +The server publishes `instructions` describing the knowledge base: what it +holds, when to reach for it, the collection names when it covers several, and +`prompts.domain_preamble` when set. Claude Code shows them to the model. Claude +Desktop does not, so every tool description stands on its own. ## Continuous ingestion diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 44a4f259..e94b7624 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -1,29 +1,79 @@ import asyncio from collections.abc import AsyncIterator from contextlib import AsyncExitStack, asynccontextmanager +from importlib import metadata from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Annotated from fastmcp import FastMCP from fastmcp.exceptions import ToolError +from mcp.types import ToolAnnotations +from pydantic import Field from haiku.rag.client import HaikuRAG from haiku.rag.config import AppConfig, get_config from haiku.rag.store.exceptions import UnknownDatabaseError from haiku.rag.store.models import Document, SearchResult +from haiku.rag.store.schema import DocumentMetaRecord from haiku.rag.tools.document import DocumentInfo from haiku.rag.utils import format_citations if TYPE_CHECKING: from haiku.rag.client.scope import DatabaseScope +_FILTER_COLUMNS = ", ".join(DocumentMetaRecord.model_fields) + +Filter = Annotated[ + str | None, + Field( + description=( + f"SQL WHERE clause over the document columns {_FILTER_COLUMNS}, " + "restricting which documents are used. `metadata` is a JSON string, " + 'so match its keys with LIKE: metadata LIKE \'%"author": "Smith"%\'. ' + "Also uri LIKE '%.pdf', title = 'Q3 report'." + ) + ), +] +Sources = Annotated[ + list[str] | None, + Field(description="Collections to use, by name. All of them by default."), +] + + +def _read_only(title: str) -> ToolAnnotations: + return ToolAnnotations(title=title, readOnlyHint=True, openWorldHint=False) + + +def _decode_image(image_base64: str) -> bytes: + import base64 + + return base64.b64decode(image_base64, validate=True) + def _decode_images(images_base64: list[str] | None) -> list[bytes] | None: if not images_base64: return None - import base64 + return [_decode_image(b64) for b64 in images_base64] - return [base64.b64decode(b64, validate=True) for b64 in images_base64] + +def _instructions(scope: "DatabaseScope", config: AppConfig) -> str: + """What the server is for, naming no tools: the client has every tool's + description from the listing.""" + lines = [ + "haiku-rag is the user's knowledge base: documents they ingested, " + "searchable by meaning and keyword, readable whole, answered with " + "citations, or computed across documents.", + "Use it whenever a question could be answered from those documents, " + "before answering from memory, and say when it had nothing relevant.", + ] + if scope.covers_multiple: + lines.append( + f"It holds several collections: {', '.join(scope.names)}. Results " + "and citations name theirs in `source`; pass `sources` to use a subset." + ) + if config.prompts.domain_preamble: + lines.append(config.prompts.domain_preamble) + return "\n".join(lines) def create_mcp_server( @@ -86,27 +136,45 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: finally: client = None - mcp = FastMCP("haiku-rag", lifespan=lifespan) + mcp = FastMCP( + "haiku-rag", + instructions=_instructions(scope, config), + version=metadata.version("haiku.rag-slim"), + lifespan=lifespan, + ) - @mcp.tool() + @mcp.tool(annotations=_read_only("Search documents")) async def search_documents( query: str, limit: int | None = None, include_images: bool = True, - sources: list[str] | None = None, + filter: Filter = None, + sources: Sources = None, ) -> list[SearchResult]: - """Search the RAG system for documents using hybrid search (vector similarity + full-text search). + """Search the knowledge base by meaning and keyword. - When include_images is True (default) and a picture-labeled chunk is - in the result set, ``SearchResult.image_data`` carries base64-encoded - PNG bytes keyed by self_ref. Set to False to omit the bytes from the - response (smaller JSON payload for plain-text consumers). - ``sources`` names the databases to search, all of them by default. + 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 + the document's id, title and collection, the section headings and the + matching passage. Scores are not comparable across queries, so read + the order, not the numbers. If nothing relevant comes back, rephrase + once or narrow with `filter` before concluding the material is absent. + + Args: + query: What to look for, in natural language or keywords. + limit: How many results to return; the server's configured default + when omitted. + include_images: Attach the bytes of pictures in the results as + base64 PNG under `image_data`. False for a smaller response. """ try: rag = await _client() return await rag.search( - query, limit=limit, include_images=include_images, sources=sources + query, + limit=limit, + filter=filter, + include_images=include_images, + sources=sources, ) except UnknownDatabaseError as e: raise ToolError(str(e)) from e @@ -121,45 +189,57 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: if get_embedder(config).supports_images: - @mcp.tool() + @mcp.tool(annotations=_read_only("Search documents by image")) async def search_documents_by_image( image_base64: str, limit: int | None = None, include_images: bool = True, - sources: list[str] | None = None, + filter: Filter = None, + sources: Sources = None, ) -> list[SearchResult]: - """Search the RAG system using an image as the query. + """Search the knowledge base with an image as the query. - ``image_base64`` is a base64-encoded image (PNG/JPEG bytes). The - image is embedded via the configured multimodal embedder and the - chunks table is searched vector-only. ``include_images`` controls - whether picture bytes are attached to picture-labeled results. - ``sources`` names the databases to search, all of them by default. + Use this when the question is about a picture rather than words. + The image is embedded and matched against document text and + figures by vector similarity alone. Results have the shape of + `search_documents` results. + + Args: + image_base64: The query image, PNG or JPEG bytes as base64. + limit: How many results to return; the server's configured + default when omitted. + include_images: Attach the bytes of pictures in the results as + base64 PNG under `image_data`. False for a smaller response. """ - import base64 - - try: - raw = base64.b64decode(image_base64) - except Exception: - return [] try: + raw = _decode_image(image_base64) rag = await _client() return await rag.search( - raw, limit=limit, include_images=include_images, sources=sources + raw, + limit=limit, + filter=filter, + include_images=include_images, + sources=sources, ) except UnknownDatabaseError as e: raise ToolError(str(e)) from e except Exception: return [] - @mcp.tool() + @mcp.tool(annotations=_read_only("Get document")) async def get_document( document_id: str, source: str | None = None ) -> Document | None: - """Get a document by its ID. + """Read one document whole, in reading order. - ``source`` names the database holding it; without one every database - is asked. + Use this after a search when a passage is not enough. Returns the + document's content, title, uri and metadata. Ids come from search + results and `list_documents`. + + Args: + document_id: The document's id. + source: The collection holding it. Without one every collection + is asked. """ try: rag = await _client() @@ -169,18 +249,21 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: except Exception: return None - @mcp.tool() + @mcp.tool(annotations=_read_only("List documents")) async def list_documents( limit: int | None = None, offset: int | None = None, - filter: str | None = None, + filter: Filter = None, ) -> list[DocumentInfo]: - """List all documents with optional pagination and filtering. + """List what the knowledge base holds. + + Use this to see which documents exist, their titles, URIs and + metadata, and so what a `filter` can match. Not a search: it returns + no passages. Args: - limit: Maximum number of documents to return. - offset: Number of documents to skip. - filter: Optional SQL WHERE clause to filter documents. + limit: How many documents to return. + offset: How many documents to skip, for paging. """ try: rag = await _client() @@ -193,30 +276,32 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: uri=doc.uri or "", created=doc.created_at.strftime("%Y-%m-%d"), source=doc.source, + metadata=doc.metadata, ) for doc in documents ] except Exception: return [] - @mcp.tool() + @mcp.tool(annotations=_read_only("Ask a question")) async def ask_question( question: str, cite: bool = False, images_base64: list[str] | None = None, - sources: list[str] | None = None, + sources: Sources = None, ) -> str: - """Ask a question using the QA agent. + """Answer a question from the documents with a retrieval agent. + + Use this when the user wants an answer rather than material to read. + It runs a model on the server and is slower than a search. Returns + the answer, followed by citations to the passages it rests on when + `cite` is set. Args: - question: The question to ask. - cite: Whether to include citations in the response. - images_base64: Base64-encoded images attached to the question - (requires a vision-capable QA model). - sources: The databases to answer from, all of them by default. - - Returns: - The answer as a string. + question: The question, in natural language. + cite: Append citations to the answer. + images_base64: Images to attach to the question, PNG or JPEG + bytes as base64. Needs a vision-capable model on the server. """ try: images = _decode_images(images_base64) @@ -232,28 +317,24 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: except Exception as e: return f"Error answering question: {e!s}" - @mcp.tool() + @mcp.tool(annotations=_read_only("Analyze documents")) async def analyze( question: str, - filter: str | None = None, + filter: Filter = None, images_base64: list[str] | None = None, - sources: list[str] | None = None, + sources: Sources = None, ) -> str: - """Answer complex questions using the analysis capability. + """Compute an answer across documents with code. - Use this for questions requiring computation, aggregation, or - structural traversal across documents. The capability can write and - execute Python code in a sandboxed interpreter. + Use this for counting, aggregation, comparison across many documents + or arithmetic over tables, where reading passages is not enough. A + model writes and runs Python in a sandbox over the selected documents. + It is the slowest tool. Returns the answer as text. Args: - question: The question to answer. - filter: Optional SQL WHERE clause to filter documents. - images_base64: Base64-encoded images attached to the question - (requires a vision-capable analysis model). - sources: The databases to analyze, all of them by default. - - Returns: - The answer as a string. + question: The question, in natural language. + images_base64: Images to attach to the question, PNG or JPEG + bytes as base64. Needs a vision-capable model on the server. """ try: images = _decode_images(images_base64) diff --git a/haiku_rag_slim/haiku/rag/tools/document.py b/haiku_rag_slim/haiku/rag/tools/document.py index 1458083a..960d5194 100644 --- a/haiku_rag_slim/haiku/rag/tools/document.py +++ b/haiku_rag_slim/haiku/rag/tools/document.py @@ -28,6 +28,7 @@ class DocumentInfo(BaseModel): uri: str created: str source: str | None = None + metadata: dict = {} class DocumentListResponse(BaseModel): diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 06bc32d5..82799789 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -56,6 +56,7 @@ async def mcp_db(temp_db_path): "Artificial intelligence is transforming industries worldwide.", title="AI Overview", uri="test://ai-overview", + metadata={"author": "Ada"}, ) await rag.create_document( "Machine learning is a subset of artificial intelligence.", @@ -104,6 +105,21 @@ class TestMCPReadTools: results = await search(query="artificial intelligence", limit=1) assert len(results) == 1 + @pytest.mark.asyncio + @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") + async def test_search_documents_with_filter(self, mcp_db): + from fastmcp import Client + + async with Client(create_mcp_server(mcp_db)) as client: + result = await client.call_tool( + "search_documents", + {"query": "artificial intelligence", "filter": "title = 'ML Basics'"}, + ) + + results = result.structured_content["result"] + assert results + assert {r["document_title"] for r in results} == {"ML Basics"} + @pytest.mark.asyncio @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") async def test_search_documents_preserves_chunk_meta_through_serialization( @@ -202,6 +218,103 @@ class TestMCPReadTools: assert len(results) == 1 assert results[0].title == "AI Overview" + @pytest.mark.asyncio + @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") + async def test_list_documents_carries_metadata(self, mcp_db): + from fastmcp import Client + + async with Client(create_mcp_server(mcp_db)) as client: + result = await client.call_tool("list_documents", {}) + + [overview] = [ + d + for d in result.structured_content["result"] + if d["title"] == "AI Overview" + ] + assert overview["metadata"] == {"author": "Ada"} + + +@pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") +class TestMCPDescribesItself: + """What a client learns from initialize and list_tools, over the wire.""" + + @pytest.mark.asyncio + async def test_instructions_and_version_are_set(self, mcp_db): + from importlib import metadata + + from fastmcp import Client + + async with Client(create_mcp_server(mcp_db)) as client: + init = client.initialize_result + + assert init.instructions + assert init.serverInfo.version == metadata.version("haiku.rag-slim") + + @pytest.mark.asyncio + async def test_instructions_name_the_collections_when_covering_several( + self, two_dbs + ): + from fastmcp import Client + + from haiku.rag.client.scope import DatabaseScope + + async with Client(_covering_all(two_dbs)) as client: + covering_both = client.initialize_result.instructions + one = DatabaseScope.resolve(two_dbs, database_name="alpha") + async with Client(_mcp_covering(one, two_dbs)) as client: + covering_one = client.initialize_result.instructions + + assert "alpha" in covering_both + assert "beta" in covering_both + assert "beta" not in covering_one + + @pytest.mark.asyncio + async def test_instructions_carry_the_domain_preamble(self, mcp_db): + from fastmcp import Client + + from haiku.rag.config import get_config + + config = get_config().model_copy(deep=True) + config.prompts.domain_preamble = "Everything here is about zebras." + + async with Client(create_mcp_server(mcp_db, config=config)) as client: + with_preamble = client.initialize_result.instructions + async with Client(create_mcp_server(mcp_db)) as client: + without = client.initialize_result.instructions + + assert "Everything here is about zebras." in with_preamble + assert "zebras" not in without + + @pytest.mark.asyncio + async def test_every_tool_is_annotated_read_only(self, mcp_db, multimodal_embedder): + from fastmcp import Client + + async with Client(create_mcp_server(mcp_db)) as client: + tools = await client.list_tools() + + assert len(tools) == 6 + for tool in tools: + assert tool.annotations is not None, tool.name + assert tool.annotations.readOnlyHint is True, tool.name + assert tool.annotations.openWorldHint is False, tool.name + assert tool.annotations.title, tool.name + + @pytest.mark.asyncio + async def test_every_parameter_is_described(self, mcp_db, multimodal_embedder): + from fastmcp import Client + + async with Client(create_mcp_server(mcp_db)) as client: + tools = await client.list_tools() + + undescribed = [ + f"{tool.name}.{name}" + for tool in tools + for name, schema in tool.inputSchema.get("properties", {}).items() + if not schema.get("description") + ] + assert len(tools) == 6 + assert undescribed == [] + class TestMCPToolSet: @pytest.mark.asyncio @@ -380,13 +493,35 @@ class TestMCPImageQuery: png = b"\x89PNG\r\n\x1a\n" results = await search_by_image( - image_base64=base64.b64encode(png).decode("ascii"), sources=["alpha"] + image_base64=base64.b64encode(png).decode("ascii"), + filter="uri LIKE 'x%'", + sources=["alpha"], ) assert results == [] assert seen["query"] == png + assert seen["filter"] == "uri LIKE 'x%'" assert seen["sources"] == ["alpha"] + @pytest.mark.asyncio + async def test_image_query_rejects_characters_outside_the_alphabet( + self, mcp_db, multimodal_embedder, monkeypatch + ): + """A lenient decoder would drop the stray characters and search.""" + searched = False + + async def fake_search(self, query, **kwargs): + nonlocal searched + searched = True + return [] + + monkeypatch.setattr(HaikuRAG, "search", fake_search) + mcp = create_mcp_server(mcp_db) + search_by_image = await _get_tool(mcp, "search_documents_by_image") + + assert await search_by_image(image_base64="AAAA!!!!") == [] + assert not searched + @pytest.mark.asyncio async def test_image_query_returns_empty_on_invalid_base64( self, mcp_db, multimodal_embedder