4b1ec096c6
Every failure reaches the client as an MCP error carrying its message; mask_error_details is set off explicitly, since FastMCP also reads it from the environment. The masking goes, and with it the filter pre-check that ran a count before every filtered call and the UnknownDatabaseError translations that existed only to survive it. Explicit domain errors stay.
460 lines
18 KiB
Python
460 lines
18 KiB
Python
import asyncio
|
|
import base64
|
|
from collections.abc import AsyncIterator
|
|
from contextlib import AsyncExitStack, asynccontextmanager
|
|
from importlib import metadata
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Annotated
|
|
|
|
from fastmcp import FastMCP
|
|
from fastmcp.exceptions import ToolError
|
|
from fastmcp.tools import ToolResult
|
|
from mcp.types import ContentBlock, ImageContent, TextContent, ToolAnnotations
|
|
from pydantic import Field
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config import AppConfig, get_config
|
|
from haiku.rag.context import build_toc
|
|
from haiku.rag.sandbox import AnalysisContext, Sandbox, recovery_hint
|
|
from haiku.rag.store.models import Document, SearchResult
|
|
from haiku.rag.store.schema import DocumentMetaRecord
|
|
from haiku.rag.tools.document import DocumentInfo, DocumentSection, OutlineNode
|
|
from haiku.rag.tools.search import collect_pictures
|
|
|
|
if TYPE_CHECKING:
|
|
from typing import Any
|
|
|
|
from haiku.rag.client.scope import DatabaseScope
|
|
from haiku.rag.store.models.document_item import DocumentItem
|
|
|
|
_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, read_only_hint=True, open_world_hint=False)
|
|
|
|
|
|
def _decode_image(image_base64: str) -> bytes:
|
|
try:
|
|
return base64.b64decode(image_base64, validate=True)
|
|
except ValueError as e:
|
|
# binascii.Error for characters outside the alphabet or bad padding,
|
|
# ValueError itself for non-ASCII input.
|
|
raise ToolError("Invalid base64 image") from e
|
|
|
|
|
|
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 or section by section, "
|
|
"or computed across with code."
|
|
]
|
|
lines.append(
|
|
"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 "
|
|
"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 _search_result(results: list[SearchResult], covers_multiple: bool) -> ToolResult:
|
|
"""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. No structured content: a client given both shows the model the
|
|
JSON and drops the text, or shows both."""
|
|
total = len(results)
|
|
text = "\n\n".join(
|
|
result.format_for_agent(
|
|
rank=rank,
|
|
total=total,
|
|
include_collection=covers_multiple,
|
|
include_document_id=True,
|
|
include_chunk_meta=True,
|
|
)
|
|
for rank, result in enumerate(results, 1)
|
|
)
|
|
content: list[ContentBlock] = [
|
|
TextContent(type="text", text=text or "No results found.")
|
|
]
|
|
pictures, _ = collect_pictures(results)
|
|
for source, chunk_id, self_ref, picture in pictures:
|
|
collection = f" in {source}" if covers_multiple and source else ""
|
|
content.append(
|
|
TextContent(
|
|
type="text",
|
|
text=f"Picture {self_ref} of search result [{chunk_id}]{collection}",
|
|
)
|
|
)
|
|
content.append(
|
|
ImageContent(
|
|
type="image",
|
|
data=base64.b64encode(picture.data).decode("ascii"),
|
|
mime_type="image/png",
|
|
)
|
|
)
|
|
return ToolResult(content=content)
|
|
|
|
|
|
def _node(toc: "dict[str, Any]") -> OutlineNode:
|
|
return OutlineNode(
|
|
id=toc["self_ref"],
|
|
title=toc["title"],
|
|
level=toc["level"],
|
|
page_numbers=toc["page_numbers"],
|
|
children=[_node(child) for child in toc["children"]],
|
|
)
|
|
|
|
|
|
def _find(toc: list["dict[str, Any]"], section_id: str) -> "dict[str, Any] | None":
|
|
for node in toc:
|
|
if node["self_ref"] == section_id:
|
|
return node
|
|
found = _find(node["children"], section_id)
|
|
if found is not None:
|
|
return found
|
|
return None
|
|
|
|
|
|
def create_mcp_server(
|
|
db_path: Path | None = None, config: AppConfig | None = None
|
|
) -> FastMCP:
|
|
"""Create an MCP server over the databases the configuration places.
|
|
|
|
Args:
|
|
db_path: Path to the database file, where `config` places none; or
|
|
None to serve the databases the configuration places. Beside
|
|
`lancedb.databases` a path raises `AmbiguousDatabaseError`.
|
|
config: Configuration to use.
|
|
"""
|
|
from haiku.rag.client.scope import DatabaseScope
|
|
|
|
config = config if config is not None else get_config()
|
|
return _covering(DatabaseScope.resolve(config, database_path=db_path), config)
|
|
|
|
|
|
def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
|
|
"""An MCP server over databases someone already resolved.
|
|
|
|
Internal, as ``HaikuRAG._covering`` is: the public factory takes a path and
|
|
resolves it, which is its own job. A caller that resolved already passes the
|
|
scope, so the configured name survives, which results carry as ``source``.
|
|
"""
|
|
client: HaikuRAG | None = None
|
|
stack = AsyncExitStack()
|
|
client_lock = asyncio.Lock()
|
|
|
|
async def _client() -> HaikuRAG:
|
|
"""The server's client, opened once.
|
|
|
|
Opening cost is per connection, and on object storage the first vector
|
|
query loads the index into the session cache, so a client per tool call
|
|
pays that repeatedly.
|
|
"""
|
|
nonlocal client
|
|
async with client_lock:
|
|
if client is None:
|
|
client = await stack.enter_async_context(
|
|
HaikuRAG._covering(scope, config, read_only=True)
|
|
)
|
|
return client
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(_server: FastMCP) -> AsyncIterator[None]:
|
|
# Opened eagerly: an unopenable database fails at startup.
|
|
nonlocal client
|
|
await _client()
|
|
try:
|
|
yield
|
|
finally:
|
|
# The lifespan can be re-entered; without the reset the next cycle
|
|
# hands out the closed client, including when aclose itself fails.
|
|
try:
|
|
await stack.aclose()
|
|
finally:
|
|
client = None
|
|
|
|
# Explicit: the setting is also read from the environment, and the contract
|
|
# is that every failure reaches the client with its message.
|
|
mcp = FastMCP(
|
|
"haiku-rag",
|
|
instructions=_instructions(scope, config),
|
|
version=metadata.version("haiku.rag-slim"),
|
|
lifespan=lifespan,
|
|
mask_error_details=False,
|
|
)
|
|
|
|
@mcp.tool(annotations=_read_only("Search documents"))
|
|
async def search_documents(
|
|
query: str,
|
|
limit: int | None = None,
|
|
include_images: bool = True,
|
|
filter: Filter = None,
|
|
sources: Sources = None,
|
|
) -> ToolResult:
|
|
"""Search the knowledge base by meaning and keyword.
|
|
|
|
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, 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.
|
|
|
|
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: Return the pictures in the results as images.
|
|
False for a smaller response.
|
|
"""
|
|
rag = await _client()
|
|
results = await rag.search(
|
|
query,
|
|
limit=limit,
|
|
filter=filter,
|
|
include_images=include_images,
|
|
sources=sources,
|
|
)
|
|
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
|
|
# open, so there is no cached embedder to read; this is the one place
|
|
# outside Store that builds one.
|
|
from haiku.rag.embeddings import get_embedder
|
|
|
|
if get_embedder(config).supports_images:
|
|
|
|
@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,
|
|
filter: Filter = None,
|
|
sources: Sources = None,
|
|
) -> ToolResult:
|
|
"""Search the knowledge base with an image as the query.
|
|
|
|
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: Return the pictures in the results as images.
|
|
False for a smaller response.
|
|
"""
|
|
raw = _decode_image(image_base64)
|
|
rag = await _client()
|
|
results = await rag.search(
|
|
raw,
|
|
limit=limit,
|
|
filter=filter,
|
|
include_images=include_images,
|
|
sources=sources,
|
|
)
|
|
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:
|
|
"""Read one document whole, in reading order.
|
|
|
|
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.
|
|
"""
|
|
rag = await _client()
|
|
document = await rag.get_document_by_id(document_id, source)
|
|
if document is None:
|
|
raise ToolError(f"No document with id {document_id!r}")
|
|
return document
|
|
|
|
async def _items_of(document_id: str, source: str | None) -> list["DocumentItem"]:
|
|
"""A document's items in reading order, from the database holding it."""
|
|
rag = await _client()
|
|
document = await rag.get_document_by_id(document_id, source)
|
|
if document is None:
|
|
raise ToolError(f"No document with id {document_id!r}")
|
|
owner = await rag.reader_for(source or document.source)
|
|
assert owner is not None, "a stored document names its database"
|
|
return await owner.document_item_repository.get_all_items(document_id)
|
|
|
|
@mcp.tool(annotations=_read_only("Document outline"))
|
|
async def get_document_outline(
|
|
document_id: str, source: str | None = None
|
|
) -> list[OutlineNode]:
|
|
"""The heading tree of a document, with page numbers.
|
|
|
|
Use this on a long document to see its structure before reading, then
|
|
pass a node's `id` to `get_document_section`. Returns the headings
|
|
nested by level; an empty list means the document has no headings,
|
|
so read it with `get_document`.
|
|
|
|
Args:
|
|
document_id: The document's id.
|
|
source: The collection holding it. Without one every collection
|
|
is asked.
|
|
"""
|
|
return [
|
|
_node(toc) for toc in build_toc(await _items_of(document_id, source), {})
|
|
]
|
|
|
|
@mcp.tool(annotations=_read_only("Document section"))
|
|
async def get_document_section(
|
|
document_id: str, section_id: str, source: str | None = None
|
|
) -> DocumentSection:
|
|
"""The text of one section of a document, subsections included.
|
|
|
|
Use this to read a part of a long document instead of the whole.
|
|
`section_id` is a node `id` from `get_document_outline`. Returns the
|
|
section's heading, page numbers and text in reading order, up to the
|
|
next heading of the same or a higher level.
|
|
|
|
Args:
|
|
document_id: The document's id.
|
|
section_id: The `id` of a node in the document's outline.
|
|
source: The collection holding it. Without one every collection
|
|
is asked.
|
|
"""
|
|
items = await _items_of(document_id, source)
|
|
node = _find(build_toc(items, {}), section_id)
|
|
if node is None:
|
|
raise ToolError(f"No section {section_id!r} in document {document_id!r}")
|
|
start, end = node["item_range"]
|
|
ordered = sorted(items, key=lambda item: item.position)
|
|
return DocumentSection(
|
|
id=node["self_ref"],
|
|
title=node["title"],
|
|
page_numbers=node["page_numbers"],
|
|
content="\n\n".join(item.text for item in ordered[start:end] if item.text),
|
|
)
|
|
|
|
@mcp.tool(annotations=_read_only("List documents"))
|
|
async def list_documents(
|
|
limit: int | None = None,
|
|
offset: int | None = None,
|
|
filter: Filter = None,
|
|
) -> list[DocumentInfo]:
|
|
"""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: How many documents to return.
|
|
offset: How many documents to skip, for paging.
|
|
"""
|
|
rag = await _client()
|
|
documents = await rag.list_documents(limit, offset, filter)
|
|
return [
|
|
DocumentInfo(
|
|
id=doc.id,
|
|
title=doc.title or "Untitled",
|
|
uri=doc.uri or "",
|
|
created=doc.created_at.strftime("%Y-%m-%d"),
|
|
source=doc.source,
|
|
metadata=doc.metadata,
|
|
)
|
|
for doc in documents
|
|
]
|
|
|
|
@mcp.tool(annotations=_read_only("Run code over the documents"))
|
|
async def execute_code(
|
|
code: str, filter: Filter = None, sources: Sources = None
|
|
) -> str:
|
|
"""Run a Python program over the documents and return what it printed.
|
|
|
|
Use this when the answer is a count, an aggregate, a comparison across
|
|
many documents, a lookup by document or chunk metadata, or a pattern
|
|
over whole documents: whatever a search cannot rank. The program runs
|
|
in a sandboxed interpreter on the server. Each call is one program,
|
|
nothing carries over between calls, and `print` is the only output.
|
|
|
|
Inside the program, `/documents/{document_id}/` holds `metadata.json`
|
|
(id, title, uri, created_at, metadata), `content.txt` (the whole text),
|
|
`items.jsonl` (one item per line: self_ref, label, text, page_numbers,
|
|
heading_level, chunk_ids), `chunks.jsonl` (one chunk per line: chunk_id,
|
|
metadata) and `toc.json` (`doc_id`, `title`, `tree`; each node has
|
|
self_ref, level, title, page_numbers, item_range as a slice into
|
|
items.jsonl, chunk_ids and children; an empty tree means no headings).
|
|
Read files with `Path.read_text()` or `open()`; a file object cannot be
|
|
iterated, use `.readlines()`. `await search(query, limit=10)` returns
|
|
dicts with chunk_id, content, document_id, document_title, document_uri,
|
|
source, score, page_numbers, headings, doc_item_refs, labels,
|
|
picture_refs (the doc_item_refs that are pictures) and chunk_meta.
|
|
`await list_documents()` returns dicts with id, title, uri, created_at,
|
|
source and metadata. Both see the documents `filter` and `sources`
|
|
select. Useful modules include json, re, math, pathlib, datetime,
|
|
collections, itertools, functools and dataclasses; decimal and
|
|
statistics do not exist. No generator functions, match statements or
|
|
class inheritance.
|
|
Files are read-only, there is no network, a call has a time limit named
|
|
in the error when it is hit, and output past a size is truncated.
|
|
|
|
Map a title or URI to a document id with one `list_documents()` call
|
|
rather than reading every `metadata.json`. The files carry no `source`,
|
|
so over several collections group by the `source` of `list_documents()`
|
|
rows. For a known document's structure read its `toc.json` before
|
|
searching: `search()` ranks across every document. A hit's
|
|
`doc_item_refs` are `self_ref` values in `items.jsonl`, which places it
|
|
in its section. `chunk_ids` on items and `chunk_id` in `chunks.jsonl`
|
|
join the two files; they are not citations.
|
|
|
|
Args:
|
|
code: The program. Use `await` on search and list_documents.
|
|
"""
|
|
rag = await _client()
|
|
sandbox = Sandbox._covering(
|
|
scope, config, AnalysisContext(filter=filter, sources=sources), rag=rag
|
|
)
|
|
try:
|
|
result = await sandbox.execute(code)
|
|
finally:
|
|
await sandbox.close()
|
|
if not result.success:
|
|
raise ToolError(
|
|
f"{result.stderr}{recovery_hint(result.stderr)}"
|
|
f"\n\nOutput: {result.stdout}"
|
|
)
|
|
return result.stdout or "No output."
|
|
|
|
return mcp
|