2a6d72171d
FastMCP masks unexpected exceptions and logs the traceback server-side, so paths and provider URLs never cross the transport. Expected failures raise ToolError with a message: unknown document, unknown collection, invalid filter, invalid base64, and agent failures naming only the exception type. No tool returns an empty value or an error string on failure any more. A filter is validated on its own before the read that would use it, with a filtered count on one of the selected databases: that is the query engine rejecting the filter and nothing else, so its message (columns and the statement) can be forwarded, while a ValueError raised later in the read stays masked and no database outside the selection is opened. Refs #599
380 lines
14 KiB
Python
380 lines
14 KiB
Python
import asyncio
|
|
import logging
|
|
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 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
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_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
|
|
|
|
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 _decode_images(images_base64: list[str] | None) -> list[bytes] | None:
|
|
if not images_base64:
|
|
return None
|
|
return [_decode_image(b64) for b64 in images_base64]
|
|
|
|
|
|
async def _check_filter(
|
|
rag: HaikuRAG, filter: str | None, sources: list[str] | None = None
|
|
) -> None:
|
|
"""Evaluate a filter on its own before the read that would use it.
|
|
|
|
A filtered count on one selected database runs the same predicate on the
|
|
same table and nothing else, so a ValueError here is the query engine
|
|
rejecting the filter; its message names columns and the statement, never
|
|
a location. A ValueError raised later in the read stays masked. Only the
|
|
selection is touched: every database shares the schema, so one suffices.
|
|
"""
|
|
if filter is None:
|
|
return
|
|
selected = await rag.clients_covering(sources)
|
|
if not selected:
|
|
return
|
|
try:
|
|
await selected[0].count_documents(filter=filter)
|
|
except ValueError as e:
|
|
raise ToolError(f"Invalid filter {filter!r}: {e}") 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, 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(
|
|
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 and citations 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
|
|
|
|
# Masking keeps paths and provider URLs out of an unexpected error's text;
|
|
# the traceback goes to the server log. A ToolError reaches the client as is.
|
|
mcp = FastMCP(
|
|
"haiku-rag",
|
|
instructions=_instructions(scope, config),
|
|
version=metadata.version("haiku.rag-slim"),
|
|
lifespan=lifespan,
|
|
mask_error_details=True,
|
|
)
|
|
|
|
@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,
|
|
) -> list[SearchResult]:
|
|
"""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
|
|
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.
|
|
"""
|
|
rag = await _client()
|
|
try:
|
|
await _check_filter(rag, filter, sources)
|
|
return await rag.search(
|
|
query,
|
|
limit=limit,
|
|
filter=filter,
|
|
include_images=include_images,
|
|
sources=sources,
|
|
)
|
|
except UnknownDatabaseError as e:
|
|
raise ToolError(str(e)) from e
|
|
|
|
# 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,
|
|
) -> list[SearchResult]:
|
|
"""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: Attach the bytes of pictures in the results as
|
|
base64 PNG under `image_data`. False for a smaller response.
|
|
"""
|
|
raw = _decode_image(image_base64)
|
|
rag = await _client()
|
|
try:
|
|
await _check_filter(rag, filter, sources)
|
|
return await rag.search(
|
|
raw,
|
|
limit=limit,
|
|
filter=filter,
|
|
include_images=include_images,
|
|
sources=sources,
|
|
)
|
|
except UnknownDatabaseError as e:
|
|
raise ToolError(str(e)) from e
|
|
|
|
@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()
|
|
try:
|
|
document = await rag.get_document_by_id(document_id, source)
|
|
except UnknownDatabaseError as e:
|
|
raise ToolError(str(e)) from e
|
|
if document is None:
|
|
raise ToolError(f"No document with id {document_id!r}")
|
|
return document
|
|
|
|
@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()
|
|
await _check_filter(rag, filter)
|
|
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("Ask a question"))
|
|
async def ask_question(
|
|
question: str,
|
|
cite: bool = False,
|
|
images_base64: list[str] | None = None,
|
|
sources: Sources = None,
|
|
) -> str:
|
|
"""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, 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.
|
|
"""
|
|
images = _decode_images(images_base64)
|
|
rag = await _client()
|
|
try:
|
|
answer, citations = await rag.ask(question, images=images, sources=sources)
|
|
except UnknownDatabaseError as e:
|
|
raise ToolError(str(e)) from e
|
|
except Exception as e:
|
|
logger.exception("ask_question failed")
|
|
raise ToolError(f"ask_question failed: {type(e).__name__}") from e
|
|
if cite and citations:
|
|
answer += "\n\n" + format_citations(
|
|
citations, include_source=rag.covers_multiple
|
|
)
|
|
return answer
|
|
|
|
@mcp.tool(annotations=_read_only("Analyze documents"))
|
|
async def analyze(
|
|
question: str,
|
|
filter: Filter = None,
|
|
images_base64: list[str] | None = None,
|
|
sources: Sources = None,
|
|
) -> str:
|
|
"""Compute an answer across documents with code.
|
|
|
|
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, 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.
|
|
"""
|
|
images = _decode_images(images_base64)
|
|
rag = await _client()
|
|
try:
|
|
result = await rag.analyze(
|
|
question, filter=filter, images=images, sources=sources
|
|
)
|
|
except UnknownDatabaseError as e:
|
|
raise ToolError(str(e)) from e
|
|
except Exception as e:
|
|
logger.exception("analyze failed")
|
|
raise ToolError(f"analyze failed: {type(e).__name__}") from e
|
|
return result.answer
|
|
|
|
return mcp
|