danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Describe the MCP server and its tools to clients

Browse source 
FastMCP gets instructions, the haiku.rag-slim version and a lifespan.
Every tool carries read-only ToolAnnotations with a title, a description
that says when to use it, and a description on every parameter. The
filter description lists the document columns from DocumentMetaRecord
and how to match metadata; filter is accepted by both search tools.
One strict base64 decoder serves every image parameter.
DocumentInfo carries metadata.

Refs #599
This commit is contained in:
Yiorgis Gozadinos 2026-09-04 11:31:45 +03:00
parent 40d40bcbf2
commit 88cd2b1d6c
No known key found for this signature in database
5 changed files with 327 additions and 97 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
CHANGELOG.md
View file

@ -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`,

71
docs/mcp.md
View file

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

209
haiku_rag_slim/haiku/rag/mcp.py
View file

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

1
haiku_rag_slim/haiku/rag/tools/document.py
View file

@ -28,6 +28,7 @@ class DocumentInfo(BaseModel):
uri: str
created: str
source: str | None = None
metadata: dict = {}
class DocumentListResponse(BaseModel):

137
tests/test_mcp.py
View file

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