virtuoso_mcp/server_fastmcp.py

"""FastMCP transport wrapper for virtuoso tool logic.

This module keeps the existing Virtuoso tool implementations untouched
and exposes them through the MCP SDK transport.
"""

from fastapi import FastAPI
from fastapi import HTTPException
from pydantic import BaseModel
from mcp.server.fastmcp import FastMCP
from mcp.server.transport_security import TransportSecuritySettings

from virtuoso_mcp import (
    TOOLS,
    TOOL_DOCS,
    VIRTUOSO_ENDPOINT,
    SPARQL_DEFAULT_LIMIT,
    SPARQL_MAX_LIMIT,
    ALLOW_EXAMPLE_LOAD,
)


mcp = FastMCP(
    "virtuoso",
    transport_security=TransportSecuritySettings(
        enable_dns_rebinding_protection=False
    ),
)


def _wrap_tool(tool_name: str):
    def _tool(input: dict | None = None):
        payload = input or {}
        return TOOLS[tool_name](payload)

    _tool.__name__ = f"tool_{tool_name}"
    _tool.__doc__ = TOOL_DOCS.get(tool_name, "")
    return _tool


for _name in sorted(TOOLS.keys()):
    mcp.add_tool(
        _wrap_tool(_name),
        name=_name,
        description=TOOL_DOCS.get(_name, ""),
    )


app = FastAPI(title="Virtuoso MCP Server")
app.mount("/mcp", mcp.sse_app())


class LegacyToolRequest(BaseModel):
    tool: str
    input: dict = {}


@app.post("/rpc")
def legacy_rpc(req: LegacyToolRequest):
    if req.tool not in TOOLS:
        raise HTTPException(status_code=400, detail=f"Unknown tool: {req.tool}")
    try:
        result = TOOLS[req.tool](req.input or {})
        return {
            "status": "ok",
            "tool": req.tool,
            "description": TOOL_DOCS.get(req.tool, ""),
            "result": result,
        }
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc))


@app.get("/")
def root():
    return {
        "status": "MCP server running",
        "transport": "fastmcp+sse",
        "mount": "/mcp",
        "tools": sorted(TOOLS.keys()),
        "virtuoso": VIRTUOSO_ENDPOINT,
        "guardrails": {
            "default_limit": SPARQL_DEFAULT_LIMIT,
            "max_limit": SPARQL_MAX_LIMIT,
            "allow_example_load": ALLOW_EXAMPLE_LOAD,
        },
    }


@app.get("/health")
def health():
    return root()