254 lines
9.2 KiB
Python
254 lines
9.2 KiB
Python
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from fastmcp import FastMCP
|
|
|
|
from haiku.rag.agents.research.models import ResearchReport
|
|
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 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.
|
|
"""
|
|
mcp = FastMCP("haiku-rag")
|
|
|
|
# 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:
|
|
async with HaikuRAG(db_path, config=config) as rag:
|
|
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:
|
|
async with HaikuRAG(db_path, config=config) as rag:
|
|
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:
|
|
async with HaikuRAG(db_path, config=config) as rag:
|
|
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:
|
|
async with HaikuRAG(db_path, config=config) as rag:
|
|
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:
|
|
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
|
|
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.
|
|
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:
|
|
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
|
|
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:
|
|
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
|
|
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:
|
|
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
|
|
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,
|
|
) -> str:
|
|
"""Ask a question using the QA agent.
|
|
|
|
Args:
|
|
question: The question to ask.
|
|
cite: Whether to include citations in the response.
|
|
|
|
Returns:
|
|
The answer as a string.
|
|
"""
|
|
try:
|
|
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
|
|
answer, citations = await rag.ask(question)
|
|
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 research_question(
|
|
question: str,
|
|
) -> ResearchReport | None:
|
|
"""Run multi-agent research to investigate a complex question.
|
|
|
|
The research process uses multiple agents to plan, search, evaluate, and synthesize
|
|
information iteratively until confidence threshold is met or max iterations reached.
|
|
|
|
Args:
|
|
question: The research question to investigate.
|
|
|
|
Returns:
|
|
A research report with findings, or None if an error occurred.
|
|
"""
|
|
try:
|
|
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
|
|
return await rag.research(question=question)
|
|
except Exception:
|
|
return None
|
|
|
|
@mcp.tool()
|
|
async def analyze(
|
|
question: str,
|
|
document: str | None = None,
|
|
filter: str | None = None,
|
|
) -> str:
|
|
"""Answer complex questions using code execution (analysis agent).
|
|
|
|
Use this for questions requiring computation, aggregation, or
|
|
complex traversal across documents. The agent can write Python
|
|
code to search, analyze, and compute answers.
|
|
|
|
Args:
|
|
question: The question to answer.
|
|
document: Optional document ID or title to pre-load for analysis.
|
|
filter: Optional SQL WHERE clause to filter documents.
|
|
|
|
Returns:
|
|
The answer as a string.
|
|
"""
|
|
try:
|
|
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
|
|
documents = [document] if document else None
|
|
result = await rag.analyze(question, documents=documents, filter=filter)
|
|
return result.answer
|
|
except Exception as e:
|
|
return f"Error running analysis agent: {e!s}"
|
|
|
|
return mcp
|