Update auto-sklearn/mcp_output/mcp_plugin/mcp_service.py

auto-sklearn/mcp_output/mcp_plugin/mcp_service.py

@@ -1,5 +1,7 @@
 import os
 import sys
+from typing import Dict, Any, List, Optional
+import json
 
 # Add the local source directory to sys.path
 source_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "source")
@@ -7,114 +9,448 @@ if source_path not in sys.path:
     sys.path.insert(0, source_path)
 
 from fastmcp import FastMCP
+import numpy as np
 
 # Import core modules from the local source directory
-from autosklearn.automl import AutoML
-from autosklearn.ensembles.ensemble_selection import EnsembleSelection
 from autosklearn.estimators import AutoSklearnClassifier, AutoSklearnRegressor
-from autosklearn.smbo import AutoMLSMBO
 
 # Create the FastMCP service application
 mcp = FastMCP("auto_sklearn_service")
 
-@mcp.tool(name="automl_fit", description="Fit the AutoML model with the given data.")
-def automl_fit(X_train: list, y_train: list, time_left_for_this_task: int) -> dict:
+# Store models by ID
+_classifiers: Dict[str, AutoSklearnClassifier] = {}
+_regressors: Dict[str, AutoSklearnRegressor] = {}
+
+
+@mcp.tool(name="get_library_info")
+def get_library_info() -> dict:
+    """
+    Get information about the auto-sklearn library.
+
+    Returns:
+        dict: Version and configuration information.
+    """
+    try:
+        from autosklearn import __version__
+        
+        return {
+            "success": True,
+            "result": {
+                "library": "auto-sklearn",
+                "version": __version__,
+                "estimators": ["AutoSklearnClassifier", "AutoSklearnRegressor"],
+                "features": [
+                    "Automated Machine Learning",
+                    "Ensemble Learning",
+                    "Meta-learning",
+                    "Hyperparameter Optimization (SMAC)",
+                ],
+            },
+            "error": None,
+        }
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
+@mcp.tool(name="create_classifier")
+def create_classifier(
+    classifier_id: str,
+    time_left_for_this_task: int = 3600,
+    per_run_time_limit: Optional[int] = None,
+    ensemble_size: int = 50,
+    ensemble_nbest: int = 50,
+    seed: int = 1,
+    memory_limit: int = 3072,
+    n_jobs: Optional[int] = None,
+) -> dict:
     """
-    Fit the AutoML model with the given training data.
+    Create a new AutoSklearnClassifier.
 
     Parameters:
-    - X_train: list of training data features
-    - y_train: list of training data labels
-    - time_left_for_this_task: time budget in seconds for the task
+        classifier_id (str): Unique identifier for the classifier.
+        time_left_for_this_task (int): Total time budget in seconds.
+        per_run_time_limit (Optional[int]): Time limit per model evaluation.
+        ensemble_size (int): Number of models in the final ensemble.
+        ensemble_nbest (int): Consider only the best n models for ensemble.
+        seed (int): Random seed.
+        memory_limit (int): Memory limit in MB.
+        n_jobs (Optional[int]): Number of parallel jobs.
 
     Returns:
-    - dict: containing success, result, or error fields
+        dict: Success status and classifier information.
     """
     try:
-        automl = AutoML(time_left_for_this_task=time_left_for_this_task)
-        automl.fit(X_train, y_train)
-        return {"success": True, "result": "Model fitted successfully"}
+        if classifier_id in _classifiers:
+            return {"success": False, "result": None, "error": f"Classifier '{classifier_id}' already exists"}
+        
+        clf = AutoSklearnClassifier(
+            time_left_for_this_task=time_left_for_this_task,
+            per_run_time_limit=per_run_time_limit,
+            ensemble_size=ensemble_size,
+            ensemble_nbest=ensemble_nbest,
+            seed=seed,
+            memory_limit=memory_limit,
+            n_jobs=n_jobs,
+        )
+        
+        _classifiers[classifier_id] = clf
+        
+        return {
+            "success": True,
+            "result": {
+                "classifier_id": classifier_id,
+                "time_budget": time_left_for_this_task,
+                "ensemble_size": ensemble_size,
+                "message": "Classifier created successfully",
+            },
+            "error": None,
+        }
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
 
-@mcp.tool(name="ensemble_selection", description="Select the best ensemble from the given models.")
-def ensemble_selection(models: list) -> dict:
+
+@mcp.tool(name="create_regressor")
+def create_regressor(
+    regressor_id: str,
+    time_left_for_this_task: int = 3600,
+    per_run_time_limit: Optional[int] = None,
+    ensemble_size: int = 50,
+    ensemble_nbest: int = 50,
+    seed: int = 1,
+    memory_limit: int = 3072,
+    n_jobs: Optional[int] = None,
+) -> dict:
     """
-    Select the best ensemble from the given models.
+    Create a new AutoSklearnRegressor.
 
     Parameters:
-    - models: list of models to select from
+        regressor_id (str): Unique identifier for the regressor.
+        time_left_for_this_task (int): Total time budget in seconds.
+        per_run_time_limit (Optional[int]): Time limit per model evaluation.
+        ensemble_size (int): Number of models in the final ensemble.
+        ensemble_nbest (int): Consider only the best n models for ensemble.
+        seed (int): Random seed.
+        memory_limit (int): Memory limit in MB.
+        n_jobs (Optional[int]): Number of parallel jobs.
 
     Returns:
-    - dict: containing success, result, or error fields
+        dict: Success status and regressor information.
     """
     try:
-        ensemble = EnsembleSelection(models)
-        best_ensemble = ensemble.select()
-        return {"success": True, "result": "Ensemble selected successfully"}
+        if regressor_id in _regressors:
+            return {"success": False, "result": None, "error": f"Regressor '{regressor_id}' already exists"}
+        
+        reg = AutoSklearnRegressor(
+            time_left_for_this_task=time_left_for_this_task,
+            per_run_time_limit=per_run_time_limit,
+            ensemble_size=ensemble_size,
+            ensemble_nbest=ensemble_nbest,
+            seed=seed,
+            memory_limit=memory_limit,
+            n_jobs=n_jobs,
+        )
+        
+        _regressors[regressor_id] = reg
+        
+        return {
+            "success": True,
+            "result": {
+                "regressor_id": regressor_id,
+                "time_budget": time_left_for_this_task,
+                "ensemble_size": ensemble_size,
+                "message": "Regressor created successfully",
+            },
+            "error": None,
+        }
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
+
 
-@mcp.tool(name="classifier_predict", description="Predict using the AutoSklearnClassifier.")
-def classifier_predict(X_test: list) -> dict:
+@mcp.tool(name="fit_classifier")
+def fit_classifier(classifier_id: str, X_train: List[List[float]], y_train: List) -> dict:
     """
-    Predict using the AutoSklearnClassifier.
+    Fit a classifier with training data.
 
     Parameters:
-    - X_test: list of test data features
+        classifier_id (str): ID of the classifier to fit.
+        X_train (List[List[float]]): Training features.
+        y_train (List): Training labels.
 
     Returns:
-    - dict: containing success, result, or error fields
+        dict: Success status and fitting information.
     """
     try:
-        classifier = AutoSklearnClassifier()
-        predictions = classifier.predict(X_test)
-        return {"success": True, "result": predictions}
+        if classifier_id not in _classifiers:
+            return {"success": False, "result": None, "error": f"Classifier '{classifier_id}' not found"}
+        
+        clf = _classifiers[classifier_id]
+        X_train_np = np.array(X_train)
+        y_train_np = np.array(y_train)
+        
+        clf.fit(X_train_np, y_train_np)
+        
+        return {
+            "success": True,
+            "result": {
+                "classifier_id": classifier_id,
+                "num_samples": len(X_train),
+                "num_features": len(X_train[0]) if X_train else 0,
+                "message": "Classifier fitted successfully",
+            },
+            "error": None,
+        }
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
+
 
-@mcp.tool(name="regressor_predict", description="Predict using the AutoSklearnRegressor.")
-def regressor_predict(X_test: list) -> dict:
+@mcp.tool(name="fit_regressor")
+def fit_regressor(regressor_id: str, X_train: List[List[float]], y_train: List[float]) -> dict:
     """
-    Predict using the AutoSklearnRegressor.
+    Fit a regressor with training data.
 
     Parameters:
-    - X_test: list of test data features
+        regressor_id (str): ID of the regressor to fit.
+        X_train (List[List[float]]): Training features.
+        y_train (List[float]): Training targets.
 
     Returns:
-    - dict: containing success, result, or error fields
+        dict: Success status and fitting information.
     """
     try:
-        regressor = AutoSklearnRegressor()
-        predictions = regressor.predict(X_test)
-        return {"success": True, "result": predictions}
+        if regressor_id not in _regressors:
+            return {"success": False, "result": None, "error": f"Regressor '{regressor_id}' not found"}
+        
+        reg = _regressors[regressor_id]
+        X_train_np = np.array(X_train)
+        y_train_np = np.array(y_train)
+        
+        reg.fit(X_train_np, y_train_np)
+        
+        return {
+            "success": True,
+            "result": {
+                "regressor_id": regressor_id,
+                "num_samples": len(X_train),
+                "num_features": len(X_train[0]) if X_train else 0,
+                "message": "Regressor fitted successfully",
+            },
+            "error": None,
+        }
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
 
-@mcp.tool(name="smbo_optimize", description="Optimize hyperparameters using SMBO.")
-def smbo_optimize(config_space: dict, max_evals: int) -> dict:
+
+@mcp.tool(name="predict_classifier")
+def predict_classifier(classifier_id: str, X_test: List[List[float]]) -> dict:
     """
-    Optimize hyperparameters using SMBO.
+    Make predictions using a fitted classifier.
 
     Parameters:
-    - config_space: dictionary defining the configuration space
-    - max_evals: maximum number of evaluations
+        classifier_id (str): ID of the classifier.
+        X_test (List[List[float]]): Test features.
 
     Returns:
-    - dict: containing success, result, or error fields
+        dict: Predictions.
     """
     try:
-        smbo = AutoMLSMBO(config_space=config_space, max_evals=max_evals)
-        best_config = smbo.optimize()
-        return {"success": True, "result": best_config}
+        if classifier_id not in _classifiers:
+            return {"success": False, "result": None, "error": f"Classifier '{classifier_id}' not found"}
+        
+        clf = _classifiers[classifier_id]
+        X_test_np = np.array(X_test)
+        
+        predictions = clf.predict(X_test_np)
+        
+        return {
+            "success": True,
+            "result": {
+                "predictions": predictions.tolist(),
+                "num_predictions": len(predictions),
+            },
+            "error": None,
+        }
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
+
+
+@mcp.tool(name="predict_regressor")
+def predict_regressor(regressor_id: str, X_test: List[List[float]]) -> dict:
+    """
+    Make predictions using a fitted regressor.
+
+    Parameters:
+        regressor_id (str): ID of the regressor.
+        X_test (List[List[float]]): Test features.
+
+    Returns:
+        dict: Predictions.
+    """
+    try:
+        if regressor_id not in _regressors:
+            return {"success": False, "result": None, "error": f"Regressor '{regressor_id}' not found"}
+        
+        reg = _regressors[regressor_id]
+        X_test_np = np.array(X_test)
+        
+        predictions = reg.predict(X_test_np)
+        
+        return {
+            "success": True,
+            "result": {
+                "predictions": predictions.tolist(),
+                "num_predictions": len(predictions),
+            },
+            "error": None,
+        }
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
+@mcp.tool(name="get_classifier_model_performance")
+def get_classifier_model_performance(classifier_id: str) -> dict:
+    """
+    Get performance statistics of explored models.
+
+    Parameters:
+        classifier_id (str): ID of the classifier.
+
+    Returns:
+        dict: Performance statistics.
+    """
+    try:
+        if classifier_id not in _classifiers:
+            return {"success": False, "result": None, "error": f"Classifier '{classifier_id}' not found"}
+        
+        clf = _classifiers[classifier_id]
+        
+        # Get leaderboard information
+        leaderboard = clf.leaderboard()
+        
+        return {
+            "success": True,
+            "result": {
+                "classifier_id": classifier_id,
+                "num_models": len(leaderboard) if leaderboard is not None else 0,
+                "leaderboard_summary": leaderboard.head(10).to_dict() if leaderboard is not None else {},
+            },
+            "error": None,
+        }
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
+@mcp.tool(name="get_regressor_model_performance")
+def get_regressor_model_performance(regressor_id: str) -> dict:
+    """
+    Get performance statistics of explored models.
+
+    Parameters:
+        regressor_id (str): ID of the regressor.
+
+    Returns:
+        dict: Performance statistics.
+    """
+    try:
+        if regressor_id not in _regressors:
+            return {"success": False, "result": None, "error": f"Regressor '{regressor_id}' not found"}
+        
+        reg = _regressors[regressor_id]
+        
+        # Get leaderboard information
+        leaderboard = reg.leaderboard()
+        
+        return {
+            "success": True,
+            "result": {
+                "regressor_id": regressor_id,
+                "num_models": len(leaderboard) if leaderboard is not None else 0,
+                "leaderboard_summary": leaderboard.head(10).to_dict() if leaderboard is not None else {},
+            },
+            "error": None,
+        }
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
+@mcp.tool(name="list_models")
+def list_models() -> dict:
+    """
+    List all stored classifiers and regressors.
+
+    Returns:
+        dict: List of model IDs.
+    """
+    try:
+        return {
+            "success": True,
+            "result": {
+                "classifiers": list(_classifiers.keys()),
+                "regressors": list(_regressors.keys()),
+                "total": len(_classifiers) + len(_regressors),
+            },
+            "error": None,
+        }
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
+@mcp.tool(name="delete_classifier")
+def delete_classifier(classifier_id: str) -> dict:
+    """
+    Delete a stored classifier.
+
+    Parameters:
+        classifier_id (str): ID of the classifier to delete.
+
+    Returns:
+        dict: Confirmation of deletion.
+    """
+    try:
+        if classifier_id not in _classifiers:
+            return {"success": False, "result": None, "error": f"Classifier '{classifier_id}' not found"}
+        
+        del _classifiers[classifier_id]
+        
+        return {
+            "success": True,
+            "result": {"message": f"Classifier '{classifier_id}' deleted"},
+            "error": None,
+        }
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
+@mcp.tool(name="delete_regressor")
+def delete_regressor(regressor_id: str) -> dict:
+    """
+    Delete a stored regressor.
+
+    Parameters:
+        regressor_id (str): ID of the regressor to delete.
+
+    Returns:
+        dict: Confirmation of deletion.
+    """
+    try:
+        if regressor_id not in _regressors:
+            return {"success": False, "result": None, "error": f"Regressor '{regressor_id}' not found"}
+        
+        del _regressors[regressor_id]
+        
+        return {
+            "success": True,
+            "result": {"message": f"Regressor '{regressor_id}' deleted"},
+            "error": None,
+        }
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
 
 def create_app() -> FastMCP:
     """
     Create and return the FastMCP application instance.
 
     Returns:
-    - FastMCP: the FastMCP application instance
+        FastMCP: the FastMCP application instance
     """
-    return mcp
+    return mcp