f45ed90b33
haiku-rag mcp --no-agents reaches create_mcp_server(agents=False): ask_question and analyze are not registered and the instructions drop the line describing them. qa.model always has a default, so a server without a usable model cannot be detected from configuration; the flag is how an operator says so. Refs #599
537 lines
20 KiB
Python
537 lines
20 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 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.store.exceptions import UnknownDatabaseError
|
|
from haiku.rag.store.models import Document, SearchResult
|
|
from haiku.rag.store.models.document_item import DocumentItem
|
|
from haiku.rag.store.schema import DocumentMetaRecord
|
|
from haiku.rag.tools.document import DocumentInfo, DocumentSection, OutlineNode
|
|
from haiku.rag.tools.search import collect_pictures
|
|
from haiku.rag.utils import format_citations
|
|
|
|
if TYPE_CHECKING:
|
|
from typing import Any
|
|
|
|
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, agents: bool) -> 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."
|
|
]
|
|
if agents:
|
|
lines.append(
|
|
"Questions can be answered from them with citations, or computed "
|
|
"across them 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 "
|
|
"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 _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."""
|
|
import base64
|
|
|
|
total = len(results)
|
|
text = "\n\n".join(
|
|
result.format_for_agent(
|
|
rank=rank,
|
|
total=total,
|
|
include_collection=covers_multiple,
|
|
include_document_id=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"),
|
|
mimeType="image/png",
|
|
)
|
|
)
|
|
return ToolResult(
|
|
content=content,
|
|
structured_content={
|
|
"result": [
|
|
result.model_dump(mode="json", exclude={"image_data"})
|
|
for result in results
|
|
]
|
|
},
|
|
)
|
|
|
|
|
|
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,
|
|
agents: bool = True,
|
|
) -> 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.
|
|
agents: Register `ask_question` and `analyze`, which run a model on
|
|
the server.
|
|
"""
|
|
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, agents
|
|
)
|
|
|
|
|
|
def _covering(
|
|
scope: "DatabaseScope", config: AppConfig, agents: bool = True
|
|
) -> 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, agents),
|
|
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,
|
|
) -> 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 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,
|
|
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()
|
|
try:
|
|
await _check_filter(rag, filter, sources)
|
|
results = await rag.search(
|
|
query,
|
|
limit=limit,
|
|
filter=filter,
|
|
include_images=include_images,
|
|
sources=sources,
|
|
)
|
|
except UnknownDatabaseError as e:
|
|
raise ToolError(str(e)) from e
|
|
return _search_result(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()
|
|
try:
|
|
await _check_filter(rag, filter, sources)
|
|
results = await rag.search(
|
|
raw,
|
|
limit=limit,
|
|
filter=filter,
|
|
include_images=include_images,
|
|
sources=sources,
|
|
)
|
|
except UnknownDatabaseError as e:
|
|
raise ToolError(str(e)) from e
|
|
return _search_result(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()
|
|
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
|
|
|
|
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()
|
|
try:
|
|
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)
|
|
except UnknownDatabaseError as e:
|
|
raise ToolError(str(e)) from e
|
|
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"]
|
|
return DocumentSection(
|
|
id=node["self_ref"],
|
|
title=node["title"],
|
|
page_numbers=node["page_numbers"],
|
|
content="\n\n".join(
|
|
item.text
|
|
for item in items
|
|
if start <= item.position < end and 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()
|
|
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
|
|
]
|
|
|
|
if agents:
|
|
|
|
@mcp.tool(annotations=_read_only("Ask a question"))
|
|
async def ask_question(
|
|
question: str,
|
|
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.
|
|
|
|
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:
|
|
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 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
|