mnemosyne/mnemosyne/mcp_server/tools/search.py

"""Search-related MCP tools: hybrid `search` and `get_chunk` for full text."""

from __future__ import annotations

from dataclasses import asdict
from typing import Any

from asgiref.sync import sync_to_async
from django.conf import settings
from django.core.files.storage import default_storage
from fastmcp.server.context import Context

from ..context import get_mcp_user
from ..metrics import record_tool_call

DEFAULT_SEARCH_TYPES = ["vector", "fulltext", "graph"]


def register_search_tools(mcp):
    @mcp.tool
    async def search(
        query: str,
        library_uid: str | None = None,
        library_type: str | None = None,
        collection_uid: str | None = None,
        limit: int = 20,
        rerank: bool = True,
        include_images: bool = True,
        search_types: list[str] | None = None,
        # workspace_id is system-injected by Daedalus's chat path. It is
        # intentionally absent from the docstring so the calling LLM is
        # never told it exists. Whatever value the LLM produces here is
        # overwritten by Daedalus before the call reaches Mnemosyne.
        workspace_id: str | None = None,
        ctx: Context | None = None,
    ) -> dict[str, Any]:
        """Hybrid retrieval over Mnemosyne: vector + full-text + concept-graph
        candidates fused by RRF and optionally re-ranked by Synesis.

        Filters: library_uid (exact library), library_type (one of fiction,
        nonfiction, technical, music, film, art, journal, business, finance),
        or collection_uid.
        Set rerank=False to skip re-ranking. search_types defaults to all three.

        Returns ranked candidates with chunk_uid (use get_chunk for full text),
        item_uid/item_title for citation, library_type, text_preview (~500 chars),
        score, and source. Also returns matching images when include_images=True.
        """
        types = search_types or DEFAULT_SEARCH_TYPES
        with record_tool_call("search"):
            user = await get_mcp_user(ctx)
            return await sync_to_async(_run_search, thread_sensitive=True)(
                user=user,
                query=query,
                library_uid=library_uid,
                library_type=library_type,
                collection_uid=collection_uid,
                workspace_id=workspace_id,
                limit=limit,
                rerank=rerank,
                include_images=include_images,
                search_types=types,
            )

    @mcp.tool
    async def get_chunk(
        chunk_uid: str,
        # System-injected; deliberately absent from the docstring.
        workspace_id: str | None = None,
        ctx: Context | None = None,
    ) -> dict[str, Any]:
        """Fetch the full text of a chunk by its uid (typically obtained from `search`).

        Returns the chunk text plus parent item context: chunk_uid, chunk_index,
        item_uid, item_title, library_type, text. Use this when the 500-character
        text_preview from `search` isn't enough.
        """
        with record_tool_call("get_chunk"):
            return await sync_to_async(_load_chunk, thread_sensitive=True)(
                chunk_uid, workspace_id
            )


def _run_search(*, user, query, library_uid, library_type, collection_uid,
                workspace_id, limit, rerank, include_images, search_types) -> dict[str, Any]:
    from library.services.search import SearchRequest, SearchService

    req = SearchRequest(
        query=query,
        library_uid=library_uid,
        library_type=library_type,
        collection_uid=collection_uid,
        workspace_id=workspace_id,
        search_types=search_types,
        limit=limit,
        vector_top_k=getattr(settings, "SEARCH_VECTOR_TOP_K", 50),
        fulltext_top_k=getattr(settings, "SEARCH_FULLTEXT_TOP_K", 30),
        rerank=rerank,
        include_images=include_images,
    )
    service = SearchService(user=user)
    response = service.search(req)
    return {
        "query": response.query,
        "candidates": [asdict(c) for c in response.candidates],
        "images": [asdict(i) for i in response.images],
        "total_candidates": response.total_candidates,
        "search_time_ms": response.search_time_ms,
        "reranker_used": response.reranker_used,
        "reranker_model": response.reranker_model,
        "search_types_used": response.search_types_used,
    }


def _load_chunk(chunk_uid: str, workspace_id: str | None = None) -> dict[str, Any]:
    from neomodel import db

    rows, _ = db.cypher_query(
        "MATCH (l:Library)-[:CONTAINS]->(:Collection)-[:CONTAINS]->"
        "(i:Item)-[:HAS_CHUNK]->(c:Chunk {uid: $uid}) "
        "WHERE ($workspace_id IS NULL AND l.workspace_id IS NULL OR "
        "       l.workspace_id = $workspace_id) "
        "RETURN c.uid, c.chunk_index, c.chunk_s3_key, "
        "i.uid, i.title, l.library_type LIMIT 1",
        {"uid": chunk_uid, "workspace_id": workspace_id},
    )
    if not rows:
        raise ValueError(f"Chunk not found: {chunk_uid}")

    c_uid, chunk_index, chunk_s3_key, item_uid, item_title, library_type = rows[0]

    text = ""
    if chunk_s3_key:
        with default_storage.open(chunk_s3_key, "rb") as fh:
            text = fh.read().decode("utf-8", errors="replace")

    return {
        "chunk_uid": c_uid,
        "chunk_index": chunk_index,
        "item_uid": item_uid,
        "item_title": item_title,
        "library_type": library_type,
        "text": text,
    }