"""
Workflow tools for Stats Compass MCP server.
These wrap stats-compass-core workflow functions with session isolation.
"""
import logging
from typing import Any, List, Optional
from fastmcp import Context, FastMCP
from stats_compass_core.workflows import (
ClassificationConfig,
EDAConfig,
PreprocessingConfig,
RegressionConfig,
TimeSeriesConfig,
run_classification,
run_eda_report,
run_preprocessing,
run_regression,
run_timeseries_forecast,
)
from stats_compass_core.workflows.classification import RunClassificationInput
from stats_compass_core.workflows.eda_report import RunEDAReportInput
from stats_compass_core.workflows.preprocessing import RunPreprocessingInput
from stats_compass_core.workflows.regression import RunRegressionInput
from stats_compass_core.workflows.timeseries import RunTimeseriesForecastInput
from stats_compass_mcp.exports import save_plot_export
from stats_compass_mcp.image_utils import with_images
from stats_compass_mcp.session import Session, SessionManager, get_session
logger = logging.getLogger(__name__)
def save_workflow_exports(result_dict: dict, session: Session, workflow_name: str) -> dict:
"""
Save all exportable artifacts from a workflow result.
- Saves plots (base64 images) to exports/plots/
- Returns modified result with download_urls added
"""
downloads = []
# Process steps to find and save images
steps = result_dict.get("steps", [])
for i, step in enumerate(steps):
# image_base64 is a top-level field on WorkflowStepResult, not inside "result"
image_b64 = step.get("image_base64")
if image_b64:
step_name = step.get("step_name", f"step_{i}")
name_prefix = f"{workflow_name}_{step_name}"
export_info = save_plot_export(
session_id=session.session_id,
image_base64=image_b64,
name_prefix=name_prefix,
)
if export_info["download_url"]:
# Add download info to the step
step["download_url"] = export_info["download_url"]
step["filename"] = export_info["filename"]
downloads.append({
"type": "plot",
"name": export_info["filename"],
"step": step_name,
"url": export_info["download_url"],
})
# Add downloads summary to result
if downloads:
result_dict["downloads"] = downloads
return result_dict
def register_workflow_tools(mcp: FastMCP, session_manager: SessionManager):
"""Register all workflow tools with the FastMCP server."""
@mcp.tool()
def run_eda_report_workflow(
ctx: Context,
dataframe_name: Optional[str] = None,
config: Optional[dict] = None
) -> Any:
"""
Run comprehensive EDA report: descriptive stats, correlations,
missing data analysis, and auto-generated visualizations.
Args:
dataframe_name: Name of DataFrame to analyze (default: active)
config: Optional EDA configuration dict
Returns:
Workflow result with steps, metrics, and charts.
"""
session = get_session(ctx, session_manager)
eda_config = EDAConfig(**config) if config else None
params = RunEDAReportInput(dataframe_name=dataframe_name, config=eda_config)
result = run_eda_report(state=session.state, params=params)
result_dict = save_workflow_exports(result.model_dump(), session, "eda")
return with_images(result_dict, summarize=True)
@mcp.tool()
def run_preprocessing_workflow(
ctx: Context,
dataframe_name: Optional[str] = None,
save_as: Optional[str] = None,
config: Optional[dict] = None
) -> Any:
"""
Run data preprocessing pipeline: analyze missing data, apply imputation,
handle outliers, and remove duplicates.
Args:
dataframe_name: Name of DataFrame to preprocess (default: active)
save_as: Name for the cleaned DataFrame (default: auto-generated)
config: Optional preprocessing configuration dict
Returns:
Workflow result with steps and cleaned DataFrame name.
"""
session = get_session(ctx, session_manager)
preproc_config = PreprocessingConfig(**config) if config else None
params = RunPreprocessingInput(
dataframe_name=dataframe_name, save_as=save_as, config=preproc_config
)
result = run_preprocessing(state=session.state, params=params)
result_dict = save_workflow_exports(result.model_dump(), session, "preprocessing")
return with_images(result_dict, summarize=True)
@mcp.tool()
def run_classification_workflow(
ctx: Context,
target_column: str,
dataframe_name: Optional[str] = None,
feature_columns: Optional[List[str]] = None,
config: Optional[dict] = None
) -> Any:
"""
Run classification workflow: train model, evaluate performance,
generate confusion matrix, ROC curve, and feature importance plots.
Args:
target_column: Column with class labels to predict
dataframe_name: Name of DataFrame (default: active)
feature_columns: Feature columns (default: all numeric except target)
config: Optional classification configuration dict
Returns:
Workflow result with metrics, model ID, and diagnostic charts.
"""
session = get_session(ctx, session_manager)
class_config = ClassificationConfig(**config) if config else None
params = RunClassificationInput(
dataframe_name=dataframe_name,
target_column=target_column,
feature_columns=feature_columns,
config=class_config
)
result = run_classification(state=session.state, params=params)
result_dict = save_workflow_exports(result.model_dump(), session, "classification")
return with_images(result_dict, summarize=True)
@mcp.tool()
def run_regression_workflow(
ctx: Context,
target_column: str,
dataframe_name: Optional[str] = None,
feature_columns: Optional[List[str]] = None,
config: Optional[dict] = None
) -> Any:
"""
Run regression workflow: train model, evaluate with RMSE/MAE/R²,
generate feature importance plots.
Args:
target_column: Column with continuous values to predict
dataframe_name: Name of DataFrame (default: active)
feature_columns: Feature columns (default: all numeric except target)
config: Optional regression configuration dict
Returns:
Workflow result with metrics, model ID, and charts.
"""
session = get_session(ctx, session_manager)
reg_config = RegressionConfig(**config) if config else None
params = RunRegressionInput(
dataframe_name=dataframe_name,
target_column=target_column,
feature_columns=feature_columns,
config=reg_config
)
result = run_regression(state=session.state, params=params)
result_dict = save_workflow_exports(result.model_dump(), session, "regression")
return with_images(result_dict, summarize=True)
@mcp.tool()
def run_timeseries_workflow(
ctx: Context,
target_column: str,
dataframe_name: Optional[str] = None,
date_column: Optional[str] = None,
config: Optional[dict] = None
) -> Any:
"""
Run time series forecasting: check stationarity, fit ARIMA model,
generate forecasts and visualization.
Args:
target_column: Column with values to forecast
dataframe_name: Name of DataFrame (default: active)
date_column: Column with dates (default: uses index)
config: Optional time series configuration dict
Returns:
Workflow result with forecasts and forecast chart.
"""
session = get_session(ctx, session_manager)
ts_config = TimeSeriesConfig(**config) if config else None
params = RunTimeseriesForecastInput(
dataframe_name=dataframe_name,
target_column=target_column,
date_column=date_column,
config=ts_config
)
result = run_timeseries_forecast(state=session.state, params=params)
result_dict = save_workflow_exports(result.model_dump(), session, "timeseries")
return with_images(result_dict, summarize=True)
