Research-Stack/4-Infrastructure/infra/shape_index_mcp.py

#!/usr/bin/env python3
"""
MCP server for the FAMM shape-index of the Lean corpus.

Given an unsolved conjecture, classify its shape and return nearest
neighbor proofs from the indexed corpus.

Tools:
  shape_query      — find nearest neighbors by shape signature
  shape_summary    — distribution of shapes in the corpus
  shape_classify   — classify a raw Lean theorem by its FAMM shape
"""

from __future__ import annotations

import json
import os
import sys
from pathlib import Path
from typing import Any

INDEX_PATH = os.environ.get(
    "SHAPE_INDEX_PATH",
    "/home/allaun/lean_corpus/shape_index.json",
)

SERVER_NAME = "shape-index-mcp"
SERVER_VERSION = "0.1.0"

_index = None


def get_index() -> dict:
    global _index
    if _index is None:
        with open(INDEX_PATH) as f:
            _index = json.load(f)
    return _index


def json_text(data: Any) -> list[dict[str, str]]:
    return [{"type": "text", "text": json.dumps(data, indent=2, sort_keys=True)}]


def tool_shape_summary(_: dict[str, Any]) -> dict[str, Any]:
    idx = get_index()
    meta = idx.get("meta", {})
    shapes = idx.get("shape_summary", {})
    return {
        "ok": True,
        "server": SERVER_NAME,
        "version": SERVER_VERSION,
        "files_scanned": meta.get("files_scanned"),
        "theorems_indexed": meta.get("theorems_indexed"),
        "shape_distribution": shapes,
        "total_theorems": sum(shapes.values()),
    }


def tool_shape_query(args: dict[str, Any]) -> dict[str, Any]:
    idx = get_index()
    target_density = args.get("sumset_density", 0.5)
    target_vars = args.get("n_vars", 3)
    target_state = args.get("target_state", "ACCEPT")
    top_k = min(args.get("top_k", 5), 20)

    scored = []
    for t in idx["all_theorems"]:
        density = t["sidon_signature"]["sumset_density"]
        n_vars = t["constraint_graph"]["n_vars"]
        state = t["famm"]["famm_state"]
        shape = t["famm"]["rrc_shape"]

        d_density = abs(density - target_density) * 2
        d_vars = abs(n_vars - target_vars) * 0.05
        d_state = 0 if state == target_state else 1.5 if state == "HOLD" else 0.5
        distance = d_density + d_vars + d_state

        scored.append({
            "distance": round(distance, 4),
            "name": t["name"],
            "kind": t["kind"],
            "shape": shape,
            "state": state,
            "sidon_density": density,
            "n_vars": n_vars,
            "path": t["path"],
            "line": t["line"],
        })

    scored.sort(key=lambda x: x["distance"])
    return {"ok": True, "query": args, "neighbors": scored[:top_k]}


def tool_shape_classify(args: dict[str, Any]) -> dict[str, Any]:
    return {"ok": False, "error": "shape_classify requires the shape_index module at runtime; call shape_query instead"}


TOOLS = {
    "shape_summary": {
        "description": "Return shape distribution of the indexed Lean corpus.",
        "inputSchema": {"type": "object", "properties": {}},
        "handler": tool_shape_summary,
    },
    "shape_query": {
        "description": "Find nearest-neighbor theorems by Sidon/FAMM shape.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "sumset_density": {"type": "number", "default": 0.5},
                "n_vars": {"type": "integer", "default": 3},
                "target_state": {"type": "string", "default": "ACCEPT"},
                "top_k": {"type": "integer", "default": 5},
            },
        },
        "handler": tool_shape_query,
    },
    "shape_classify": {
        "description": "Classify a Lean theorem by its FAMM shape.",
        "inputSchema": {
            "type": "object",
            "required": ["content"],
            "properties": {
                "content": {"type": "string"},
                "path": {"type": "string", "default": "inline.lean"},
            },
        },
        "handler": tool_shape_classify,
    },
}


def handle(message: dict[str, Any]) -> dict[str, Any] | None:
    method = message.get("method")
    msg_id = message.get("id")
    if method == "initialize":
        return {
            "jsonrpc": "2.0", "id": msg_id,
            "result": {
                "protocolVersion": "2024-11-05",
                "capabilities": {"tools": {}},
                "serverInfo": {"name": SERVER_NAME, "version": SERVER_VERSION},
            },
        }
    if method == "tools/list":
        return {
            "jsonrpc": "2.0", "id": msg_id,
            "result": {
                "tools": [
                    {"name": name, "description": data["description"],
                     "inputSchema": data["inputSchema"]}
                    for name, data in TOOLS.items()
                ]
            },
        }
    if method == "tools/call":
        params = message.get("params") or {}
        name = params.get("name")
        args = params.get("arguments") or {}
        if name not in TOOLS:
            result = {"ok": False, "error": f"unknown tool: {name}"}
        else:
            result = TOOLS[name]["handler"](args)
        return {"jsonrpc": "2.0", "id": msg_id, "result": {"content": json_text(result)}}
    if method and method.startswith("notifications/"):
        return None
    return {"jsonrpc": "2.0", "id": msg_id, "error": {"code": -32601, "message": f"method not found: {method}"}}


def main() -> None:
    for line in sys.stdin:
        if not line.strip():
            continue
        try:
            response = handle(json.loads(line))
        except Exception as exc:
            response = {
                "jsonrpc": "2.0", "id": None,
                "error": {"code": -32000, "message": f"{type(exc).__name__}: {exc}"},
            }
        if response is not None:
            print(json.dumps(response, separators=(",", ":")), flush=True)


if __name__ == "__main__":
    main()