mcp_server.py

"""
mcp_server.py
=============
Exposes RAGPipeline as MCP tools, so any MCP-aware host (Claude Desktop,
Claude Code, or any custom MCP client) can use this MCP-docs assistant as a
tool -- without knowing the REST API, just by calling "ask_mcp_docs".

Can be run two ways:

1. Standalone, stdio transport (local MCP host, e.g. Claude Desktop config):
       python mcp_server.py

2. Mounted inside FastAPI, streamable-http transport (remote MCP clients):
       via main.py -> app.mount("/mcp", mcp.streamable_http_app())

Both modes use the same underlying `rag_core.get_pipeline()` singleton,
so the vector store / cache / checkpointer stay shared whether a request
comes in via REST or via MCP.
"""

from __future__ import annotations

from mcp.server.fastmcp import FastMCP

import rag_core
from rag_pipeline.logging_setup import logger

mcp = FastMCP(
    name="mcp-docs-assistant",
    instructions=(
        "Ask questions about the Model Context Protocol (MCP) -- architecture, building "
        "servers/clients, tools/resources/prompts spec, security best practices. Answers are "
        "grounded in the official MCP documentation via retrieval-augmented generation, with "
        "guardrails, PII masking, and hallucination checking applied."
    ),
)


@mcp.tool(
    name="ask_mcp_docs",
    description=(
        "Ask a natural-language question about the Model Context Protocol (MCP) and get a "
        "grounded answer generated from the official MCP docs. Supports multi-turn "
        "conversations: pass the same thread_id across calls so follow-up questions "
        "(e.g. 'show a Python example of that') resolve correctly. Input is safety-checked "
        "and PII-masked; the answer is checked for hallucination against the retrieved context."
    ),
)
async def ask_mcp_docs(question: str, thread_id: str | None = None) -> dict:
    pipeline = await rag_core.get_pipeline()
    return await pipeline.chat(question, thread_id)


@mcp.tool(
    name="search_mcp_docs",
    description=(
        "Retrieve and rerank the top MCP documentation chunks most relevant to a query, "
        "without generation or guardrails. Use this when you (the calling agent) want the "
        "raw grounded context to reason over yourself, instead of a generated answer."
    ),
)
async def search_mcp_docs(query: str, top_n: int = 5) -> dict:
    pipeline = await rag_core.get_pipeline()
    chunks = await pipeline.search(query, top_n=top_n)
    return {"query": query, "chunks": chunks}


@mcp.tool(
    name="get_conversation_history",
    description="Get the stored (user, assistant) turn history for a given thread_id.",
)
async def get_conversation_history(thread_id: str) -> dict:
    pipeline = await rag_core.get_pipeline()
    return {"thread_id": thread_id, "history": await pipeline.get_history(thread_id)}


@mcp.tool(
    name="cache_stats",
    description="Get semantic-cache stats (entry count, TTL, max size) for observability.",
)
async def cache_stats() -> dict:
    pipeline = await rag_core.get_pipeline()
    return pipeline.cache_stats()


def main():
    """Standalone entrypoint -- stdio transport, for local MCP hosts like
    Claude Desktop (see README.md for the config snippet)."""
    logger.info("Starting mcp-docs-assistant MCP server (stdio transport)...")
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()