basicmachines-co-basic-memory/src/basic_memory/mcp/tools/chatgpt_tools.py

"""ChatGPT-compatible MCP tools for Basic Memory.

These adapters expose Basic Memory's search/fetch functionality using the exact
tool names and response structure OpenAI's MCP clients expect: each call returns
a list containing a single `{"type": "text", "text": "{...json...}"}` item.
"""

import json
from typing import Any, Dict, List, Optional

from fastmcp import Context
from loguru import logger

from basic_memory.mcp.server import mcp
from basic_memory.mcp.tools.read_note import read_note
from basic_memory.mcp.tools.search import search_notes
from basic_memory.schemas.search import SearchResponse, SearchResult


def _format_search_results_for_chatgpt(
    results: SearchResponse | list[SearchResult] | list[dict[str, Any]] | dict[str, Any],
) -> List[Dict[str, Any]]:
    """Format search results according to ChatGPT's expected schema.

    Returns a list of result objects with id, title, and url fields.
    """
    if isinstance(results, SearchResponse):
        raw_results: list[SearchResult] | list[dict[str, Any]] = results.results
    elif isinstance(results, dict):
        nested_results = results.get("results")
        raw_results = nested_results if isinstance(nested_results, list) else []
    else:
        raw_results = results

    formatted_results = []

    for result in raw_results:
        if isinstance(result, SearchResult):
            title = result.title
            permalink = result.permalink
        elif isinstance(result, dict):
            title = result.get("title")
            permalink = result.get("permalink")
        else:
            raise TypeError(f"Unexpected result type: {type(result).__name__}")

        formatted_result = {
            "id": permalink or f"doc-{len(formatted_results)}",
            "title": title if isinstance(title, str) and title.strip() else "Untitled",
            "url": permalink or "",
        }
        formatted_results.append(formatted_result)

    return formatted_results


def _format_document_for_chatgpt(
    content: str, identifier: str, title: Optional[str] = None
) -> Dict[str, Any]:
    """Format document content according to ChatGPT's expected schema.

    Returns a document object with id, title, text, url, and metadata fields.
    """
    # Extract title from markdown content if not provided
    if not title and isinstance(content, str):
        lines = content.split("\n")
        if lines and lines[0].startswith("# "):
            title = lines[0][2:].strip()
        else:
            title = identifier.split("/")[-1].replace("-", " ").title()

    # Ensure title is never None
    if not title:
        title = "Untitled Document"

    # Handle error cases
    if isinstance(content, str) and content.lstrip().startswith("# Note Not Found"):
        return {
            "id": identifier,
            "title": title or "Document Not Found",
            "text": content,
            "url": identifier,
            "metadata": {"error": "Document not found"},
        }

    return {
        "id": identifier,
        "title": title or "Untitled Document",
        "text": content,
        "url": identifier,
        "metadata": {"format": "markdown"},
    }


@mcp.tool(
    description="Search for content across the knowledge base",
    annotations={"readOnlyHint": True, "openWorldHint": False},
)
async def search(
    query: str,
    context: Context | None = None,
) -> List[Dict[str, Any]]:
    """ChatGPT/OpenAI MCP search adapter returning a single text content item.

    Args:
        query: Search query (full-text syntax supported by `search_notes`)
        context: Optional FastMCP context passed through for auth/session data

    Returns:
        List with one dict: `{ "type": "text", "text": "{...JSON...}" }`
        where the JSON body contains `results`, `total_count`, and echo of `query`.
    """
    logger.info(f"ChatGPT search request: query='{query}'")

    try:
        # Let search_notes resolve the default project via get_project_client(),
        # which works in both local mode (ConfigManager) and cloud mode (database).
        results = await search_notes(
            query=query,
            page=1,
            page_size=10,
            output_format="json",
            context=context,
        )

        # Handle string error responses from search_notes
        if isinstance(results, str):
            logger.warning(f"Search failed with error: {results[:100]}...")
            search_results = {
                "results": [],
                "error": "Search failed",
                "error_details": results[:500],  # Truncate long error messages
            }
        else:
            # Format successful results for ChatGPT
            raw_results = results.get("results", []) if isinstance(results, dict) else []
            formatted_results = _format_search_results_for_chatgpt(raw_results)
            search_results = {
                "results": formatted_results,
                "total_count": len(raw_results),  # Use actual count from results
                "query": query,
            }
            logger.info(f"Search completed: {len(formatted_results)} results returned")

        # Return in MCP content array format as required by OpenAI
        return [{"type": "text", "text": json.dumps(search_results, ensure_ascii=False)}]

    except Exception as e:
        logger.error(f"ChatGPT search failed for query '{query}': {e}")
        error_results = {
            "results": [],
            "error": "Internal search error",
            "error_message": str(e)[:200],
        }
        return [{"type": "text", "text": json.dumps(error_results, ensure_ascii=False)}]


@mcp.tool(
    description="Fetch the full contents of a search result document",
    annotations={"readOnlyHint": True, "openWorldHint": False},
)
async def fetch(
    id: str,
    context: Context | None = None,
) -> List[Dict[str, Any]]:
    """ChatGPT/OpenAI MCP fetch adapter returning a single text content item.

    Args:
        id: Document identifier (permalink, title, or memory URL)
        context: Optional FastMCP context passed through for auth/session data

    Returns:
        List with one dict: `{ "type": "text", "text": "{...JSON...}" }`
        where the JSON body includes `id`, `title`, `text`, `url`, and metadata.
    """
    logger.info(f"ChatGPT fetch request: id='{id}'")

    try:
        # Let read_note resolve the default project via get_project_client(),
        # which works in both local mode (ConfigManager) and cloud mode (database).
        content = str(
            await read_note(
                identifier=id,
                page=1,
                page_size=10,
                context=context,
            )
        )

        # Format the document for ChatGPT
        document = _format_document_for_chatgpt(content, id)

        logger.info(f"Fetch completed: id='{id}', content_length={len(document.get('text', ''))}")

        # Return in MCP content array format as required by OpenAI
        return [{"type": "text", "text": json.dumps(document, ensure_ascii=False)}]

    except Exception as e:
        logger.error(f"ChatGPT fetch failed for id '{id}': {e}")
        error_document = {
            "id": id,
            "title": "Fetch Error",
            "text": f"Failed to fetch document: {str(e)[:200]}",
            "url": id,
            "metadata": {"error": "Fetch failed"},
        }
        return [{"type": "text", "text": json.dumps(error_document, ensure_ascii=False)}]