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

auto-sklearn/mcp_output/mcp_plugin/mcp_service.py

@@ -1,260 +1,451 @@
+import os
+import sys
+from typing import Dict, Any, List, Optional
+import json
+
 from fastmcp import FastMCP
+import numpy as np
+
+# Import core modules from installed auto-sklearn package
+from autosklearn.estimators import AutoSklearnClassifier, AutoSklearnRegressor
 
 # Create the FastMCP service application
 mcp = FastMCP("auto_sklearn_service")
 
-# Define tools here following the AgML MCP structure
+# Store models by ID
+_classifiers: Dict[str, AutoSklearnClassifier] = {}
+_regressors: Dict[str, AutoSklearnRegressor] = {}
+
 
-# Example tool
-@mcp.tool(name="example_tool", description="Example tool description")
-def example_tool() -> dict:
+@mcp.tool(name="get_library_info")
+def get_library_info() -> dict:
     """
-    Example tool function.
+    Get information about the auto-sklearn library.
 
     Returns:
-    - dict: Example response.
+        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:
     """
-    return {"success": True, "message": "This is an example tool."}
+    Create a new AutoSklearnClassifier.
+
+    Parameters:
+        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.
 
-@mcp.tool(name="load_dataset", description="Load a dataset using auto-sklearn")
-def load_dataset(dataset_name: str) -> dict:
+    Returns:
+        dict: Success status and classifier information.
+    """
+    try:
+        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, "result": None, "error": str(e)}
+
+
+@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:
     """
-    Load a dataset using auto-sklearn.
+    Create a new AutoSklearnRegressor.
 
     Parameters:
-    - dataset_name: Name of the dataset to load (e.g., 'breast_cancer')
+        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: Information about the loaded dataset.
+        dict: Success status and regressor information.
     """
     try:
-        import sklearn.datasets
-        X, y = sklearn.datasets.fetch_openml(data_id=dataset_name, return_X_y=True, as_frame=True)
+        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,
-            "dataset_name": dataset_name,
-            "num_samples": len(X),
-            "num_features": X.shape[1],
-            "num_classes": len(set(y))
+            "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="create_classifier", description="Create an AutoSklearnClassifier")
-def create_classifier(time_limit: int, per_run_time_limit: int, memory_limit: int = 3072) -> dict:
+@mcp.tool(name="fit_classifier")
+def fit_classifier(classifier_id: str, X_train: List[List[float]], y_train: List) -> dict:
     """
-    Create an AutoSklearnClassifier instance.
+    Fit a classifier with training data.
 
     Parameters:
-    - time_limit: Total time limit for the AutoML process.
-    - per_run_time_limit: Time limit for each model training.
-    - memory_limit: Memory limit for each model training (default: 3072MB).
+        classifier_id (str): ID of the classifier to fit.
+        X_train (List[List[float]]): Training features.
+        y_train (List): Training labels.
 
     Returns:
-    - dict: Information about the created classifier.
+        dict: Success status and fitting information.
     """
     try:
-        import autosklearn.classification
-        classifier = autosklearn.classification.AutoSklearnClassifier(
-            time_left_for_this_task=time_limit,
-            per_run_time_limit=per_run_time_limit,
-            memory_limit=memory_limit
-        )
+        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,
-            "message": "Classifier created successfully",
-            "time_limit": time_limit,
-            "per_run_time_limit": per_run_time_limit,
-            "memory_limit": memory_limit
+            "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="fit_classifier", description="Fit the AutoSklearnClassifier")
-def fit_classifier(classifier, X_train, y_train, dataset_name: str) -> dict:
+@mcp.tool(name="fit_regressor")
+def fit_regressor(regressor_id: str, X_train: List[List[float]], y_train: List[float]) -> dict:
     """
-    Fit the AutoSklearnClassifier on the training data.
+    Fit a regressor with training data.
 
     Parameters:
-    - classifier: The AutoSklearnClassifier instance.
-    - X_train: Training features.
-    - y_train: Training labels.
-    - dataset_name: Name of the dataset.
+        regressor_id (str): ID of the regressor to fit.
+        X_train (List[List[float]]): Training features.
+        y_train (List[float]): Training targets.
 
     Returns:
-    - dict: Information about the fitting process.
+        dict: Success status and fitting information.
     """
     try:
-        classifier.fit(X_train, y_train, dataset_name=dataset_name)
+        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,
-            "message": "Classifier fitted successfully",
-            "dataset_name": dataset_name
+            "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="predict", description="Make predictions using the trained classifier")
-def predict(classifier, X_test) -> dict:
+@mcp.tool(name="predict_classifier")
+def predict_classifier(classifier_id: str, X_test: List[List[float]]) -> dict:
     """
-    Make predictions using the trained AutoSklearnClassifier.
+    Make predictions using a fitted classifier.
 
     Parameters:
-    - classifier: The trained AutoSklearnClassifier instance.
-    - X_test: Test features.
+        classifier_id (str): ID of the classifier.
+        X_test (List[List[float]]): Test features.
 
     Returns:
-    - dict: Predictions and success status.
+        dict: Predictions.
     """
     try:
-        predictions = classifier.predict(X_test)
+        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,
-            "predictions": predictions.tolist()
+            "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="optimize_hyperparameters", description="Optimize hyperparameters using AutoSklearn")
-def optimize_hyperparameters(X_train, y_train, time_limit: int, per_run_time_limit: int, memory_limit: int = 3072) -> dict:
+@mcp.tool(name="predict_regressor")
+def predict_regressor(regressor_id: str, X_test: List[List[float]]) -> dict:
     """
-    Optimize hyperparameters using AutoSklearn.
+    Make predictions using a fitted regressor.
 
     Parameters:
-    - X_train: Training features.
-    - y_train: Training labels.
-    - time_limit: Total time limit for the AutoML process.
-    - per_run_time_limit: Time limit for each model training.
-    - memory_limit: Memory limit for each model training (default: 3072MB).
+        regressor_id (str): ID of the regressor.
+        X_test (List[List[float]]): Test features.
 
     Returns:
-    - dict: Optimization results and best model information.
+        dict: Predictions.
     """
     try:
-        import autosklearn.classification
-        automl = autosklearn.classification.AutoSklearnClassifier(
-            time_left_for_this_task=time_limit,
-            per_run_time_limit=per_run_time_limit,
-            memory_limit=memory_limit
-        )
-        automl.fit(X_train, y_train)
+        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,
-            "message": "Hyperparameter optimization completed successfully",
-            "best_model": automl.show_models(),
-            "statistics": automl.sprint_statistics()
+            "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="evaluate_model", description="Evaluate a trained model on test data")
-def evaluate_model(classifier, X_test, y_test) -> dict:
+@mcp.tool(name="get_classifier_model_performance")
+def get_classifier_model_performance(classifier_id: str) -> dict:
     """
-    Evaluate a trained model on test data.
+    Get performance statistics of explored models.
 
     Parameters:
-    - classifier: The trained AutoSklearnClassifier instance.
-    - X_test: Test features.
-    - y_test: Test labels.
+        classifier_id (str): ID of the classifier.
 
     Returns:
-    - dict: Evaluation metrics.
+        dict: Performance statistics.
     """
     try:
-        import sklearn.metrics
-        predictions = classifier.predict(X_test)
-        accuracy = sklearn.metrics.accuracy_score(y_test, predictions)
+        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,
-            "accuracy": accuracy,
-            "message": "Model evaluation completed successfully"
+            "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, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
 
 
-@mcp.tool(name="get_pipeline_components", description="Get pipeline components used by AutoSklearn")
-def get_pipeline_components() -> dict:
+@mcp.tool(name="get_regressor_model_performance")
+def get_regressor_model_performance(regressor_id: str) -> dict:
     """
-    Get pipeline components used by AutoSklearn.
+    Get performance statistics of explored models.
+
+    Parameters:
+        regressor_id (str): ID of the regressor.
 
     Returns:
-    - dict: Information about pipeline components.
+        dict: Performance statistics.
     """
     try:
-        import autosklearn.pipeline.components.classification as classification_components
-        import autosklearn.pipeline.components.feature_preprocessing as preprocessing_components
+        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)}
 
-        classifiers = classification_components.ClassifierChoice.get_components()
-        preprocessors = preprocessing_components.PreprocessorChoice.get_components()
 
+@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,
-            "classifiers": list(classifiers.keys()),
-            "preprocessors": list(preprocessors.keys())
+            "result": {
+                "classifiers": list(_classifiers.keys()),
+                "regressors": list(_regressors.keys()),
+                "total": len(_classifiers) + len(_regressors),
+            },
+            "error": None,
         }
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
 
 
-@mcp.tool(name="meta_learning", description="Perform meta-learning using AutoSklearn")
-def meta_learning(meta_features: dict, time_limit: int, per_run_time_limit: int) -> dict:
+@mcp.tool(name="delete_classifier")
+def delete_classifier(classifier_id: str) -> dict:
     """
-    Perform meta-learning using AutoSklearn.
+    Delete a stored classifier.
 
     Parameters:
-    - meta_features: A dictionary of meta-features for the dataset.
-    - time_limit: Total time limit for the meta-learning process.
-    - per_run_time_limit: Time limit for each meta-learning iteration.
+        classifier_id (str): ID of the classifier to delete.
 
     Returns:
-    - dict: Meta-learning results.
+        dict: Confirmation of deletion.
     """
     try:
-        from autosklearn.metalearning import MetaLearning
-
-        meta_learner = MetaLearning(meta_features)
-        meta_learner.run(time_limit=time_limit, per_run_time_limit=per_run_time_limit)
-
+        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,
-            "message": "Meta-learning completed successfully",
-            "recommendations": meta_learner.get_recommendations()
+            "result": {"message": f"Classifier '{classifier_id}' deleted"},
+            "error": None,
         }
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
 
 
-@mcp.tool(name="get_model_leaderboard", description="Retrieve the leaderboard of models")
-def get_model_leaderboard(classifier) -> dict:
+@mcp.tool(name="delete_regressor")
+def delete_regressor(regressor_id: str) -> dict:
     """
-    Retrieve the leaderboard of models from AutoSklearn.
+    Delete a stored regressor.
 
     Parameters:
-    - classifier: The AutoSklearnClassifier instance.
+        regressor_id (str): ID of the regressor to delete.
 
     Returns:
-    - dict: Leaderboard information.
+        dict: Confirmation of deletion.
     """
     try:
-        leaderboard = classifier.leaderboard()
+        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,
-            "leaderboard": leaderboard
+            "result": {"message": f"Regressor '{regressor_id}' deleted"},
+            "error": None,
         }
     except Exception as e:
-        return {"success": False, "error": str(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