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

auto-sklearn/mcp_output/mcp_plugin/mcp_service.py

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