dataiku_mcp/tools/machine_learning/analyses.py

# Copyright 2026 Dataiku SAS
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Inspection tools for generic ML analyses and trained models."""

from __future__ import annotations

from typing import Annotated, Any

from fastmcp import Context
from pydantic import Field

from ...instance_policy import PROJECT_READ_POLICY
from ...server import mcp
from ...auth import get_dss_client
from ...executors import run_blocking
from ..utils.serialization import columnar, compact_json, omit_empty
from ..utils.validation import require_non_empty_string as _require_non_empty_string
from .shared.common import find_analysis_input_dataset, require_single_ml_task
from .shared.summaries import build_ml_task_summary


def get_single_task_bundle(project_key: str, analysis_id: str):
    project = get_dss_client().get_project(project_key)
    analysis = project.get_analysis(analysis_id)
    task_ref = require_single_ml_task(analysis)
    mltask_id = task_ref["mlTaskId"]
    mltask = analysis.get_ml_task(mltask_id)
    return project, analysis, mltask_id, mltask


def slim_model_data(data: Any) -> Any:
    if isinstance(data, dict):
        return {
            key: slim_model_data(value)
            for key, value in data.items()
            if not key.endswith("PerFeature")
        }
    if isinstance(data, list):
        return [slim_model_data(item) for item in data]
    return data


def slim_mltask_settings(raw_settings: dict) -> dict:
    slimmed = dict(raw_settings)
    modeling = slimmed.get("modeling")
    if isinstance(modeling, dict):
        slimmed["modeling"] = {
            key: value
            for key, value in modeling.items()
            if not (
                isinstance(value, dict) and "enabled" in value and not value["enabled"]
            )
        }
    return slimmed


@mcp.tool(
    meta=PROJECT_READ_POLICY,
    title="List ML Analyses",
    annotations={
        "readOnlyHint": True,
        "destructiveHint": False,
        "openWorldHint": False,
    },
)
async def list_ml_analyses(
    project_key: str,
    ctx: Context,
    input_dataset: Annotated[
        str | None, Field(description="Keeps only analyses trained on this dataset.")
    ] = None,
) -> str:
    """Find a project's ML analyses, what they predict, and from which dataset."""
    project_key = _require_non_empty_string(project_key, "project_key")
    if input_dataset is not None:
        input_dataset = _require_non_empty_string(input_dataset, "input_dataset")
    await ctx.info(f"Listing ML analyses in {project_key}...")

    def _list() -> dict[str, Any]:
        project = get_dss_client().get_project(project_key)
        analyses = []
        for analysis_item in project.list_analyses():
            if input_dataset and analysis_item.get("inputDataset") != input_dataset:
                continue
            analysis_id = analysis_item["analysisId"]
            analysis = project.get_analysis(analysis_id)
            analysis_definition = analysis.get_definition().get_raw()
            task_ref = require_single_ml_task(analysis)
            mltask_id = task_ref["mlTaskId"]
            mltask = analysis.get_ml_task(mltask_id)
            settings = mltask.get_settings()
            analyses.append(
                build_ml_task_summary(
                    project_key=project_key,
                    analysis_id=analysis_id,
                    analysis_name=analysis_definition.get("name"),
                    input_dataset=analysis_item.get("inputDataset"),
                    mltask_id=mltask_id,
                    mltask=mltask,
                    settings=settings,
                    include_feature_details=False,
                    include_trained_models=False,
                )
            )
        return {
            "analyses": columnar(
                analyses,
                [
                    "analysis_id",
                    "analysis_name",
                    "input_dataset",
                    "mltask_id",
                    "task_type",
                    "prediction_type",
                    "target_column",
                    "backend_type",
                    "metric",
                    "metric_details",
                    "enabled_algorithms",
                    "guessing",
                    "training",
                    "trained_model_count",
                    "split",
                    "hyperparameter_search",
                    "weighting",
                    "diagnostics",
                    "feature_role_counts",
                    "features",
                ],
            )
        }

    return compact_json(await run_blocking(_list))


@mcp.tool(
    meta=PROJECT_READ_POLICY,
    title="Get ML Analysis Summary",
    annotations={
        "readOnlyHint": True,
        "destructiveHint": False,
        "openWorldHint": False,
    },
)
async def get_ml_analysis_summary(
    project_key: str,
    analysis_id: str,
    ctx: Context,
) -> str:
    """Read an ML analysis's target, task type, and training setup at a glance."""
    project_key = _require_non_empty_string(project_key, "project_key")
    analysis_id = _require_non_empty_string(analysis_id, "analysis_id")
    await ctx.info(f"Loading ML analysis summary for {analysis_id}...")

    def _get() -> dict[str, Any]:
        project, analysis, mltask_id, mltask = get_single_task_bundle(
            project_key, analysis_id
        )
        analysis_definition = analysis.get_definition().get_raw()
        settings = mltask.get_settings()
        return build_ml_task_summary(
            project_key=project_key,
            analysis_id=analysis_id,
            analysis_name=analysis_definition.get("name"),
            input_dataset=find_analysis_input_dataset(project, analysis_id),
            mltask_id=mltask_id,
            mltask=mltask,
            settings=settings,
            include_feature_details=True,
            include_trained_models=False,
        )

    return compact_json(await run_blocking(_get))


@mcp.tool(
    meta=PROJECT_READ_POLICY,
    title="Get ML Analysis Settings",
    annotations={
        "readOnlyHint": True,
        "destructiveHint": False,
        "openWorldHint": False,
    },
)
async def get_ml_analysis_settings(
    project_key: str,
    analysis_id: str,
    ctx: Context,
) -> str:
    """Read an ML analysis's raw feature handling, algorithms, and split settings."""
    project_key = _require_non_empty_string(project_key, "project_key")
    analysis_id = _require_non_empty_string(analysis_id, "analysis_id")
    await ctx.info(f"Loading ML analysis settings for {analysis_id}...")

    def _get() -> dict[str, Any]:
        project, analysis, mltask_id, mltask = get_single_task_bundle(
            project_key, analysis_id
        )
        analysis_definition = analysis.get_definition().get_raw()
        settings = mltask.get_settings()
        result = {
            "analysis_name": analysis_definition.get("name"),
            "input_dataset": find_analysis_input_dataset(project, analysis_id),
            "mltask_id": mltask_id,
            "mltask_settings": slim_mltask_settings(settings.get_raw()),
        }
        return omit_empty(result)

    return compact_json(await run_blocking(_get))


@mcp.tool(
    meta=PROJECT_READ_POLICY,
    title="List ML Analysis Models",
    annotations={
        "readOnlyHint": True,
        "destructiveHint": False,
        "openWorldHint": False,
    },
)
async def list_ml_analysis_models(
    project_key: str,
    analysis_id: str,
    ctx: Context,
) -> str:
    """Compare the models trained in an analysis and their performance."""
    project_key = _require_non_empty_string(project_key, "project_key")
    analysis_id = _require_non_empty_string(analysis_id, "analysis_id")
    await ctx.info(f"Listing trained models for ML analysis {analysis_id}...")

    def _list() -> dict[str, Any]:
        _, _, mltask_id, mltask = get_single_task_bundle(project_key, analysis_id)
        trained_model_ids = mltask.get_trained_models_ids()
        snippets = (
            mltask.get_trained_model_snippet(ids=trained_model_ids)
            if trained_model_ids
            else {}
        )
        models = []
        for trained_model_id in trained_model_ids:
            snippet = snippets.get(trained_model_id, {})
            models.append(
                {
                    "trained_model_id": trained_model_id,
                    "snippet": slim_model_data(snippet),
                }
            )
        return {
            "mltask_id": mltask_id,
            "models": columnar(models, ["trained_model_id", "snippet"]),
        }

    return compact_json(await run_blocking(_list))


@mcp.tool(
    meta=PROJECT_READ_POLICY,
    title="Get ML Model Details",
    annotations={
        "readOnlyHint": True,
        "destructiveHint": False,
        "openWorldHint": False,
    },
)
async def get_ml_model_details(
    project_key: str,
    analysis_id: str,
    trained_model_id: str,
    ctx: Context,
) -> str:
    """Read one trained model's metrics, features, and hyperparameters."""
    project_key = _require_non_empty_string(project_key, "project_key")
    analysis_id = _require_non_empty_string(analysis_id, "analysis_id")
    trained_model_id = _require_non_empty_string(trained_model_id, "trained_model_id")
    await ctx.info(
        f"Loading trained model {trained_model_id} from analysis {analysis_id}..."
    )

    def _get() -> dict[str, Any]:
        _, _, mltask_id, mltask = get_single_task_bundle(project_key, analysis_id)
        details = mltask.get_trained_model_details(trained_model_id)
        return {
            "mltask_id": mltask_id,
            "full_model_id": details.full_id,
            "details_class": details.__class__.__name__,
            "snippet": details.get_raw_snippet(),
        }

    return compact_json(await run_blocking(_get))