8a4d488a72
All ten tool bodies opened their own `HaikuRAG`, so every tool call paid a connection open and, on object storage, refetched the index the previous call had just cached. The client is now opened once, lazily so that calling a tool function directly still works, and eagerly from the lifespan so an unopenable database fails startup instead of every call. Teardown clears the cached client in a finally, since `_lifespan_manager` can be re-entered and would otherwise hand out a closed connection, including when the close itself fails. `delete_document` no longer opens its own connection with `skip_validation=True`. Keeping it separate broke consistency once connections became long-lived: the delete committed on one connection while reads served from another, which with a 30s consistency interval showed the deleted document as still present. A connection always sees its own writes, so sharing one is what makes delete visible to the next read. So the server no longer opts out of embedding-config validation. Drift that validation rejects now fails MCP startup, where before the server started and only `delete_document` worked while every read returned empty. Same-dimension identity drift still starts a read-only server, matching every other read verb. Delete under drift is now a CLI operation; CLAUDE.md and the CHANGELOG record it.
281 lines
9.8 KiB
Python
281 lines
9.8 KiB
Python
import asyncio
|
|
from collections.abc import AsyncIterator
|
|
from contextlib import AsyncExitStack, asynccontextmanager
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from fastmcp import FastMCP
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config import AppConfig, Config
|
|
from haiku.rag.store.models import Document, SearchResult
|
|
from haiku.rag.tools.document import DocumentInfo
|
|
from haiku.rag.utils import format_citations
|
|
|
|
|
|
def _decode_images(images_base64: list[str] | None) -> list[bytes] | None:
|
|
if not images_base64:
|
|
return None
|
|
import base64
|
|
|
|
return [base64.b64decode(b64, validate=True) for b64 in images_base64]
|
|
|
|
|
|
def create_mcp_server(
|
|
db_path: Path, config: AppConfig = Config, read_only: bool = False
|
|
) -> FastMCP:
|
|
"""Create an MCP server with the specified database path.
|
|
|
|
Args:
|
|
db_path: Path to the database file.
|
|
config: Configuration to use.
|
|
read_only: If True, write tools (add_document_*, delete_document) are not registered.
|
|
"""
|
|
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(db_path, config=config, read_only=read_only)
|
|
)
|
|
return client
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(_server: FastMCP) -> AsyncIterator[None]:
|
|
# Open eagerly so an unopenable database fails startup rather than
|
|
# every tool call.
|
|
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
|
|
|
|
mcp = FastMCP("haiku-rag", lifespan=lifespan)
|
|
|
|
# Write tools - only registered when not in read-only mode
|
|
if not read_only:
|
|
|
|
@mcp.tool()
|
|
async def add_document_from_file(
|
|
file_path: str,
|
|
metadata: dict[str, Any] | None = None,
|
|
title: str | None = None,
|
|
) -> str | None:
|
|
"""Add a document to the RAG system from a file path."""
|
|
try:
|
|
rag = await _client()
|
|
result = await rag.create_document_from_source(
|
|
Path(file_path), title=title, metadata=metadata or {}
|
|
)
|
|
# Handle both single document and list of documents (directories)
|
|
if isinstance(result, list):
|
|
return result[0].id if result else None
|
|
return result.id
|
|
except Exception:
|
|
return None
|
|
|
|
@mcp.tool()
|
|
async def add_document_from_url(
|
|
url: str, metadata: dict[str, Any] | None = None, title: str | None = None
|
|
) -> str | None:
|
|
"""Add a document to the RAG system from a URL."""
|
|
try:
|
|
rag = await _client()
|
|
result = await rag.create_document_from_source(
|
|
url, title=title, metadata=metadata or {}
|
|
)
|
|
# Handle both single document and list of documents
|
|
if isinstance(result, list):
|
|
return result[0].id if result else None
|
|
return result.id
|
|
except Exception:
|
|
return None
|
|
|
|
@mcp.tool()
|
|
async def add_document_from_text(
|
|
content: str,
|
|
uri: str | None = None,
|
|
metadata: dict[str, Any] | None = None,
|
|
title: str | None = None,
|
|
) -> str | None:
|
|
"""Add a document to the RAG system from text content."""
|
|
try:
|
|
rag = await _client()
|
|
document = await rag.create_document(
|
|
content, uri, title=title, metadata=metadata or {}
|
|
)
|
|
return document.id
|
|
except Exception:
|
|
return None
|
|
|
|
@mcp.tool()
|
|
async def delete_document(document_id: str) -> bool:
|
|
"""Delete a document by its ID."""
|
|
try:
|
|
rag = await _client()
|
|
return await rag.delete_document(document_id)
|
|
except Exception:
|
|
return False
|
|
|
|
# Read tools - always registered
|
|
@mcp.tool()
|
|
async def search_documents(
|
|
query: str, limit: int | None = None, include_images: bool = True
|
|
) -> list[SearchResult]:
|
|
"""Search the RAG system for documents using hybrid search (vector similarity + full-text search).
|
|
|
|
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).
|
|
"""
|
|
try:
|
|
rag = await _client()
|
|
return await rag.search(query, limit=limit, include_images=include_images)
|
|
except Exception:
|
|
return []
|
|
|
|
# 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()
|
|
async def search_documents_by_image(
|
|
image_base64: str,
|
|
limit: int | None = None,
|
|
include_images: bool = True,
|
|
) -> list[SearchResult]:
|
|
"""Search the RAG system using 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.
|
|
"""
|
|
import base64
|
|
|
|
try:
|
|
raw = base64.b64decode(image_base64)
|
|
except Exception:
|
|
return []
|
|
try:
|
|
rag = await _client()
|
|
return await rag.search(raw, limit=limit, include_images=include_images)
|
|
except Exception:
|
|
return []
|
|
|
|
@mcp.tool()
|
|
async def get_document(document_id: str) -> Document | None:
|
|
"""Get a document by its ID."""
|
|
try:
|
|
rag = await _client()
|
|
return await rag.get_document_by_id(document_id)
|
|
except Exception:
|
|
return None
|
|
|
|
@mcp.tool()
|
|
async def list_documents(
|
|
limit: int | None = None,
|
|
offset: int | None = None,
|
|
filter: str | None = None,
|
|
) -> list[DocumentInfo]:
|
|
"""List all documents with optional pagination and filtering.
|
|
|
|
Args:
|
|
limit: Maximum number of documents to return.
|
|
offset: Number of documents to skip.
|
|
filter: Optional SQL WHERE clause to filter documents.
|
|
"""
|
|
try:
|
|
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"),
|
|
)
|
|
for doc in documents
|
|
]
|
|
except Exception:
|
|
return []
|
|
|
|
@mcp.tool()
|
|
async def ask_question(
|
|
question: str,
|
|
cite: bool = False,
|
|
images_base64: list[str] | None = None,
|
|
) -> str:
|
|
"""Ask a question using the QA agent.
|
|
|
|
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).
|
|
|
|
Returns:
|
|
The answer as a string.
|
|
"""
|
|
try:
|
|
images = _decode_images(images_base64)
|
|
rag = await _client()
|
|
answer, citations = await rag.ask(question, images=images)
|
|
if cite and citations:
|
|
answer += "\n\n" + format_citations(citations)
|
|
return answer
|
|
except Exception as e:
|
|
return f"Error answering question: {e!s}"
|
|
|
|
@mcp.tool()
|
|
async def analyze(
|
|
question: str,
|
|
filter: str | None = None,
|
|
images_base64: list[str] | None = None,
|
|
) -> str:
|
|
"""Answer complex questions using the analysis capability.
|
|
|
|
Use this for questions requiring computation, aggregation, or
|
|
structural traversal across documents. The capability can write and
|
|
execute Python code in a sandboxed interpreter.
|
|
|
|
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).
|
|
|
|
Returns:
|
|
The answer as a string.
|
|
"""
|
|
try:
|
|
images = _decode_images(images_base64)
|
|
rag = await _client()
|
|
result = await rag.analyze(question, filter=filter, images=images)
|
|
return result.answer
|
|
except Exception as e:
|
|
return f"Error running analysis capability: {e!s}"
|
|
|
|
return mcp
|