server/mcp.py

#!/usr/bin/env python3
"""
Text-to-CNC MCP Server
======================
Exposes the text-to-CNC pipeline as MCP tools over stdio transport.

Tools:
  - generate_cnc_model: Text prompt → CadQuery code → 3D solid → STEP/STL
  - validate_cnc_model: Run CNC manufacturability checks on CadQuery code
  - execute_cadquery:   Run arbitrary CadQuery code and get geometry info
  - list_models:        List previously generated models in the output dir

Usage:
  python -m server.mcp                    # stdio transport (default)
  python -m server.mcp --transport sse    # SSE transport on port 8000
"""

import json
import os
from pathlib import Path

from mcp.server.fastmcp import FastMCP

from core.cadquery_prompts import build_messages, CADQUERY_SYSTEM_PROMPT
from core.executor import execute_cadquery, export_all
from core.validator import validate_for_cnc


def _exec_dump(result) -> dict:
    return result.model_dump(by_alias=True)


def _val_dump(result) -> dict:
    return result.model_dump()


# ── Server Setup ──────────────────────────────────────────────────────────

mcp = FastMCP(
    "text-to-cnc",
    instructions=(
        "Generate CNC-machinable 3D models from text descriptions. "
        "Converts natural language → CadQuery code → validated STEP/STL files. "
        "Version 1.0.0"
    ),
)

from config.settings import settings
from core.backend_factory import BackendFactory

DEFAULT_OUTPUT_DIR = settings.output_dir
if not DEFAULT_OUTPUT_DIR.is_absolute():
    DEFAULT_OUTPUT_DIR = Path(__file__).parent.parent / DEFAULT_OUTPUT_DIR
DEFAULT_OUTPUT_DIR.mkdir(exist_ok=True)


# ── Helper: LLM Backend Selection ────────────────────────────────────────

def get_backend(backend_name: str = "mock"):
    """Get LLM backend by name, using the factory registry.

    NeuralCADBackend is a special case not registered in the factory.
    """
    if backend_name == "neural":
        from core.backends import NeuralCADBackend
        return NeuralCADBackend()
    return BackendFactory.create_safe(backend_name)


# ── Tool: generate_cnc_model ─────────────────────────────────────────────

@mcp.tool()
def generate_cnc_model(
    prompt: str,
    part_name: str = "",
    backend: str = "mock",
    max_retries: int = 2,
    output_format: str = "both",
) -> str:
    """
    Generate a CNC-machinable 3D model from a text description.

    Takes a natural language description of a mechanical part, generates
    CadQuery Python code via an LLM, executes it to produce a 3D solid,
    validates it for CNC manufacturability, and exports STEP/STL files.

    Args:
        prompt: Natural language description of the part to generate.
                Example: "A mounting bracket with four M6 bolt holes, 80mm wide"
        part_name: Optional name for the part (used in filenames).
                   If empty, auto-generated from the prompt.
        backend: LLM backend to use: "mock" (no API key), "anthropic", or "openai".
        max_retries: Number of retry attempts if code generation fails (0-3).
        output_format: Export format: "step", "stl", or "both".

    Returns:
        JSON string with generation results including:
        - generated_code: The CadQuery Python code
        - execution: Success/failure status and geometry metadata
        - validation: CNC manufacturability analysis
        - exported_files: Paths to generated STEP/STL files
    """
    from core.pipeline import run_pipeline

    if not part_name:
        part_name = prompt[:40].strip().replace(" ", "_").lower()
        part_name = "".join(c for c in part_name if c.isalnum() or c == "_")

    llm_backend = get_backend(backend)

    result = run_pipeline(
        prompt=prompt,
        backend=llm_backend,
        output_dir=DEFAULT_OUTPUT_DIR,
        max_retries=min(max_retries, 3),
        export=True,
        validate=True,
        part_name=part_name,
    )

    # Build response
    response = {
        "success": result.execution.success,
        "prompt": prompt,
        "part_name": part_name,
        "retries": result.retry_count,
        "generated_code": result.generated_code,
        "execution": _exec_dump(result.execution),
    }

    if result.validation:
        response["validation"] = _val_dump(result.validation)

    if result.exported_files:
        response["exported_files"] = {
            fmt: str(path) for fmt, path in result.exported_files.items()
        }

    return json.dumps(response, indent=2)


# ── Tool: validate_cnc_model ─────────────────────────────────────────────

@mcp.tool()
def validate_cnc_model(
    cadquery_code: str,
    part_name: str = "Part",
    min_wall_thickness_mm: float = 1.5,
    max_part_size_mm: float = 500.0,
) -> str:
    """
    Validate CadQuery code for CNC manufacturability without generating new code.

    Executes the provided CadQuery code, then runs manufacturability checks
    including wall thickness, tool access, aspect ratios, and surface complexity.

    Args:
        cadquery_code: Valid CadQuery Python code that assigns result to `result`.
                       Example: 'import cadquery as cq\\nresult = cq.Workplane("XY").box(10,10,10)'
        part_name: Name for the part in the validation report.
        min_wall_thickness_mm: Minimum acceptable wall thickness in mm (default 1.5).
        max_part_size_mm: Maximum part dimension in mm (default 500).

    Returns:
        JSON string with execution status and CNC validation results including
        machinable flag, axis recommendation, and list of issues.
    """
    exec_result = execute_cadquery(cadquery_code)

    response = {
        "execution_success": exec_result.success,
        "error": exec_result.error,
        "volume_mm3": exec_result.volume,
        "bounding_box_mm": list(exec_result.bounding_box) if exec_result.bounding_box else [],
    }

    if exec_result.success:
        config = {
            "min_wall_thickness_mm": min_wall_thickness_mm,
            "max_part_size_mm": max_part_size_mm,
        }
        validation = validate_for_cnc(exec_result.result, part_name=part_name, config=config)
        response["validation"] = {
            **_val_dump(validation),
            "summary": validation.summary(),
        }

    return json.dumps(response, indent=2)


# ── Tool: execute_cadquery ───────────────────────────────────────────────

@mcp.tool()
def execute_cadquery_code(
    code: str,
    export_path: str = "",
) -> str:
    """
    Execute CadQuery Python code and return geometry information.

    Runs CadQuery code in a sandboxed environment and returns metadata
    about the resulting 3D solid (volume, bounding box, face/edge counts).
    Optionally exports to STEP/STL.

    Args:
        code: CadQuery Python code. Must assign the final solid to a variable
              called `result`. Example:
              'import cadquery as cq\\nresult = cq.Workplane("XY").box(20,20,20).hole(8)'
        export_path: Optional base file path for STEP/STL export (without extension).
                     Example: "output/my_part" → creates my_part.step and my_part.stl

    Returns:
        JSON string with execution results including success status,
        geometry metadata, stdout output, and export file paths if requested.
    """
    exec_result = execute_cadquery(code)

    response = {
        **_exec_dump(exec_result),
        "stdout": exec_result.stdout,
    }

    if exec_result.success and export_path:
        try:
            files = export_all(exec_result.result, export_path)
            response["exported_files"] = {fmt: str(p) for fmt, p in files.items()}
        except Exception as e:
            response["export_error"] = str(e)

    return json.dumps(response, indent=2)


# ── Tool: list_models ────────────────────────────────────────────────────

@mcp.tool()
def list_models(output_dir: str = "") -> str:
    """
    List all previously generated CNC models in the output directory.

    Returns a list of generated STEP and STL files with their sizes.

    Args:
        output_dir: Directory to scan. Defaults to the server's output directory.

    Returns:
        JSON string with a list of model files and their sizes in bytes.
    """
    scan_dir = Path(output_dir) if output_dir else DEFAULT_OUTPUT_DIR

    if not scan_dir.exists():
        return json.dumps({"error": f"Directory not found: {scan_dir}"})

    models = {}
    for ext in ("*.step", "*.stl"):
        for f in scan_dir.glob(ext):
            name = f.stem
            if name not in models:
                models[name] = {"name": name, "files": {}}
            models[name]["files"][f.suffix.lstrip(".")] = {
                "path": str(f),
                "size_bytes": f.stat().st_size,
            }

    return json.dumps({
        "output_dir": str(scan_dir),
        "model_count": len(models),
        "models": list(models.values()),
    }, indent=2)


# ── Tool: generate_from_image ───────────────────────────────────────────

@mcp.tool()
def generate_from_image(
    image_path: str,
    text_hint: str = "",
    part_name: str = "",
    backend: str = "anthropic",
    max_retries: int = 2,
) -> str:
    """
    Generate a CNC-machinable 3D model from a photo or sketch image.

    Sends the image to a vision-capable LLM (Claude or GPT-4o) along with
    the CadQuery system prompt to generate code, then executes, validates,
    and exports the result.

    Args:
        image_path: Path to an image file (photo, sketch, or CAD screenshot).
        text_hint: Optional text to guide generation alongside the image.
                   Example: "This is a mounting bracket — add M6 bolt holes"
        part_name: Optional name for the part (used in filenames).
        backend: LLM backend: "anthropic" or "openai". Must support vision.
        max_retries: Number of retry attempts if code execution fails (0-3).

    Returns:
        JSON string with generation results including generated code,
        execution status, validation, and exported file paths.
    """
    if not Path(image_path).exists():
        return json.dumps({"success": False, "error": f"Image not found: {image_path}"})

    if not part_name:
        part_name = Path(image_path).stem

    llm_backend = get_backend(backend)

    # Build prompt with optional text hint
    prompt = "Generate CadQuery code for the mechanical part shown in this image."
    if text_hint:
        prompt += f"\n\nAdditional context: {text_hint}"

    messages = build_messages(prompt)

    # Use vision-capable generate_with_image
    generated_code = llm_backend.generate_with_image(messages, image_path)

    # Run through standard execution/validation/export
    exec_result = execute_cadquery(generated_code)
    retry_count = 0

    while not exec_result.success and retry_count < min(max_retries, 3):
        retry_count += 1
        error_feedback = (
            f"The previous code failed with this error:\n"
            f"```\n{exec_result.error}\n```\n\n"
            f"Please fix the code and return only the corrected Python code."
        )
        retry_messages = build_messages(error_feedback)
        generated_code = llm_backend.generate_with_image(retry_messages, image_path)
        exec_result = execute_cadquery(generated_code)

    response = {
        "success": exec_result.success,
        "image_path": image_path,
        "text_hint": text_hint,
        "part_name": part_name,
        "backend": backend,
        "retries": retry_count,
        "generated_code": generated_code,
        "execution": _exec_dump(exec_result),
    }

    if exec_result.success:
        validation = validate_for_cnc(exec_result.result, part_name=part_name)
        response["validation"] = _val_dump(validation)

        base_path = DEFAULT_OUTPUT_DIR / part_name
        try:
            exported = export_all(exec_result.result, base_path)
            response["exported_files"] = {fmt: str(p) for fmt, p in exported.items()}
        except Exception as e:
            response["export_error"] = str(e)

    return json.dumps(response, indent=2)


# ── Tool: chat_turn ─────────────────────────────────────────────────────

@mcp.tool()
def chat_turn(
    message: str,
    history: str = "[]",
    mentions: str = "[]",
    backend: str = "mock",
) -> str:
    """
    Multi-agent chat turn for collaborative CAD design.

    Send a message to the design team agents (Design, Engineering, CNC, CAD Coder).
    Agents collaborate to help you design a mechanical part step by step.

    Args:
        message: Your message to the design team.
                 Use @design, @engineering, @cnc, or @cad to address specific agents.
        history: JSON string of previous messages. Format:
                 [{"role": "user"|"agent", "agent_id": "design", "content": "..."}]
        mentions: JSON string of agent IDs to address. Format: ["design", "engineering"]
                  Empty list = auto-route based on message content.
        backend: LLM backend: "mock", "gemini", "anthropic", "openai".

    Returns:
        JSON string with agent responses and optional 3D preview data.
    """
    import json as json_mod

    from agents.orchestrator import get_orchestrator
    from agents.crew_orchestrator import CrewOrchestrator
    from agents.prompts import parse_mentions

    history_list = json_mod.loads(history) if isinstance(history, str) else history
    mentions_list = json_mod.loads(mentions) if isinstance(mentions, str) else mentions

    # Parse @mentions from message if not provided
    if not mentions_list:
        message, mentions_list = parse_mentions(message)

    mentions_or_none = mentions_list if mentions_list else None

    if backend in ("anthropic", "openai"):
        orchestrator = CrewOrchestrator(backend_name=backend, output_dir=DEFAULT_OUTPUT_DIR)
    else:
        orchestrator = get_orchestrator(backend, output_dir=DEFAULT_OUTPUT_DIR)

    result = orchestrator.chat_turn(
        message=message,
        history=history_list,
        mentions=mentions_or_none,
    )

    return json_mod.dumps(result, indent=2)


# ── Resource: System prompt (for transparency) ───────────────────────────

@mcp.resource("text-to-cnc://system-prompt")
def get_system_prompt() -> str:
    """The CadQuery generation system prompt used by the LLM."""
    return CADQUERY_SYSTEM_PROMPT


@mcp.resource("text-to-cnc://capabilities")
def get_capabilities() -> str:
    """Server capabilities and configuration."""
    backends = ["mock (always available)", "neural (local models — requires trained weights)"]
    if os.environ.get("ANTHROPIC_API_KEY"):
        backends.append("anthropic (API key detected)")
    if os.environ.get("OPENAI_API_KEY"):
        backends.append("openai (API key detected)")
    if os.environ.get("GEMINI_API_KEY"):
        backends.append("gemini (API key detected)")

    return json.dumps({
        "name": "text-to-cnc",
        "version": "1.0.0",
        "available_backends": backends,
        "output_dir": str(DEFAULT_OUTPUT_DIR),
        "export_formats": ["STEP", "STL"],
        "cnc_validation": True,
        "max_retries": 3,
    }, indent=2)


# ── Entry Point ──────────────────────────────────────────────────────────

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Text-to-CNC MCP Server")
    parser.add_argument(
        "--transport", choices=["stdio", "sse"], default="stdio",
        help="MCP transport (default: stdio)"
    )
    parser.add_argument(
        "--port", type=int, default=8000,
        help="Port for SSE transport (default: 8000)"
    )
    args = parser.parse_args()

    if args.transport == "sse":
        mcp.run(transport="sse")
    else:
        mcp.run(transport="stdio")