QML-MCP

"""Quantum Machine Learning MCP Server. This server provides tools for quantum machine learning using Qiskit. """ import json import logging import sys from typing import Any, Dict, Optional import numpy as np from mcp.server import Server from mcp.types import Tool, TextContent from qiskit import QuantumCircuit from qiskit.qasm3 import loads as qasm3_loads from config import config from qml import ( QuantumCircuitRunner, QuantumKernelComputer, VQCTrainer, ModelEvaluator, serialize_numpy, ) # Configure logging logging.basicConfig( level=getattr(logging, config.log_level), format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", stream=sys.stderr, ) logger = logging.getLogger(__name__) # Initialize MCP server server = Server("qml-mcp") # Initialize quantum ML components circuit_runner = QuantumCircuitRunner( max_qubits=config.quantum.max_qubits, max_shots=config.quantum.max_shots ) kernel_computer = QuantumKernelComputer(max_qubits=config.quantum.max_qubits) vqc_trainer = VQCTrainer(max_qubits=config.quantum.max_qubits) # ModelEvaluator only has static methods, no instance needed def format_error(error: Exception) -> Dict[str, Any]: """Format error for JSON response.""" error_dict = { "success": False, "error": str(error), "error_type": type(error).__name__, } if config.enable_detailed_errors: import traceback error_dict["traceback"] = traceback.format_exc() return error_dict @server.list_tools() async def list_tools() -> list[Tool]: """List available quantum ML tools.""" return [ Tool( name="run_quantum_circuit", description=( "Execute a quantum circuit and return measurement results. " "Accepts circuit in QASM3 format or as a quantum circuit specification. " f"Maximum {config.quantum.max_qubits} qubits and " f"{config.quantum.max_shots} shots allowed." ), inputSchema={ "type": "object", "properties": { "qasm": { "type": "string", "description": "Quantum circuit in QASM3 format" }, "shots": { "type": "integer", "description": f"Number of shots (default: {config.quantum.default_shots}, max: {config.quantum.max_shots})", "default": config.quantum.default_shots } }, "required": ["qasm"] } ), Tool( name="compute_quantum_kernel", description=( "Compute quantum kernel matrix for machine learning. " "Uses ZZ feature map with fidelity-based kernel. " "Returns kernel matrix for train data or train-test data." ), inputSchema={ "type": "object", "properties": { "train_data": { "type": "array", "description": "Training data as 2D array [[x1, x2, ...], ...]", "items": {"type": "array", "items": {"type": "number"}} }, "test_data": { "type": "array", "description": "Optional test data as 2D array", "items": {"type": "array", "items": {"type": "number"}} }, "feature_dimension": { "type": "integer", "description": "Number of features (defaults to data dimension)" } }, "required": ["train_data"] } ), Tool( name="train_vqc", description=( "Train a Variational Quantum Classifier (VQC) on provided data. " "Uses ZZ feature map and RealAmplitudes ansatz. " "Returns trained model as base64-encoded joblib serialization." ), inputSchema={ "type": "object", "properties": { "X_train": { "type": "array", "description": "Training features as 2D array [[x1, x2, ...], ...]", "items": {"type": "array", "items": {"type": "number"}} }, "y_train": { "type": "array", "description": "Training labels as 1D array [label1, label2, ...]", "items": {"type": "number"} }, "feature_dimension": { "type": "integer", "description": "Number of features (defaults to X_train dimension)" }, "max_iter": { "type": "integer", "description": "Maximum optimization iterations (default: 100)", "default": 100 } }, "required": ["X_train", "y_train"] } ), Tool( name="evaluate_model", description=( "Evaluate a trained quantum ML model on test data. " "Model must be provided as base64-encoded joblib serialization. " "Returns predictions and accuracy (if labels provided)." ), inputSchema={ "type": "object", "properties": { "model": { "type": "string", "description": "Base64-encoded trained model (from train_vqc)" }, "X_test": { "type": "array", "description": "Test features as 2D array [[x1, x2, ...], ...]", "items": {"type": "array", "items": {"type": "number"}} }, "y_test": { "type": "array", "description": "Optional test labels for accuracy computation", "items": {"type": "number"} } }, "required": ["model", "X_test"] } ), ] @server.call_tool() async def call_tool(name: str, arguments: Any) -> list[TextContent]: """Handle tool execution requests.""" logger.info(f"Tool called: {name}") try: if name == "run_quantum_circuit": result = await handle_run_circuit(arguments) elif name == "compute_quantum_kernel": result = await handle_compute_kernel(arguments) elif name == "train_vqc": result = await handle_train_vqc(arguments) elif name == "evaluate_model": result = await handle_evaluate_model(arguments) else: result = format_error(ValueError(f"Unknown tool: {name}")) # Serialize result to JSON result_json = json.dumps(serialize_numpy(result), indent=2) return [TextContent(type="text", text=result_json)] except Exception as e: logger.error(f"Tool execution failed: {e}", exc_info=True) error_result = format_error(e) return [TextContent(type="text", text=json.dumps(error_result, indent=2))] async def handle_run_circuit(arguments: Dict[str, Any]) -> Dict[str, Any]: """Handle quantum circuit execution.""" qasm = arguments.get("qasm") shots = arguments.get("shots", config.quantum.default_shots) if not qasm: raise ValueError("Missing required parameter: qasm") # Parse QASM3 to circuit try: circuit = qasm3_loads(qasm) except Exception as e: raise ValueError(f"Failed to parse QASM3: {e}") # Run circuit result = await circuit_runner.run_circuit(circuit, shots=shots) logger.info(f"Circuit executed successfully with {shots} shots") return result async def handle_compute_kernel(arguments: Dict[str, Any]) -> Dict[str, Any]: """Handle quantum kernel computation.""" train_data = arguments.get("train_data") test_data = arguments.get("test_data") feature_dimension = arguments.get("feature_dimension") if train_data is None: raise ValueError("Missing required parameter: train_data") # Convert to numpy arrays train_data = np.array(train_data, dtype=float) if test_data is not None: test_data = np.array(test_data, dtype=float) # Compute kernel result = await kernel_computer.compute_kernel( train_data=train_data, test_data=test_data, feature_dimension=feature_dimension ) logger.info("Kernel computed successfully") return result async def handle_train_vqc(arguments: Dict[str, Any]) -> Dict[str, Any]: """Handle VQC training.""" X_train = arguments.get("X_train") y_train = arguments.get("y_train") feature_dimension = arguments.get("feature_dimension") max_iter = arguments.get("max_iter", 100) if X_train is None or y_train is None: raise ValueError("Missing required parameters: X_train and y_train") # Convert to numpy arrays X_train = np.array(X_train, dtype=float) y_train = np.array(y_train, dtype=int) # Train VQC result = await vqc_trainer.train( X_train=X_train, y_train=y_train, feature_dimension=feature_dimension, max_iter=max_iter ) logger.info(f"VQC trained successfully with score: {result['train_score']:.4f}") return result async def handle_evaluate_model(arguments: Dict[str, Any]) -> Dict[str, Any]: """Handle model evaluation.""" model = arguments.get("model") X_test = arguments.get("X_test") y_test = arguments.get("y_test") if model is None or X_test is None: raise ValueError("Missing required parameters: model and X_test") # Convert to numpy arrays X_test = np.array(X_test, dtype=float) if y_test is not None: y_test = np.array(y_test, dtype=int) # Evaluate model result = await ModelEvaluator.evaluate( model_base64=model, X_test=X_test, y_test=y_test ) logger.info("Model evaluated successfully") return result async def main(): """Run the MCP server.""" from mcp.server.stdio import stdio_server logger.info("Starting QML MCP Server") logger.info(f"Max qubits: {config.quantum.max_qubits}") logger.info(f"Max shots: {config.quantum.max_shots}") async with stdio_server() as (read_stream, write_stream): await server.run( read_stream, write_stream, server.create_initialization_options() ) if __name__ == "__main__": import asyncio asyncio.run(main())