110adf23ee
build_toc stored item positions, while the sandbox instructions describe item_range as a slice into items.jsonl; a gap in positions pulled the next heading into a section. Ranges are now indices into the position-ordered items, and get_document_section slices by index too. docs/mcp.md names the tools that take sources.
498 lines
19 KiB
Python
498 lines
19 KiB
Python
import asyncio
|
|
import logging
|
|
import re
|
|
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.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
|
|
|
|
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, read_only_hint=True, open_world_hint=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
|
|
|
|
|
|
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:
|
|
# The engine lists its own columns too, lance internals among them.
|
|
reason = re.sub(r"\s*Valid fields are .*", "", str(e), flags=re.DOTALL)
|
|
raise ToolError(
|
|
f"Invalid filter {filter!r}: {reason.rstrip('. ')}. "
|
|
f"Columns: {_FILTER_COLUMNS}."
|
|
) 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 "
|
|
"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, 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."""
|
|
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,
|
|
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 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,
|
|
) -> 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()
|
|
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(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()
|
|
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(
|
|
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()
|
|
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"]
|
|
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()
|
|
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("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` (the section tree, each node with an item_range
|
|
slice into items.jsonl). 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 and chunk_meta. `await list_documents()`
|
|
returns dicts with id, title, uri, created_at, source and metadata.
|
|
Modules: json, re, math, pathlib. Not available: generators, class
|
|
inheritance, match statements, decorators, collections.
|
|
|
|
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:
|
|
await _check_filter(rag, filter, sources)
|
|
result = await sandbox.execute(code)
|
|
except UnknownDatabaseError as e:
|
|
raise ToolError(str(e)) from e
|
|
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
|