"""
Executor for sktime MCP.
Responsible for instantiating estimators, loading datasets,
and running fit/predict operations.
"""
import asyncio
import inspect
import logging
import uuid
from collections import deque
from typing import Any
import pandas as pd
from sktime_mcp.registry.interface import get_registry
from sktime_mcp.runtime.handles import get_handle_manager
from sktime_mcp.runtime.jobs import JobStatus, get_job_manager
logger = logging.getLogger(__name__)
# Dynamically discover all available sktime demo datasets at import time.
# This replaces the old hardcoded dictionary and automatically exposes every
# load_* function in sktime.datasets to the MCP server.
def _discover_demo_datasets() -> dict:
"""Return a mapping of dataset name -> dotted module path for every
zero-argument ``load_*`` function exported by ``sktime.datasets``."""
try:
import sktime.datasets as _ds_module
return {
name.removeprefix("load_"): f"sktime.datasets.{name}"
for name, obj in inspect.getmembers(_ds_module, inspect.isfunction)
if name.startswith("load_")
and all(
p.default is not inspect.Parameter.empty
for p in inspect.signature(obj).parameters.values()
if p.kind not in (p.VAR_POSITIONAL, p.VAR_KEYWORD)
)
}
except Exception: # pragma: no cover
return {} # fallback: empty dict if sktime not installed
_DEMO_DATASETS: dict | None = None
def _get_demo_datasets() -> dict:
"""Lazy singleton — discovers datasets only on first call."""
global _DEMO_DATASETS
if _DEMO_DATASETS is None:
_DEMO_DATASETS = _discover_demo_datasets()
return _DEMO_DATASETS
def _to_period_index_if_possible(obj: Any) -> Any:
"""Return *obj* with a ``PeriodIndex`` when its index is a regular datetime index.
Seasonal sktime forecasters coerce the series index to a ``PeriodIndex``
internally (``index.to_period(freq)``) and raise on offset frequencies such
as ``MonthBegin`` ("MS"), which is what ``load_data_source`` produces for
monthly data. Demo datasets carry a ``PeriodIndex`` and work, so we
normalise handle-loaded data to match. ``to_period()`` is called with no
argument so pandas maps the offset to its period alias (MS -> "M"); passing
the offset string back in would re-raise the same error.
No-op for non-datetime indexes, ``PeriodIndex`` already, or an index with no
determinable frequency.
"""
if obj is None or not hasattr(obj, "index"):
return obj
idx = obj.index
if isinstance(idx, pd.PeriodIndex) or not isinstance(idx, pd.DatetimeIndex):
return obj
try:
if idx.freq is None:
inferred = pd.infer_freq(idx)
if inferred is None:
return obj
idx = pd.DatetimeIndex(idx, freq=inferred)
converted = obj.copy()
converted.index = idx.to_period()
return converted
except (ValueError, TypeError):
return obj
# Max forecast rows returned inline before truncation (NB-22). Normal horizons
# (<= a few dozen) are never affected; a 1000-step forecast would otherwise
# flood the client with ~30KB+ of inline JSON.
_MAX_PREDICTION_ROWS = 500
# Dunder methods that are safe and useful to call via call_method (e.g. __call__
# for callable metrics/aligners). Everything else starting with "_" is blocked
# (BUG-11) — notably __reduce__/__class__/__getattribute__ and private methods.
_ALLOWED_DUNDERS = frozenset({"__call__", "__len__", "__repr__", "__str__"})
def _is_sktime_object(obj: Any) -> bool:
"""True if *obj* is a genuine sktime estimator/object, not a bare value.
craft evaluates arbitrary specs, so a spec like "42" returns an int. Such
non-objects should not receive an estimator handle (BUG-10). We accept
anything deriving from skbase's BaseObject, falling back to a duck-typed
check for get_params + a scitype tag.
"""
try:
from skbase.base import BaseObject
if isinstance(obj, BaseObject):
return True
except Exception: # pragma: no cover - skbase always present with sktime
pass
return hasattr(obj, "get_params") and hasattr(obj, "get_class_tag")
def _cap_prediction_rows(result: dict) -> tuple[dict, dict | None]:
"""Cap an index-keyed prediction dict, returning (capped, truncation_note)."""
if not isinstance(result, dict) or len(result) <= _MAX_PREDICTION_ROWS:
return result, None
total = len(result)
kept = dict(list(result.items())[:_MAX_PREDICTION_ROWS])
note = {
"shown": _MAX_PREDICTION_ROWS,
"total": total,
"note": (
"forecast truncated; request a smaller horizon or use save_data to write "
"the full series to a file"
),
}
return kept, note
def _get_index_frequency_metadata(
index: pd.Index,
fallback: str | None = None,
) -> str | None:
"""Return a stable frequency label for metadata without assuming datetime-only indexes."""
if isinstance(index, (pd.DatetimeIndex, pd.PeriodIndex)):
freq = getattr(index, "freq", None)
if freq is not None:
return str(freq)
inferred = pd.infer_freq(index)
if inferred is not None:
return inferred
return fallback
def _resolve_metric_scoring(metric_name: str) -> Any | None:
"""Return an instantiated sktime forecasting metric by name, or None if not found."""
try:
from sktime.registry import all_estimators
except ImportError: # pragma: no cover
return None
try:
metrics_df = all_estimators("metric", as_dataframe=True)
row = metrics_df[metrics_df["name"] == metric_name]
if row.empty:
return None
return row.iloc[0]["object"]()
except Exception as e:
logger.warning(f"Failed to resolve metric '{metric_name}': {e}")
return None
def _run_evaluate(
instance: Any,
y: Any,
X: Any,
cv_folds: int,
scoring: Any | None,
initial_window: int | None,
) -> tuple[list[dict[str, Any]], dict[str, float], dict[str, dict[str, float]]]:
"""
Run sktime.evaluate with an expanding-window splitter and summarize results.
Returns
-------
fold_results : list of dict
Per-fold rows from sktime.evaluate.
metrics : dict
Mean value per ``test_*`` metric column.
summary : dict
Mean, std, min, max per ``test_*`` metric column.
"""
from sktime.forecasting.model_evaluation import evaluate
try:
from sktime.split import ExpandingWindowSplitter
except ImportError: # pragma: no cover - sktime < 0.29
from sktime.forecasting.model_selection import ExpandingWindowSplitter
n = len(y)
if initial_window is not None:
if not 1 <= initial_window < n:
raise ValueError(
f"initial_window must be between 1 and n-1={n - 1} "
f"(series has {n} observations), got {initial_window}"
)
win = initial_window
else:
folds = int(cv_folds)
if not 1 <= folds <= n - 1:
raise ValueError(
f"cv_folds must be between 1 and n-1={n - 1} "
f"(series has {n} observations), got {folds}"
)
win = n - folds
cv = ExpandingWindowSplitter(initial_window=win, step_length=1, fh=[1])
# error_score="raise" — sktime's default (np.nan) swallows per-fold
# exceptions and reports success with all-NaN metrics
results = evaluate(
forecaster=instance, y=y, X=X, cv=cv, scoring=scoring, error_score="raise"
)
if "estimator" in results.columns:
results = results.drop(columns=["estimator"])
fold_results = results.to_dict(orient="records")
metric_cols = [
c for c in results.select_dtypes(include="number").columns if c.startswith("test_")
]
metrics = {c: float(results[c].mean()) for c in metric_cols}
summary = {
c: {
"mean": float(results[c].mean()),
"std": float(results[c].std()),
"min": float(results[c].min()),
"max": float(results[c].max()),
}
for c in metric_cols
}
return fold_results, metrics, summary
def _merge_adapter_validation_warnings(
validation_report: dict[str, Any],
metadata: dict[str, Any],
) -> dict[str, Any]:
"""Merge warnings added during adapter conversion into validation output."""
metadata_validation = metadata.get("validation")
if not isinstance(metadata_validation, dict):
return validation_report
metadata_warnings = metadata_validation.get("warnings", [])
if not metadata_warnings:
return validation_report
merged = validation_report.copy()
existing_warnings = list(merged.get("warnings", []))
for warning in metadata_warnings:
if warning not in existing_warnings:
existing_warnings.append(warning)
merged["warnings"] = existing_warnings
return merged
[docs]
class Executor:
"""
Execution runtime for sktime estimators.
Handles instantiation, fitting, and prediction.
"""
[docs]
def __init__(self):
self._registry = get_registry()
self._handle_manager = get_handle_manager()
self._job_manager = get_job_manager()
self._data_handles: dict[str, Any] = {}
# Tombstones for data handles evicted under the cap (see _cleanup_oldest_data).
self._evicted_data: deque[str] = deque(maxlen=1024)
from sktime_mcp.config import settings
self._max_data_handles = settings.max_data_handles
self._auto_format_enabled = settings.auto_format
def _cleanup_oldest_data(self, count: int = 10) -> None:
to_remove = list(self._data_handles.keys())[:count]
for handle_id in to_remove:
del self._data_handles[handle_id]
self._evicted_data.append(handle_id)
logger.info("Evicted data handle %s (limit %d reached)", handle_id, self._max_data_handles)
[docs]
def data_handle_missing(self, handle_id: str) -> dict[str, Any]:
"""Error body for a missing data handle — distinguishes evicted from unknown.
Returns the ``error`` string plus the capped available-handles summary,
so callers can splat it into a not-found response.
"""
if handle_id in self._evicted_data:
error = (
f"Data handle '{handle_id}' was evicted (handle limit "
f"{self._max_data_handles} reached); reload the source."
)
else:
error = f"Data handle '{handle_id}' not found"
return {"error": error, **self.summarize_available_handles()}
def _register_data_handle(self, handle_id: str, data: dict[str, Any]) -> None:
if len(self._data_handles) >= self._max_data_handles:
self._cleanup_oldest_data(count=max(1, self._max_data_handles // 5))
self._data_handles[handle_id] = data
[docs]
def summarize_available_handles(self, limit: int = 5) -> dict[str, Any]:
"""Capped view of data-handle ids for not-found error responses.
Returns the *limit* most recent handles plus the total count, so
error responses stay small and don't enumerate every handle in the
process.
"""
handle_ids = list(self._data_handles.keys())
return {
"available_handles": handle_ids[-limit:],
"n_available_handles": len(handle_ids),
}
def _resolve_source(self, source: str, prefer: str = "y") -> dict[str, Any]:
"""Resolve a source id to a series, trying data_handle then demo dataset.
``prefer`` selects which component of a demo dataset to return
("y" or "X"); the other is the fallback when the preferred one is
absent. Data handles always resolve to their primary series.
"""
if source in self._data_handles:
return {"success": True, "data": self._data_handles[source]["y"]}
res = self.load_dataset(source)
if res["success"]:
first, second = ("X", "y") if prefer == "X" else ("y", "X")
data = res[first] if res[first] is not None else res[second]
return {"success": True, "data": data}
return res
[docs]
def instantiate(
self,
spec: str,
) -> dict[str, Any]:
"""Instantiate an estimator or pipeline from a spec and return a handle."""
import importlib
importlib.invalidate_caches()
try:
from sktime.utils.dependencies._dependencies import _get_installed_packages_private
_get_installed_packages_private.cache_clear()
except ImportError:
pass
import numpy as np
import pandas as pd
import sktime.registry._craft as _craft_module
from sktime.registry import craft
# Temporarily patch all_estimators to inject standard libraries into craft's registry.
# This allows users to pass callables like `numpy.exp` into estimators
# like CurveFitForecaster via the craft spec.
original_all = _craft_module.all_estimators
def mock_all_estimators(*args, **kwargs):
results = original_all(*args, **kwargs)
# results is a list of tuples: [(name, class), ...]
# We append numpy and pandas so they enter the register dict!
results.append(("np", np))
results.append(("numpy", np))
results.append(("pd", pd))
results.append(("pandas", pd))
return results
_craft_module.all_estimators = mock_all_estimators
try:
try:
instance = craft(spec)
finally:
_craft_module.all_estimators = original_all
# Reject specs that don't produce an sktime object — e.g. "42",
# "[1,2,3]", "None" otherwise got est_ handles that failed
# confusingly downstream (BUG-10).
if not _is_sktime_object(instance):
return {
"success": False,
"error": (
f"Spec did not produce an sktime estimator, got "
f"{type(instance).__name__}. Provide a craft spec such as "
"'NaiveForecaster(sp=12)' or 'Detrender() * ARIMA()'."
),
}
estimator_name = type(instance).__name__
handle_id = self._handle_manager.create_handle(
estimator_name=estimator_name,
instance=instance,
params={"spec": spec},
)
return {
"success": True,
"handle": handle_id,
"estimator": estimator_name,
"spec": spec,
}
except Exception as e:
import sys
error_msg = str(e)
if (
"requires package" in error_msg
or "pip install" in error_msg
or "ModuleNotFoundError" in type(e).__name__
):
error_msg += f"\n\n(Hint for AI: To install missing dependencies, use the server's exact python environment by running: `{sys.executable} -m pip install <package_name>`)"
logger.error("fit failed: %s", e, exc_info=True)
return {
"success": False,
"error": error_msg,
}
# L-7: We can also add custom load_dataset functions here
[docs]
def load_dataset(self, name: str) -> dict[str, Any]:
"""Load a demo dataset.
Returns canonical keys with one consistent meaning for every
dataset family: ``y`` is always the target/labels, ``X`` is always
the features/panel (or None).
"""
demo_datasets = _get_demo_datasets()
if name not in demo_datasets:
return {
"success": False,
"error": f"Unknown dataset: {name}",
"available": list(demo_datasets.keys()),
}
try:
module_path = demo_datasets[name]
parts = module_path.rsplit(".", 1)
module = __import__(parts[0], fromlist=[parts[1]])
loader = getattr(module, parts[1])
data = loader()
if isinstance(data, tuple):
# sktime classifier/clusterer datasets return (X-panel, y-labels)
# whereas forecaster datasets return (y-target, X-exog)
if name in (
"arrow_head",
"italy_power_demand",
"basic_motions",
"gunpoint",
"osuleaf",
"plaid",
):
X, y = data[0], data[1] if len(data) > 1 else None
primary = X
else:
y, X = data[0], data[1] if len(data) > 1 else None
primary = y
else:
y, X = data, None
primary = y
return {
"success": True,
"name": name,
"shape": primary.shape if hasattr(primary, "shape") else len(primary),
"type": str(type(primary).__name__),
"y": y,
"X": X,
}
except Exception as e:
return {"success": False, "error": str(e)}
[docs]
def fit(
self,
handle_id: str,
y: Any,
X: Any | None = None,
fh: Any | None = None,
) -> dict[str, Any]:
"""Fit an estimator."""
try:
handle_info = self._handle_manager.get_info(handle_id)
instance = handle_info.instance
except KeyError:
return {"success": False, "error": self._handle_manager.describe_missing(handle_id)}
obj_type = getattr(instance, "get_class_tag", lambda x, y: "")("object_type", "")
if not hasattr(instance, "fit"):
return {
"success": False,
"error": f"The {obj_type or 'estimator'} scitype does not support fit(). Please use the 'call_method' tool to interact with its native methods.",
}
# Check scitype to determine how to call fit
# By default in sktime:
# - Forecasters: fit(y, X=None, fh=None)
# - Classifiers/Regressors: fit(X, y)
# - Transformers/Clusterers: fit(X, y=None)
is_classifier_or_regressor = False
is_transformer = False
if hasattr(instance, "get_class_tag"):
obj_type = instance.get_class_tag("object_type", "")
if obj_type in ("classifier", "regressor"):
is_classifier_or_regressor = True
elif obj_type == "transformer":
is_transformer = True
try:
if is_classifier_or_regressor:
# With decoupled X and y handles, X is features and y is labels
instance.fit(X, y)
elif is_transformer:
if X is not None:
instance.fit(y, X)
else:
instance.fit(y)
elif obj_type == "clusterer":
if y is not None:
instance.fit(X, y)
else:
instance.fit(X)
else:
# Assume forecaster or similar default
if fh is not None:
instance.fit(y, X=X, fh=fh)
elif X is not None:
instance.fit(y, X=X)
else:
instance.fit(y)
self._handle_manager.mark_fitted(handle_id)
return {"success": True, "handle": handle_id, "fitted": True}
except Exception as e:
logger.error("%s failed: %s", type(e).__name__, e, exc_info=True)
return {"success": False, "error": str(e)}
[docs]
def predict(
self,
handle_id: str,
fh: int | list[int] | None = None,
X: Any | None = None,
y: Any | None = None,
mode: str = "predict",
coverage: float | list[float] = 0.9,
alpha: float | list[float] | None = None,
) -> dict[str, Any]:
"""Generate predictions."""
try:
instance = self._handle_manager.get_instance(handle_id)
except KeyError:
return {"success": False, "error": self._handle_manager.describe_missing(handle_id)}
obj_type = getattr(instance, "get_class_tag", lambda x, y: "")("object_type", "")
if (
not hasattr(instance, "predict")
and mode == "predict"
and not (hasattr(instance, "transform") and obj_type == "transformer")
):
return {
"success": False,
"error": f"The {obj_type or 'estimator'} scitype does not support predict(). Please use the 'call_method' tool to interact with its native methods.",
}
if not self._handle_manager.is_fitted(handle_id):
return {"success": False, "error": "Estimator not fitted"}
is_classifier_or_regressor = False
is_transformer = False
if hasattr(instance, "get_class_tag"):
obj_type = instance.get_class_tag("object_type", "")
if obj_type in ("classifier", "regressor"):
is_classifier_or_regressor = True
elif obj_type in ("transformer", "clusterer"):
is_transformer = True
dropped_y_warning = None
try:
if fh is None and not (is_classifier_or_regressor or is_transformer):
fh = list(range(1, 13))
kwargs = {}
if X is not None:
kwargs["X"] = X
if y is not None:
# y at predict is only for annotators; forwarding it to a
# forecaster raised a raw "unexpected keyword argument 'y'"
# TypeError (NB-18). Only pass it when predict accepts it.
accepts_y = False
try:
accepts_y = "y" in inspect.signature(instance.predict).parameters
except (ValueError, TypeError):
accepts_y = False
if accepts_y:
kwargs["y"] = y
else:
dropped_y_warning = (
f"y was ignored: {obj_type or 'this estimator'}.predict() does not "
"accept y (it is only used by annotators/detectors)."
)
if is_classifier_or_regressor:
# Classifiers take X in predict (X is the feature matrix)
# But instance.predict(X) is the signature.
# Since kwargs["X"] has it, we can just pass X positionally
if mode == "predict":
predictions = instance.predict(X)
elif mode == "predict_proba":
predictions = instance.predict_proba(X)
else:
return {"success": False, "error": f"Mode {mode} not supported for {obj_type}"}
elif is_transformer:
if mode == "predict":
if obj_type == "clusterer":
predictions = (
instance.predict(X) if X is not None else instance.predict(fh=fh)
) # some clusterers might use predict(X)
else:
# For transformer, transform is basically the predict equivalent if X is passed
if X is not None:
predictions = instance.transform(X)
else:
return {"success": False, "error": "Transform requires X"}
else:
return {"success": False, "error": f"Mode {mode} not supported for {obj_type}"}
else:
if mode == "predict":
predictions = instance.predict(fh=fh, **kwargs)
elif mode == "predict_interval":
predictions = instance.predict_interval(fh=fh, coverage=coverage, **kwargs)
elif mode == "predict_quantiles":
predictions = instance.predict_quantiles(fh=fh, alpha=alpha, **kwargs)
elif mode == "predict_proba":
predictions = instance.predict_proba(fh=fh, **kwargs)
elif mode == "predict_var":
predictions = instance.predict_var(fh=fh, **kwargs)
else:
return {"success": False, "error": f"Unknown prediction mode: {mode}"}
from sktime_mcp.server import sanitize_for_json
truncated_note = None
if isinstance(predictions, pd.Series):
predictions_copy = predictions.copy()
predictions_copy.index = predictions_copy.index.astype(str)
result, truncated_note = _cap_prediction_rows(predictions_copy.to_dict())
elif isinstance(predictions, pd.DataFrame):
predictions_copy = predictions.copy()
predictions_copy.index = predictions_copy.index.astype(str)
# Flatten multiindex columns (predict_interval/quantiles) for JSON.
if isinstance(predictions_copy.columns, pd.MultiIndex):
predictions_copy.columns = [
"_".join(map(str, col)) for col in predictions_copy.columns.values
]
# orient="index" keeps the time index as the key so interval /
# variance values map to time points, consistent with predict
# (NB-21). orient="list" dropped the index entirely.
result, truncated_note = _cap_prediction_rows(
predictions_copy.to_dict(orient="index")
)
else:
result = sanitize_for_json(predictions)
out = {
"success": True,
"mode": mode,
}
# horizon is only meaningful for forecasters; echoing it for
# classifiers/regressors/transformers implied a truncation that
# didn't happen (N-01).
if not (is_classifier_or_regressor or is_transformer):
out["horizon"] = len(fh) if hasattr(fh, "__len__") else fh
if mode == "predict":
out["predictions"] = result
elif mode == "predict_interval":
out["intervals"] = result
out["coverage"] = coverage
elif mode == "predict_quantiles":
out["quantiles"] = result
out["alpha"] = alpha
else:
out["predictions"] = result
if truncated_note:
out["predictions_truncated"] = truncated_note
if dropped_y_warning:
out["warnings"] = [dropped_y_warning]
return out
except Exception as e:
return {"success": False, "error": str(e)}
[docs]
async def predict_async(
self,
handle_id: str,
*,
horizon: int = 12,
mode: str = "predict",
coverage: float | list[float] = 0.9,
alpha: float | list[float] | None = None,
X_dataset: str | None = None,
y_dataset: str | None = None,
X_handle: str | None = None,
y_handle: str | None = None,
job_id: str | None = None,
) -> dict[str, Any]:
"""Async version of predict with job tracking."""
try:
self._job_manager.update_job(job_id, status=JobStatus.RUNNING)
# Step 1: Load data
self._job_manager.update_job(job_id, completed_steps=0, current_step="Loading data...")
await asyncio.sleep(0.01)
X = None
y = None
if X_handle:
if X_handle not in self._data_handles:
raise ValueError(f"Unknown X data handle: {X_handle}")
X = self._data_handles[X_handle]["y"]
if y_handle:
if y_handle not in self._data_handles:
raise ValueError(f"Unknown y data handle: {y_handle}")
y = self._data_handles[y_handle]["y"]
if X_dataset and X_dataset == y_dataset:
data_res = self.load_dataset(X_dataset)
if not data_res["success"]:
raise ValueError(data_res.get("error", "Failed to load dataset"))
y = data_res["y"]
X = data_res["X"]
else:
if X_dataset:
data_res = self.load_dataset(X_dataset)
if not data_res["success"]:
raise ValueError(data_res.get("error", "Failed to load dataset"))
X = data_res["X"] if data_res["X"] is not None else data_res["y"]
if y_dataset:
data_res = self.load_dataset(y_dataset)
if not data_res["success"]:
raise ValueError(data_res.get("error", "Failed to load dataset"))
y = data_res["y"]
fh = list(range(1, horizon + 1))
# Step 2: Generate predictions
self._job_manager.update_job(
job_id, completed_steps=1, current_step="Generating predictions..."
)
await asyncio.sleep(0.01)
loop = asyncio.get_running_loop()
result = await loop.run_in_executor(
None,
lambda: self.predict(
handle_id, fh=fh, X=X, y=y, mode=mode, coverage=coverage, alpha=alpha
),
)
if not result.get("success"):
self._job_manager.update_job(
job_id,
status=JobStatus.FAILED,
current_step="Prediction failed.",
errors=[result.get("error", "Unknown error")],
)
return result
self._job_manager.update_job(
job_id,
status=JobStatus.COMPLETED,
completed_steps=2,
current_step="Prediction completed.",
result=result,
)
return result
except Exception as e:
self._job_manager.update_job(
job_id,
status=JobStatus.FAILED,
current_step="Prediction failed.",
errors=[str(e)], # traceback logged server-side, not leaked to the client
)
return {"success": False, "error": str(e)}
[docs]
def call_method(
self,
handle_id: str,
method_name: str,
kwargs: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""Dynamically call a method on the underlying estimator."""
try:
instance = self._handle_manager.get_instance(handle_id)
except KeyError:
return {"success": False, "error": self._handle_manager.describe_missing(handle_id)}
# Block private/dunder methods: they are not part of the estimator API
# and expose internals — e.g. __reduce__ dumps __dict__ including the
# fitted _y/_X training data to any caller (BUG-11).
if method_name.startswith("_") and method_name not in _ALLOWED_DUNDERS:
return {
"success": False,
"error": (
f"Method '{method_name}' is private and not callable via call_method. "
"Only public estimator methods are exposed."
),
}
if not hasattr(instance, method_name):
obj_type = getattr(instance, "get_class_tag", lambda x, y: "")("object_type", "")
return {
"success": False,
"error": f"The {obj_type or 'estimator'} does not have a method '{method_name}'.",
}
kwargs = kwargs or {}
try:
method = getattr(instance, method_name)
# Map data_handle and dataset from kwargs if they exist
# This allows the LLM to pass 'dataset': 'airline' and we inject the actual data
for k, v in list(kwargs.items()):
if k.endswith("_dataset") and isinstance(v, str):
data_res = self.load_dataset(v)
if not data_res.get("success"):
error_res = {
"success": False,
"error": data_res.get("error", f"Unknown dataset: {v}"),
}
if "available" in data_res:
error_res["available"] = data_res["available"]
return error_res
# Replace the kwarg with the actual data (e.g. y_dataset -> y);
# the prefix selects the dataset component
actual_key = k.replace("_dataset", "")
if actual_key == "X":
value = data_res["X"] if data_res["X"] is not None else data_res["y"]
else:
value = data_res["y"]
kwargs[actual_key] = value
del kwargs[k]
elif k.endswith("_data_handle") and isinstance(v, str):
if v in self._data_handles:
actual_key = k.replace("_data_handle", "")
kwargs[actual_key] = self._data_handles[v]["y"]
del kwargs[k]
else:
return {"success": False, "error": f"Unknown data handle: {v}"}
result = method(**kwargs)
# Materialize generators (e.g. splitter.split) so the caller gets
# the actual values instead of a useless repr string
if inspect.isgenerator(result):
result = list(result)
from sktime_mcp.server import sanitize_for_json
if hasattr(result, "to_dict"):
if isinstance(result, __import__("pandas").DataFrame) and isinstance(
result.columns, __import__("pandas").MultiIndex
):
result.columns = ["_".join(map(str, col)) for col in result.columns.values]
sanitized = result.to_dict(orient="list")
else:
sanitized = result.to_dict()
else:
sanitized = sanitize_for_json(result)
return {"success": True, "result": sanitized}
except Exception as e:
logger.error("%s failed: %s", type(e).__name__, e, exc_info=True)
return {"success": False, "error": str(e)}
[docs]
def update(
self,
handle_id: str,
y: Any,
X: Any | None = None,
update_params: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""Update a fitted estimator with new data."""
try:
instance = self._handle_manager.get_instance(handle_id)
except KeyError:
return {"success": False, "error": self._handle_manager.describe_missing(handle_id)}
if not self._handle_manager.is_fitted(handle_id):
return {"success": False, "error": "Estimator not fitted"}
if y is None:
return {
"success": False,
"error": (
"update requires new data — provide y_handle or y_dataset "
"(and optionally X_handle/X_dataset)."
),
}
# update mutates the live instance in place; snapshot fitted state so
# a rejected update does not leave the handle un-fitted
import copy
snapshot = copy.deepcopy(instance)
try:
kwargs = update_params or {}
if X is not None:
instance.update(y, X=X, **kwargs)
else:
instance.update(y, **kwargs)
return {
"success": True,
"handle": handle_id,
"message": "Estimator updated successfully",
}
except Exception as e:
self._handle_manager.replace_instance(handle_id, snapshot)
return {"success": False, "error": str(e)}
[docs]
def get_fitted_params(self, handle_id: str) -> dict[str, Any]:
"""Get fitted parameters from an estimator."""
try:
instance = self._handle_manager.get_instance(handle_id)
except KeyError:
return {"success": False, "error": self._handle_manager.describe_missing(handle_id)}
if not self._handle_manager.is_fitted(handle_id):
return {"success": False, "error": "Estimator not fitted"}
try:
from sktime_mcp.server import sanitize_for_json
params = instance.get_fitted_params()
return {"success": True, "fitted_params": sanitize_for_json(params)}
except Exception as e:
return {"success": False, "error": str(e)}
[docs]
async def fit_async(
self,
handle_id: str,
X_dataset: str | None = None,
y_dataset: str | None = None,
X_handle: str | None = None,
y_handle: str | None = None,
fh: Any | None = None,
job_id: str | None = None,
) -> dict[str, Any]:
"""Async version of fit with job tracking."""
try:
import asyncio
from sktime_mcp.runtime.jobs import JobStatus
# Update status to RUNNING
self._job_manager.update_job(job_id, status=JobStatus.RUNNING)
# Step 1: Load data
self._job_manager.update_job(
job_id,
completed_steps=0,
current_step="Loading data...",
)
await asyncio.sleep(0.01)
X = None
y = None
if X_handle:
if X_handle not in self._data_handles:
raise ValueError(f"Unknown X data handle: {X_handle}")
X = self._data_handles[X_handle]["y"]
if y_handle:
if y_handle not in self._data_handles:
raise ValueError(f"Unknown y data handle: {y_handle}")
y = self._data_handles[y_handle]["y"]
if X_dataset and X_dataset == y_dataset:
data_res = self.load_dataset(X_dataset)
if not data_res["success"]:
raise ValueError(data_res["error"])
y = data_res["y"]
X = data_res["X"]
else:
if X_dataset:
data_res = self.load_dataset(X_dataset)
if not data_res["success"]:
raise ValueError(data_res["error"])
X = data_res["X"] if data_res["X"] is not None else data_res["y"]
if y_dataset:
data_res = self.load_dataset(y_dataset)
if not data_res["success"]:
raise ValueError(data_res["error"])
y = data_res["y"]
# Step 2: Fit model
self._job_manager.update_job(
job_id,
completed_steps=1,
current_step="Fitting model (this may take a while)...",
)
# Run fit in thread pool so it doesn't block async loop
loop = asyncio.get_running_loop()
import concurrent.futures
with concurrent.futures.ThreadPoolExecutor() as pool:
def run_fit():
return self.fit(handle_id, y, X=X, fh=fh)
fit_result = await loop.run_in_executor(pool, run_fit)
if not fit_result["success"]:
raise ValueError(fit_result["error"])
if X_dataset or y_dataset:
try:
handle_info = self._handle_manager.get_info(handle_id)
handle_info.metadata["training_dataset"] = y_dataset or X_dataset
except Exception:
pass
self._job_manager.update_job(
job_id,
status=JobStatus.COMPLETED,
completed_steps=2,
current_step="Training completed successfully.",
result={"success": True, "handle": handle_id, "fitted": True},
)
return {"success": True, "handle": handle_id}
except Exception as e:
from sktime_mcp.runtime.jobs import JobStatus
self._job_manager.update_job(
job_id,
status=JobStatus.FAILED,
current_step="Training failed.",
errors=[str(e)], # traceback logged server-side, not leaked to the client
)
return {"success": False, "error": str(e)}
[docs]
async def evaluate_async(
self,
handle_id: str,
y: str,
*,
X: str | None = None,
cv_folds: int = 3,
metric: str | None = None,
initial_window: int | None = None,
job_id: str | None = None,
) -> dict[str, Any]:
"""Async version of evaluate with job tracking."""
try:
self._job_manager.update_job(job_id, status=JobStatus.RUNNING)
# Step 1: Load data
self._job_manager.update_job(job_id, completed_steps=0, current_step="Loading data...")
await asyncio.sleep(0.01)
try:
instance = self._handle_manager.get_instance(handle_id)
except KeyError as err:
raise ValueError(self._handle_manager.describe_missing(handle_id)) from err
y_res = self._resolve_source(y)
if not y_res["success"]:
raise ValueError(y_res["error"])
_y = y_res["data"]
_X = None
if X:
x_res = self._resolve_source(X, prefer="X")
if not x_res["success"]:
raise ValueError(x_res["error"])
_X = x_res["data"]
scoring = None
if metric:
scoring = _resolve_metric_scoring(metric)
if scoring is None:
raise ValueError(
f"Unknown metric: {metric}. "
"Check available metrics with query_registry(task='metric')."
)
# Step 2: Run cross-validation
self._job_manager.update_job(
job_id, completed_steps=1, current_step="Running cross-validation..."
)
await asyncio.sleep(0.01)
loop = asyncio.get_running_loop()
fold_results, metrics, summary = await loop.run_in_executor(
None,
lambda: _run_evaluate(instance, _y, _X, cv_folds, scoring, initial_window),
)
# Step 3: Summarize results
self._job_manager.update_job(
job_id, completed_steps=2, current_step="Summarizing results..."
)
await asyncio.sleep(0.01)
result = {
"success": True,
"metrics": metrics,
"fold_results": fold_results,
"summary": summary,
"cv_folds_run": len(fold_results),
"cv_folds_requested": cv_folds,
}
self._job_manager.update_job(
job_id,
status=JobStatus.COMPLETED,
completed_steps=3,
current_step="Evaluation completed.",
result=result,
)
return result
except Exception as e:
self._job_manager.update_job(
job_id,
status=JobStatus.FAILED,
current_step="Evaluation failed.",
errors=[str(e)], # traceback logged server-side, not leaked to the client
)
return {"success": False, "error": str(e)}
# L-9: We can add more methods here to handle diverse use cases and their pipelines
[docs]
def list_datasets(self) -> list[str]:
"""List available demo datasets."""
return list(_get_demo_datasets().keys())
[docs]
def load_data_source(self, config: dict[str, Any]) -> dict[str, Any]:
"""
Load data from any source (pandas, SQL, file, etc.).
Args:
config: Data source configuration with 'type' key
Examples:
- {"type": "pandas", "data": df, "time_column": "date", "target_column": "value"}
- {"type": "sql", "connection_string": "...", "query": "...", "time_column": "date"}
- {"type": "file", "path": "/path/to/data.csv", "time_column": "date"}
Returns:
Dictionary with:
- success: bool
- data_handle: str (handle ID for the loaded data)
- metadata: dict (information about the data)
- validation: dict (validation results)
"""
try:
from sktime_mcp.data import DataSourceRegistry
# Create adapter
adapter = DataSourceRegistry.create_adapter(config)
# Load data
data = adapter.load()
# Validate
is_valid, validation_report = adapter.validate(data)
if not is_valid:
return {
"success": False,
"error": "Data validation failed",
"validation": validation_report,
}
# Convert to sktime format
y, X = adapter.to_sktime_format(data)
# Update metadata to reflect the target and used columns
metadata = adapter.get_metadata().copy()
validation_report = _merge_adapter_validation_warnings(validation_report, metadata)
metadata["columns"] = [y.name if hasattr(y, "name") and y.name else "target"]
if X is not None:
metadata["exog_columns"] = list(X.columns)
# Inject column dtypes so LLMs can distinguish time index vs target
metadata["dtypes"] = {col: str(dtype) for col, dtype in data.dtypes.items()}
# Generate handle
data_handle = f"data_{uuid.uuid4().hex[:8]}"
# Store (enforces max_data_handles limit)
self._register_data_handle(
data_handle,
{
"y": y,
"X": X,
"metadata": metadata,
"validation": validation_report,
"config": config,
},
)
# Apply auto-formatting if enabled
if getattr(self, "_auto_format_enabled", True):
try:
format_result = self.format_data_handle(
data_handle,
auto_infer_freq=True,
fill_missing=True,
remove_duplicates=True,
release_original=True,
)
if format_result["success"]:
return {
"success": True,
"data_handle": format_result["data_handle"],
"metadata": format_result["metadata"],
"validation": validation_report,
"formatted": True,
"changes_made": format_result["changes_made"],
}
except Exception as e:
logger.warning(f"Auto-formatting failed: {e}")
# Continue with unformatted data if formatting fails
# Auto-format disabled or failed: still normalise the stored handle to
# a PeriodIndex where possible so seasonal forecasters work (#531).
stored = self._data_handles.get(data_handle)
if stored is not None:
stored["y"] = _to_period_index_if_possible(stored["y"])
if stored.get("X") is not None:
stored["X"] = _to_period_index_if_possible(stored["X"])
_final_meta = adapter.get_metadata().copy()
_final_meta["dtypes"] = {col: str(dtype) for col, dtype in data.dtypes.items()}
return {
"success": True,
"data_handle": data_handle,
"metadata": _final_meta,
"validation": validation_report,
}
except Exception as e:
logger.exception("Error loading data source")
return {
"success": False,
"error": str(e),
"error_type": type(e).__name__,
}
[docs]
async def load_data_source_async(
self,
config: dict[str, Any],
job_id: str | None = None,
) -> dict[str, Any]:
"""
Async version of load_data_source with job tracking.
Runs data loading in the background without blocking the
MCP server. Progress is tracked via the JobManager.
Args:
config: Data source configuration
job_id: Optional job ID (created if not provided)
Returns:
Dictionary with data_handle and metadata
"""
source_type = config.get("type", "unknown")
if job_id is None:
job_id = self._job_manager.create_job(
job_type="data_loading",
estimator_handle="",
dataset_name=source_type,
total_steps=3,
)
try:
self._job_manager.update_job(job_id, status=JobStatus.RUNNING)
# Step 1: Load raw data
self._job_manager.update_job(
job_id, completed_steps=0, current_step=f"Loading data from '{source_type}'..."
)
await asyncio.sleep(0.01)
from sktime_mcp.data import DataSourceRegistry
loop = asyncio.get_event_loop()
adapter = DataSourceRegistry.create_adapter(config)
data = await loop.run_in_executor(None, adapter.load)
# Step 2: Validate
self._job_manager.update_job(
job_id, completed_steps=1, current_step="Validating data..."
)
await asyncio.sleep(0.01)
is_valid, validation_report = adapter.validate(data)
if not is_valid:
self._job_manager.update_job(
job_id, status=JobStatus.FAILED, errors=["Data validation failed"]
)
return {
"success": False,
"error": "Data validation failed",
"validation": validation_report,
}
# Step 3: Convert, store, and format
self._job_manager.update_job(
job_id, completed_steps=2, current_step="Converting to sktime format..."
)
await asyncio.sleep(0.01)
y, X = adapter.to_sktime_format(data)
metadata = adapter.get_metadata().copy()
validation_report = _merge_adapter_validation_warnings(validation_report, metadata)
metadata["columns"] = [y.name if hasattr(y, "name") and y.name else "target"]
if X is not None:
metadata["exog_columns"] = list(X.columns)
# Inject column dtypes so LLMs can distinguish time index vs target
metadata["dtypes"] = {col: str(dtype) for col, dtype in data.dtypes.items()}
data_handle = f"data_{uuid.uuid4().hex[:8]}"
self._register_data_handle(
data_handle,
{
"y": y,
"X": X,
"metadata": metadata,
"validation": validation_report,
"config": config,
},
)
# auto-format if enabled
if getattr(self, "_auto_format_enabled", True):
try:
format_result = self.format_data_handle(
data_handle,
auto_infer_freq=True,
fill_missing=True,
remove_duplicates=True,
release_original=True,
)
if format_result["success"]:
data_handle = format_result["data_handle"]
metadata = format_result["metadata"]
except Exception as e:
logger.warning(f"Auto-formatting failed: {e}")
result = {
"success": True,
"data_handle": data_handle,
"metadata": metadata,
"validation": validation_report,
}
# mark completed with the data_handle in the result
self._job_manager.update_job(
job_id,
status=JobStatus.COMPLETED,
completed_steps=3,
current_step="Completed",
result=result,
)
return result
except Exception as e:
logger.exception(f"Error in async data loading for job {job_id}")
self._job_manager.update_job(job_id, status=JobStatus.FAILED, errors=[str(e)])
return {
"success": False,
"error": str(e),
"job_id": job_id,
}
[docs]
def list_data_handles(self) -> dict[str, Any]:
"""
List all loaded data handles.
Returns:
Dictionary with list of data handles and their metadata
"""
handles = []
for handle_id, data_info in self._data_handles.items():
handles.append(
{
"handle": handle_id,
"metadata": data_info["metadata"],
"validation": data_info["validation"],
}
)
return {
"success": True,
"count": len(handles),
"handles": handles,
}
[docs]
def release_data_handle(self, data_handle: str) -> dict[str, Any]:
"""
Release a data handle and free memory.
Args:
data_handle: Data handle to release
Returns:
Dictionary with success status
"""
if data_handle in self._data_handles:
del self._data_handles[data_handle]
return {
"success": True,
"message": f"Data handle '{data_handle}' released",
}
else:
return {
"success": False,
"error": self.data_handle_missing(data_handle)["error"],
}
_executor_instance: Executor | None = None
[docs]
def get_executor() -> Executor:
global _executor_instance
if _executor_instance is None:
_executor_instance = Executor()
return _executor_instance