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

#!/usr/bin/env python3
"""
MCP server for the DeepSeek-Prover-V2-7B integration.

Exposes generate, prove, prove_batch, and status tools.
Runs as stdio JSON-RPC 2.0 (MCP protocol).

Environment:
  DEEPSEEK_PROVER_BACKEND — "ollama" | "huggingface" | "deepseek_api"
  OLLAMA_URL             — default http://localhost:11434
  DEEPSEEK_API_KEY       — for deepseek_api backend
  LAKE_WORKDIR           — default /home/allaun/Research Stack/0-Core-Formalism/lean/Semantics
"""

from __future__ import annotations

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

SRC_DIR = Path(__file__).resolve().parents[1] / "shim"
sys.path.insert(0, str(SRC_DIR))

try:
    from lean_proof import DeepSeekProver, ProverEngine
except ImportError as exc:
    print(
        json.dumps(
            {
                "jsonrpc": "2.0",
                "id": None,
                "error": {
                    "code": -32000,
                    "message": f"lean_proof package not found: {exc}",
                },
            }
        ),
        flush=True,
    )
    sys.exit(1)

SERVER_NAME = "deepseek-prover-mcp"
SERVER_VERSION = "0.1.0"
DEFAULT_WORKDIR = str(
    Path(__file__).resolve().parents[2]
    / "0-Core-Formalism"
    / "lean"
    / "Semantics"
)


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


def _engine() -> ProverEngine:
    backend = os.environ.get("DEEPSEEK_PROVER_BACKEND", "ollama")
    prover = DeepSeekProver(backend=backend)
    return ProverEngine(prover)


def tool_generate(args: dict[str, Any]) -> dict[str, Any]:
    theorem = args.get("theorem_statement", "")
    context = args.get("context", "")
    error_feedback = args.get("error_feedback", "")
    num_candidates = args.get("num_candidates", 4)
    model = args.get("model", "")

    prover = DeepSeekProver(
        backend=os.environ.get("DEEPSEEK_PROVER_BACKEND", "ollama")
    )
    if model:
        prover.config.model = model
    if num_candidates:
        prover.config.num_candidates = int(num_candidates)

    candidates = prover.generate(theorem, context, error_feedback)
    return {
        "ok": True,
        "count": len(candidates),
        "candidates": [
            {
                "code": c.code,
                "model": c.model,
                "latency_ms": c.latency_ms,
                "iteration": c.iteration,
            }
            for c in candidates
        ],
    }


def tool_prove(args: dict[str, Any]) -> dict[str, Any]:
    theorem = args.get("theorem_statement", "")
    context = args.get("context", "")
    max_iterations = args.get("max_iterations", 5)
    lake_workdir = args.get("lake_workdir") or os.environ.get(
        "LAKE_WORKDIR", DEFAULT_WORKDIR
    )
    emit_receipt = args.get("emit_receipt", True)

    engine = _engine()
    result = engine.prove(
        theorem_statement=theorem,
        context=context,
        max_iterations=int(max_iterations),
        lake_workdir=lake_workdir,
        emit_receipt=bool(emit_receipt),
    )
    return {
        "ok": result.passed,
        "theorem": result.theorem,
        "passed": result.passed,
        "iterations": result.iterations,
        "model": result.model,
        "latency_ms": result.latency_ms,
        "code": result.code,
        "receipt_sha256": result.receipt_sha256,
        "compile_log": result.compile_log[-2000:] if result.compile_log else "",
    }


def tool_prove_batch(args: dict[str, Any]) -> dict[str, Any]:
    theorems = args.get("theorems", [])
    context = args.get("context", "")
    lake_workdir = args.get("lake_workdir") or os.environ.get(
        "LAKE_WORKDIR", DEFAULT_WORKDIR
    )
    engine = _engine()
    results = engine.prove_batch(
        theorems=[(t.get("name", ""), t.get("statement", "")) for t in theorems],
        context=context,
        lake_workdir=lake_workdir,
    )
    return {
        "ok": True,
        "count": len(results),
        "results": [
            {
                "theorem": r.theorem,
                "passed": r.passed,
                "iterations": r.iterations,
                "model": r.model,
                "latency_ms": r.latency_ms,
            }
            for r in results
        ],
    }


def tool_status(_: dict[str, Any]) -> dict[str, Any]:
    backend = os.environ.get("DEEPSEEK_PROVER_BACKEND", "ollama")
    workdir = os.environ.get("LAKE_WORKDIR", DEFAULT_WORKDIR)
    return {
        "ok": True,
        "server": SERVER_NAME,
        "version": SERVER_VERSION,
        "backend": backend,
        "lake_workdir": workdir,
        "has_api_key": bool(os.environ.get("DEEPSEEK_API_KEY", "")),
        "ollama_url": os.environ.get("OLLAMA_URL", "http://localhost:11434"),
    }


TOOLS = {
    "prover_generate": {
        "description": "Generate proof candidates via DeepSeek-Prover-V2-7B.",
        "inputSchema": {
            "type": "object",
            "required": ["theorem_statement"],
            "properties": {
                "theorem_statement": {
                    "type": "string",
                    "description": "The Lean theorem to prove",
                },
                "context": {
                    "type": "string",
                    "description": "Surrounding Lean module code",
                },
                "error_feedback": {
                    "type": "string",
                    "description": "Previous compilation error feedback",
                },
                "num_candidates": {
                    "type": "integer",
                    "default": 4,
                    "description": "Number of candidates to generate",
                },
                "model": {
                    "type": "string",
                    "description": "Model override (e.g. deepseek-prover-v2:7b)",
                },
            },
        },
        "handler": tool_generate,
    },
    "prover_prove": {
        "description": "Iterative proof search: generate → compile → feedback → retry.",
        "inputSchema": {
            "type": "object",
            "required": ["theorem_statement"],
            "properties": {
                "theorem_statement": {
                    "type": "string",
                    "description": "The Lean theorem to prove",
                },
                "context": {
                    "type": "string",
                    "description": "Surrounding Lean module code",
                },
                "max_iterations": {
                    "type": "integer",
                    "default": 5,
                    "description": "Max generate-compile cycles",
                },
                "lake_workdir": {
                    "type": "string",
                    "description": "Working directory for lake build",
                },
                "emit_receipt": {
                    "type": "boolean",
                    "default": True,
                    "description": "Write JSON receipt",
                },
            },
        },
        "handler": tool_prove,
    },
    "prover_prove_batch": {
        "description": "Prove multiple theorems in sequence.",
        "inputSchema": {
            "type": "object",
            "required": ["theorems"],
            "properties": {
                "theorems": {
                    "type": "array",
                    "items": {
                        "type": "object",
                        "properties": {
                            "name": {"type": "string"},
                            "statement": {"type": "string"},
                        },
                        "required": ["statement"],
                    },
                    "description": "List of {name, statement} pairs",
                },
                "context": {
                    "type": "string",
                    "description": "Shared context for all theorems",
                },
                "lake_workdir": {
                    "type": "string",
                    "description": "Working directory for lake build",
                },
            },
        },
        "handler": tool_prove_batch,
    },
    "prover_status": {
        "description": "Return backend health and configuration.",
        "inputSchema": {"type": "object", "properties": {}},
        "handler": tool_status,
    },
}


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()