# AI-ML-Lab-MCP

An [MCP](https://modelcontextprotocol.io/) server, built with [FastMCP](https://github.com/jlowin/fastmcp), that turns Claude (or any MCP-compatible client) into an interactive machine-learning lab. Load a dataset, inspect it, clean it, explore it, engineer features, train and compare models, and look at the results, all through natural language.

> **Status: early development.** The data, preprocessing, regression, classification and model-comparison parts work and have been smoke-tested. Prediction, tuning, explainability and experiment tracking are not built yet (see [Roadmap](#roadmap)).

## Example workflow

Ask your MCP client something like:

> Load `cars.csv`, remove duplicates and missing values, encode the categorical columns, split the data to predict `selling_price`, compare all the regression models, and plot the feature importance of the best one.

Behind the scenes the client chains tools such as `load_csv` -> `remove_duplicates` -> `handle_missing_values` -> `encode_categorical` -> `train_test_split` -> `compare_models` -> `feature_importance_plot`.

## Features

The server currently registers **97 tools**, grouped as follows.

| Area | Tools |
|---|---|
| **Dataset management** | `load_csv`, `load_excel`, `load_json`, `load_parquet`, `load_datasets_into_dataframe`, `list_datasets`, `get_dataset_metadata`, `export_dataset`, `delete_dataset`, `undo_last_change` |
| **Inspection** | `dataset_info`, `dataset_shape`, `list_columns`, `column_types`, `describe_dataset`, `unique_values`, `value_counts`, `missing_values`, `duplicate_rows`, `memory_usage` |
| **Cleaning** | `remove_duplicates`, `handle_missing_values`, `fill_missing_mean` / `median` / `mode`, `drop_missing`, `convert_data_types`, `detect_outliers`, `remove_outliers`, `cap_outliers`, `rename_column`, `drop_column` |
| **EDA** | `analyze_distribution`, `analyze_correlation`, `analyze_relationship`, `analyze_categorical`, `find_high_correlations`, `find_target_relationships` |
| **Visualization** | `histogram`, `boxplot`, `scatterplot`, `barplot`, `lineplot`, `heatmap`, `pairplot`, `countplot`, `violinplot`, `auto_visualize`; model plots: `feature_importance_plot`, `residual_plot`, `prediction_plot`, `roc_curve`, `precision_recall_curve`, `confusion_matrix_plot`, `learning_curve` |
| **Preprocessing** | `train_test_split`, `split_features_target`, `get_split_info`, `one_hot_encode`, `label_encode`, `encode_categorical`, `scale_features`, `standardize_features`, `normalize_features`, `impute_missing`, `polynomial_features`, `feature_selection` |
| **Feature engineering** | `create_feature`, `transform_feature`, `log_transform`, `bin_feature`, `extract_date_features`, `interaction_features`, `select_features`, `suggest_features` |
| **Regression** | `linear_regression`, `ridge_regression`, `lasso_regression`, `elastic_net`, `decision_tree_regressor`, `random_forest_regressor`, `gradient_boosting_regressor`, `xgboost_regressor`, `regression_metrics` |
| **Classification** | `logistic_regression`, `knn_classifier`, `decision_tree_classifier`, `random_forest_classifier`, `svm_classifier`, `gradient_boosting_classifier`, `xgboost_classifier`, `classification_report`, `confusion_matrix` |
| **Model comparison** | `compare_models`, `benchmark_models` (cross-validation on the training set only), `rank_models` |

### Design notes

- **No data leakage by default.** Transform tools accept either a dataset id or a split id. For a split, transformers are fitted on `X_train` only and then applied to `X_test`.
- **Undo support.** Cleaning operations record history, so `undo_last_change` can roll a dataset back.
- **Stored steps.** Every preprocessing and feature-engineering step is recorded on the split (`Splits[id]["steps"]`) and snapshotted onto each trained model, which is the groundwork for replaying them at prediction time.
- **Guard rails.** Tools return clear error messages instead of tracebacks (for example when a feature is not numeric, contains NaN, or a classification target has more than 100 distinct values). Training tools add warnings for overfitting, class imbalance and results that are no better than a majority-class baseline.
- **Safe feature expressions.** `create_feature` evaluates row-wise expressions only, and rejects method calls, `@`, dunder access and assignments.
- **In-memory state.** Datasets, splits and models live in memory (`storage.py`) for the lifetime of the server process.

## Requirements

- Python 3.13 or newer
- [uv](https://docs.astral.sh/uv/)
- Core dependencies: `fastmcp`, `pandas` (3.x), `matplotlib`, `scikit-learn`
- Optional: `xgboost` (XGBoost models; skipped gracefully if missing), `openpyxl` (Excel), `pyarrow` (Parquet)

## Installation

```bash
git clone https://github.com/SyedHassan904/ai-ml-lab-mcp.git
cd ai-ml-lab-mcp
uv sync
uv add xgboost          # optional, enables the XGBoost tools
```

> `uv sync` removes packages that are not listed in `pyproject.toml`. If you also use a Jupyter kernel in this folder, add it as a dev dependency first: `uv add --dev ipykernel`.

## Running the server

```bash
uv run python server.py
```

> The `ml-mcp` console script is not wired up yet (see [Known limitations](#known-limitations)). Use `server.py` for now.

### Connect to Claude Desktop

Add this to your Claude Desktop MCP configuration (adjust the path to your project):

```json
{
  "mcpServers": {
    "ai-ml-lab": {
      "command": "uv",
      "args": ["--directory", "D:\\ml_mcp", "run", "python", "server.py"]
    }
  }
}
```

Restart Claude Desktop and the tools will appear in the tools menu. Any other MCP client that can launch a stdio server works the same way.

## Project layout

```
.
├── server.py                  # creates the FastMCP app and registers all tool groups
├── storage.py                 # in-memory stores: DataSets, Splits, Models
├── experiments.py             # experiment store (stub)
├── pyproject.toml             # dependencies and project metadata
├── LICENSE                    # MIT License
├── tools/
│   ├── utils.py               # id generation, validation, undo history helpers
│   ├── dataset_management.py
│   ├── dataset_inspection.py
│   ├── data_cleaning.py
│   ├── EDA.py
│   ├── visualization.py
│   ├── preprocessing.py
│   ├── feature_engineering.py
│   ├── regression.py
│   ├── classification.py
│   └── model_comparison.py
└── src/ml_mcp/                # package placeholder (see Known limitations)
```

Each module exposes a `register_*(mcp)` function that defines its tools with `@mcp.tool`.

## Known limitations

- **Not yet tested with a real MCP client.** Tools were smoke-tested in-process through `fastmcp.Client`, not through Claude Desktop or the MCP Inspector.
- **Packaging is unfinished.** `storage.py` and `tools/` live outside `src/ml_mcp/`, so the project cannot be installed as a package and the `ml-mcp` script is a placeholder. Run from the repository root.
- **No automated test suite yet.** Verification so far was done with throwaway smoke scripts. Converting them to `pytest` tests is planned.
- **No prediction on new data yet.** Models can be trained and evaluated on a split, but preprocessing steps are not yet replayed on new rows, and models cannot be saved or loaded.
- **Plots are saved to the server's working directory** as PNG files with random names.
- **State is not persistent.** Restarting the server clears all datasets, splits and models.

## Roadmap

- [ ] Prediction tools (`predict`, `batch_predict`) with preprocessing replay
- [ ] Save / load models (joblib)
- [ ] `pytest` suite built from the current smoke tests
- [ ] Package restructure into `src/ml_mcp/` and a working `ml-mcp` command
- [ ] Cross-validation and hyperparameter tuning (grid, random, Optuna)
- [ ] AutoML (`auto_ml`)
- [ ] Explainability (SHAP) and error analysis
- [ ] Experiment tracking (`experiments.py`)
- [ ] Clustering, anomaly detection, PCA
- [ ] Report generation, code generation, FastAPI/Streamlit/Docker deployment helpers
- [ ] Dedicated output folder for plots

## Contributing

Issues and pull requests are welcome. If you add a tool, put it in the matching module (or a new one with a `register_*` function wired into `server.py`) and make sure its name does not collide with an existing tool, since FastMCP tool names must be unique. Please include a test once the `tests/` folder exists.

## License

This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.

## Author

Syed Hassan Zaheer