QuickLab-MCP 0.1.0__py3-none-any.whl

ml_mcp/mcp/prompts.py

@@ -0,0 +1,38 @@
+"""MCP Prompts providing guided workflows for agentic experimentation."""
+
+
+def experiment_design_prompt(dataset_description: str, task_type: str) -> str:
+    """Prompt guiding agents to design high-yield ML experiments."""
+    return f"""You are the ModelLab ML Experimentation Architect.
+Task: Design a rigorous, leak-free experimentation plan.
+
+Dataset Description:
+{dataset_description}
+
+Task Type:
+{task_type}
+
+Guidelines:
+1. Select 2-3 candidate architectures from approved tabular models: Logistic Regression, Random Forest, XGBoost, LightGBM, CatBoost, Linear SVM, MLP.
+2. Recommend appropriate split strategy (Stratified K-Fold for imbalanced classification, Time Series Split for temporal data).
+3. Specify primary and secondary evaluation metrics (e.g. PR AUC for severe class imbalance, ROC AUC for balanced ranking).
+4. Propose non-overlapping hyperparameter variations to test.
+
+Use `list_models` and `get_model` to inspect hyperparameter bounds, then submit via `create_experiment`.
+"""
+
+
+def error_diagnosis_prompt(experiment_id: str, problem_summary: str) -> str:
+    """Prompt assisting agents in diagnosing training regressions or anomalies."""
+    return f"""You are the ModelLab Diagnostic Auditor.
+Task: Inspect model performance anomalies for Experiment: {experiment_id}
+
+Problem Summary:
+{problem_summary}
+
+Step-by-step diagnostic workflow:
+1. Call `get_experiment_metrics` for both train and validation splits to evaluate the generalization gap.
+2. Call `analyze_experiment` to inspect automated overfitting and leakage signals.
+3. Call `analyze_model_errors` to review false positives vs false negatives and residual skew.
+4. Formulate 2 actionable next experiments addressing the root cause.
+"""

ml_mcp/mcp/resources.py

@@ -0,0 +1,35 @@
+"""MCP Resources exposing system metadata and catalog information."""
+
+import json
+
+from ml_mcp.application.models.service import ModelService
+from ml_mcp.infrastructure.postgres.session import get_db_manager
+
+
+async def read_models_resource() -> str:
+    """Read full catalog of approved models (modellab://models)."""
+    async with get_db_manager().session() as sess:
+        service = ModelService(sess)
+        models = await service.list_models()
+        return json.dumps(models, indent=2)
+
+
+async def read_model_details_resource(model_id: str) -> str:
+    """Read detailed specification of a model (modellab://models/{model_id})."""
+    async with get_db_manager().session() as sess:
+        service = ModelService(sess)
+        details = await service.get_model(model_id)
+        return json.dumps(details, indent=2)
+
+
+async def read_experiment_metrics_resource(
+    experiment_id: str, tenant_id: str = "default-tenant"
+) -> str:
+    """Read metrics for an experiment (modellab://experiments/{experiment_id}/metrics)."""
+    async with get_db_manager().session() as sess:
+        from ml_mcp.infrastructure.postgres.repositories.metrics import MetricRepository
+
+        repo = MetricRepository(sess)
+        metrics = await repo.get_metrics(tenant_id, experiment_id)
+        payload = {m.metric_name: m.metric_value for m in metrics}
+        return json.dumps(payload, indent=2)

ml_mcp/mcp/schemas/__init__.py

@@ -0,0 +1 @@
+"""MCP tool and resource schemas."""

ml_mcp/mcp/schemas/tools.py

@@ -0,0 +1,124 @@
+"""Strict Pydantic schemas for MCP tool arguments and responses."""
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+# 1. Discovery Schemas
+class ListModelsInput(BaseModel):
+    pass
+
+
+class GetModelInput(BaseModel):
+    model_id: str = Field(..., description="Unique model identifier, e.g. 'random_forest'")
+
+
+class ListModelVersionsInput(BaseModel):
+    model_id: str = Field(..., description="Unique model identifier")
+
+
+class ListDatasetsInput(BaseModel):
+    project_id: str | None = Field(default=None, description="Optional project filter")
+
+
+class GetDatasetInput(BaseModel):
+    dataset_id: str = Field(..., description="Unique dataset identifier")
+
+
+class ListDatasetVersionsInput(BaseModel):
+    dataset_id: str = Field(..., description="Unique dataset identifier")
+
+
+# 2. Dataset Management Schemas
+class RegisterDatasetInput(BaseModel):
+    project_id: str = Field(..., description="Project ID to register dataset under")
+    name: str = Field(..., description="Human-readable dataset name")
+    description: str = Field(default="", description="Detailed description")
+    format: str = Field(..., description="Format: 'csv', 'parquet', or 'jsonl'")
+    data_base64: str = Field(..., description="Base64-encoded dataset file content")
+    version: str = Field(default="1.0", description="Semantic dataset version")
+
+
+class ValidateDatasetInput(BaseModel):
+    format: str = Field(..., description="Format: 'csv', 'parquet', or 'jsonl'")
+    data_base64: str = Field(..., description="Base64-encoded dataset file content")
+
+
+class InspectDatasetInput(BaseModel):
+    dataset_id: str = Field(..., description="Dataset ID to inspect")
+    version: str = Field(default="1.0", description="Dataset version")
+    sample_rows: int = Field(
+        default=5, ge=1, le=50, description="Number of preview records to return"
+    )
+
+
+# 3. Experiment Schemas
+class CreateExperimentInput(BaseModel):
+    project_id: str = Field(..., description="Target project ID")
+    dataset_version_id: str = Field(..., description="Registered DatasetVersion ID")
+    model_version_id: str = Field(..., description="Registered ModelVersion ID")
+    task_type: str = Field(
+        ..., description="'binary_classification', 'multiclass_classification', or 'regression'"
+    )
+    target_column: str = Field(..., description="Name of column to predict")
+    feature_columns: list[str] | None = Field(
+        default=None, description="Subset of features to train on"
+    )
+    primary_metric: str = Field(default="roc_auc", description="Primary optimization metric")
+    additional_metrics: list[str] = Field(default_factory=list)
+    hyperparameters: dict[str, Any] = Field(default_factory=dict)
+    random_seed: int = Field(default=42)
+    idempotency_key: str | None = Field(default=None)
+
+
+class GetExperimentInput(BaseModel):
+    experiment_id: str = Field(..., description="Experiment ID to fetch")
+
+
+class CancelExperimentInput(BaseModel):
+    experiment_id: str = Field(..., description="Experiment ID to cancel")
+
+
+class ListExperimentsInput(BaseModel):
+    project_id: str | None = None
+    status: str | None = None
+    limit: int = Field(default=50, ge=1, le=100)
+    offset: int = Field(default=0, ge=0)
+
+
+# 4. Results Schemas
+class GetExperimentMetricsInput(BaseModel):
+    experiment_id: str = Field(..., description="Target experiment ID")
+    split: str | None = Field(default=None, description="'train' or 'validation'")
+
+
+class GetExperimentPredictionsInput(BaseModel):
+    experiment_id: str = Field(..., description="Target experiment ID")
+
+
+class ListExperimentArtifactsInput(BaseModel):
+    experiment_id: str = Field(..., description="Target experiment ID")
+
+
+class ReadExperimentArtifactInput(BaseModel):
+    artifact_id: str = Field(..., description="Target artifact ID")
+
+
+class CompareExperimentsInput(BaseModel):
+    experiment_ids: list[str] = Field(
+        ..., min_length=2, max_length=10, description="List of experiment IDs to compare"
+    )
+
+
+# 5. Analysis Schemas
+class AnalyzeExperimentInput(BaseModel):
+    experiment_id: str = Field(..., description="Experiment ID to analyze")
+
+
+class AnalyzeModelErrorsInput(BaseModel):
+    experiment_id: str = Field(..., description="Experiment ID to inspect errors for")
+
+
+class CheckExperimentValidityInput(BaseModel):
+    experiment_id: str = Field(..., description="Experiment ID to validate")

ml_mcp/mcp/tools/__init__.py

@@ -0,0 +1,63 @@
+"""MCP tools module."""
+
+from ml_mcp.mcp.tools.analysis import (
+    handle_analyze_experiment,
+    handle_analyze_model_errors,
+    handle_check_experiment_validity,
+)
+from ml_mcp.mcp.tools.datasets import (
+    handle_inspect_dataset,
+    handle_register_dataset,
+    handle_validate_dataset,
+)
+from ml_mcp.mcp.tools.discovery import (
+    handle_get_dataset,
+    handle_get_model,
+    handle_list_dataset_versions,
+    handle_list_datasets,
+    handle_list_model_versions,
+    handle_list_models,
+)
+from ml_mcp.mcp.tools.experiments import (
+    handle_cancel_experiment,
+    handle_create_experiment,
+    handle_get_experiment,
+    handle_list_experiments,
+)
+from ml_mcp.mcp.tools.projects import (
+    handle_create_project,
+    handle_list_projects,
+)
+from ml_mcp.mcp.tools.results import (
+    handle_compare_experiments,
+    handle_get_experiment_metrics,
+    handle_get_experiment_predictions,
+    handle_list_experiment_artifacts,
+    handle_read_experiment_artifact,
+)
+
+__all__ = [
+    "handle_create_project",
+    "handle_list_projects",
+    "handle_list_models",
+    "handle_get_model",
+    "handle_list_model_versions",
+    "handle_list_datasets",
+    "handle_get_dataset",
+    "handle_list_dataset_versions",
+    "handle_register_dataset",
+    "handle_validate_dataset",
+    "handle_inspect_dataset",
+    "handle_create_experiment",
+    "handle_get_experiment",
+    "handle_cancel_experiment",
+    "handle_list_experiments",
+    "handle_get_experiment_metrics",
+    "handle_get_experiment_predictions",
+    "handle_list_experiment_artifacts",
+    "handle_read_experiment_artifact",
+    "handle_compare_experiments",
+    "handle_analyze_experiment",
+    "handle_analyze_model_errors",
+    "handle_check_experiment_validity",
+]

ml_mcp/mcp/tools/analysis.py

@@ -0,0 +1,37 @@
+"""Analysis MCP tool handlers."""
+
+from typing import Any
+
+from ml_mcp.application.analysis.service import AnalysisService
+from ml_mcp.domain.policies import Principal, Scope
+from ml_mcp.infrastructure.postgres.session import get_db_manager
+
+
+async def handle_analyze_experiment(
+    experiment_id: str,
+    principal: Principal,
+) -> dict[str, Any]:
+    principal.enforce_permission(Scope.ANALYSIS_CREATE)
+    async with get_db_manager().session() as sess:
+        service = AnalysisService(sess)
+        return await service.analyze_experiment(principal.tenant_id, experiment_id)
+
+
+async def handle_analyze_model_errors(
+    experiment_id: str,
+    principal: Principal,
+) -> dict[str, Any]:
+    principal.enforce_permission(Scope.ANALYSIS_CREATE)
+    async with get_db_manager().session() as sess:
+        service = AnalysisService(sess)
+        return await service.analyze_model_errors(principal.tenant_id, experiment_id)
+
+
+async def handle_check_experiment_validity(
+    experiment_id: str,
+    principal: Principal,
+) -> dict[str, Any]:
+    principal.enforce_permission(Scope.ANALYSIS_CREATE)
+    async with get_db_manager().session() as sess:
+        service = AnalysisService(sess)
+        return await service.check_experiment_validity(principal.tenant_id, experiment_id)

ml_mcp/mcp/tools/datasets.py

@@ -0,0 +1,82 @@
+import base64
+import binascii
+from typing import Any
+
+from ml_mcp.application.datasets.service import DatasetService
+from ml_mcp.domain.errors import InvalidInputError
+from ml_mcp.domain.policies import Principal, Scope, validate_tenant_access
+from ml_mcp.infrastructure.postgres.session import get_db_manager
+
+
+def _safe_b64decode(data_base64: str) -> bytes:
+    """Safely decode base64 payload with strict validation."""
+    clean_b64 = data_base64.strip()
+    if not clean_b64:
+        raise InvalidInputError("Dataset payload is empty.")
+    try:
+        return base64.b64decode(clean_b64, validate=True)
+    except (binascii.Error, ValueError) as exc:
+        raise InvalidInputError(f"Invalid base64 payload: {exc}") from exc
+
+
+async def handle_register_dataset(
+    project_id: str,
+    name: str,
+    description: str,
+    format: str,
+    data_base64: str,
+    principal: Principal,
+    version: str = "1.0",
+) -> dict[str, Any]:
+    principal.enforce_permission(Scope.DATASETS_WRITE)
+    validate_tenant_access(principal, principal.tenant_id, project_id)
+
+    raw_bytes = _safe_b64decode(data_base64)
+    async with get_db_manager().session() as sess:
+        from ml_mcp.domain.errors import ResourceNotFoundError
+        from ml_mcp.infrastructure.postgres.repositories.projects import ProjectRepository
+
+        proj_repo = ProjectRepository(sess)
+        proj = await proj_repo.get_project(principal.tenant_id, project_id)
+        if not proj:
+            raise ResourceNotFoundError("Project", project_id)
+
+        service = DatasetService(sess)
+        return await service.register_dataset(
+            tenant_id=principal.tenant_id,
+            project_id=project_id,
+            name=name,
+            description=description,
+            data_format=format,
+            raw_bytes=raw_bytes,
+            version=version,
+        )
+
+
+async def handle_validate_dataset(
+    format: str,
+    data_base64: str,
+    principal: Principal,
+) -> dict[str, Any]:
+    principal.enforce_permission(Scope.DATASETS_READ)
+    raw_bytes = _safe_b64decode(data_base64)
+    async with get_db_manager().session() as sess:
+        service = DatasetService(sess)
+        return await service.validate_only(raw_bytes, format)
+
+
+async def handle_inspect_dataset(
+    dataset_id: str,
+    principal: Principal,
+    version: str = "1.0",
+    sample_rows: int = 5,
+) -> dict[str, Any]:
+    principal.enforce_permission(Scope.DATASETS_READ)
+    async with get_db_manager().session() as sess:
+        service = DatasetService(sess)
+        return await service.inspect_dataset(
+            tenant_id=principal.tenant_id,
+            dataset_id=dataset_id,
+            version=version,
+            sample_rows=sample_rows,
+        )

ml_mcp/mcp/tools/discovery.py

@@ -0,0 +1,82 @@
+"""Discovery MCP tools for models and datasets."""
+
+from typing import Any
+
+from ml_mcp.application.datasets.service import DatasetService
+from ml_mcp.application.models.service import ModelService
+from ml_mcp.domain.policies import Principal, Scope
+from ml_mcp.infrastructure.postgres.session import get_db_manager
+
+
+async def handle_list_models(principal: Principal) -> list[dict[str, Any]]:
+    principal.enforce_permission(Scope.MODELS_READ)
+    async with get_db_manager().session() as sess:
+        service = ModelService(sess)
+        return await service.list_models()
+
+
+async def handle_get_model(model_id: str, principal: Principal) -> dict[str, Any]:
+    principal.enforce_permission(Scope.MODELS_READ)
+    async with get_db_manager().session() as sess:
+        service = ModelService(sess)
+        return await service.get_model(model_id)
+
+
+async def handle_list_model_versions(model_id: str, principal: Principal) -> list[dict[str, Any]]:
+    principal.enforce_permission(Scope.MODELS_READ)
+    async with get_db_manager().session() as sess:
+        service = ModelService(sess)
+        return await service.list_model_versions(model_id)
+
+
+async def handle_list_datasets(
+    principal: Principal, project_id: str | None = None
+) -> list[dict[str, Any]]:
+    principal.enforce_permission(Scope.DATASETS_READ)
+    if project_id:
+        from ml_mcp.domain.policies import validate_tenant_access
+
+        validate_tenant_access(principal, principal.tenant_id, project_id)
+
+    async with get_db_manager().session() as sess:
+        if project_id:
+            from ml_mcp.domain.errors import ResourceNotFoundError
+            from ml_mcp.infrastructure.postgres.repositories.projects import ProjectRepository
+
+            proj_repo = ProjectRepository(sess)
+            proj = await proj_repo.get_project(principal.tenant_id, project_id)
+            if not proj:
+                raise ResourceNotFoundError("Project", project_id)
+
+        service = DatasetService(sess)
+        return await service.list_datasets(principal.tenant_id, project_id)
+
+
+async def handle_get_dataset(dataset_id: str, principal: Principal) -> dict[str, Any]:
+    principal.enforce_permission(Scope.DATASETS_READ)
+    async with get_db_manager().session() as sess:
+        service = DatasetService(sess)
+        return await service.get_dataset(principal.tenant_id, dataset_id)
+
+
+async def handle_list_dataset_versions(
+    dataset_id: str, principal: Principal
+) -> list[dict[str, Any]]:
+    principal.enforce_permission(Scope.DATASETS_READ)
+    async with get_db_manager().session() as sess:
+        from ml_mcp.infrastructure.postgres.repositories.datasets import DatasetRepository
+
+        repo = DatasetRepository(sess)
+        versions = await repo.list_versions(principal.tenant_id, dataset_id)
+        return [
+            {
+                "version_id": v.id,
+                "version": v.version,
+                "row_count": v.row_count,
+                "column_count": v.column_count,
+                "content_hash": v.content_hash,
+                "validation_status": v.validation_status,
+                "created_at": v.created_at.isoformat(),
+            }
+            for v in versions
+        ]

ml_mcp/mcp/tools/experiments.py

@@ -0,0 +1,96 @@
+"""Experiment MCP tool handlers."""
+
+from typing import Any
+
+from ml_mcp.application.experiments.service import ExperimentService
+from ml_mcp.domain.policies import Principal, Scope, validate_tenant_access
+from ml_mcp.infrastructure.postgres.session import get_db_manager
+
+
+async def handle_create_experiment(
+    project_id: str,
+    dataset_version_id: str,
+    model_version_id: str,
+    task_type: str,
+    target_column: str,
+    principal: Principal,
+    feature_columns: list[str] | None = None,
+    primary_metric: str = "roc_auc",
+    additional_metrics: list[str] | None = None,
+    hyperparameters: dict[str, Any] | None = None,
+    random_seed: int = 42,
+    idempotency_key: str | None = None,
+) -> dict[str, Any]:
+    principal.enforce_permission(Scope.EXPERIMENTS_CREATE)
+    validate_tenant_access(principal, principal.tenant_id, project_id)
+
+    async with get_db_manager().session() as sess:
+        from ml_mcp.domain.errors import ResourceNotFoundError
+        from ml_mcp.infrastructure.postgres.repositories.projects import ProjectRepository
+
+        proj_repo = ProjectRepository(sess)
+        proj = await proj_repo.get_project(principal.tenant_id, project_id)
+        if not proj:
+            raise ResourceNotFoundError("Project", project_id)
+
+        service = ExperimentService(sess)
+        return await service.create_experiment(
+            tenant_id=principal.tenant_id,
+            project_id=project_id,
+            dataset_version_id=dataset_version_id,
+            model_version_id=model_version_id,
+            task_type=task_type,
+            target_column=target_column,
+            principal_id=principal.principal_id,
+            feature_columns=feature_columns,
+            primary_metric=primary_metric,
+            additional_metrics=additional_metrics,
+            hyperparameters=hyperparameters,
+            random_seed=random_seed,
+            idempotency_key=idempotency_key,
+        )
+
+
+async def handle_get_experiment(experiment_id: str, principal: Principal) -> dict[str, Any]:
+    principal.enforce_permission(Scope.EXPERIMENTS_READ)
+    async with get_db_manager().session() as sess:
+        service = ExperimentService(sess)
+        return await service.get_experiment(principal.tenant_id, experiment_id)
+
+
+async def handle_cancel_experiment(experiment_id: str, principal: Principal) -> dict[str, Any]:
+    principal.enforce_permission(Scope.EXPERIMENTS_CANCEL)
+    async with get_db_manager().session() as sess:
+        service = ExperimentService(sess)
+        return await service.cancel_experiment(principal.tenant_id, experiment_id)
+
+
+async def handle_list_experiments(
+    principal: Principal,
+    project_id: str | None = None,
+    status: str | None = None,
+    limit: int = 50,
+    offset: int = 0,
+) -> list[dict[str, Any]]:
+    principal.enforce_permission(Scope.EXPERIMENTS_READ)
+    if project_id:
+        validate_tenant_access(principal, principal.tenant_id, project_id)
+
+    async with get_db_manager().session() as sess:
+        if project_id:
+            from ml_mcp.domain.errors import ResourceNotFoundError
+            from ml_mcp.infrastructure.postgres.repositories.projects import ProjectRepository
+
+            proj_repo = ProjectRepository(sess)
+            proj = await proj_repo.get_project(principal.tenant_id, project_id)
+            if not proj:
+                raise ResourceNotFoundError("Project", project_id)
+
+        service = ExperimentService(sess)
+        return await service.list_experiments(
+            tenant_id=principal.tenant_id,
+            project_id=project_id,
+            status=status,
+            limit=limit,
+            offset=offset,
+        )

ml_mcp/mcp/tools/projects.py

@@ -0,0 +1,39 @@
+"""Project MCP tool handlers."""
+
+from typing import Any
+
+from ml_mcp.application.projects.service import ProjectService
+from ml_mcp.domain.policies import Principal, Scope
+from ml_mcp.infrastructure.postgres.session import get_db_manager
+
+
+async def handle_create_project(
+    name: str,
+    principal: Principal,
+    description: str = "",
+) -> dict[str, Any]:
+    """Handle create_project tool call."""
+    principal.enforce_permission(Scope.PROJECTS_WRITE)
+    async with get_db_manager().session() as sess:
+        service = ProjectService(sess)
+        return await service.create_project(
+            tenant_id=principal.tenant_id,
+            name=name,
+            description=description,
+        )
+
+
+async def handle_list_projects(
+    principal: Principal,
+    limit: int = 50,
+    offset: int = 0,
+) -> list[dict[str, Any]]:
+    """Handle list_projects tool call."""
+    principal.enforce_permission(Scope.PROJECTS_READ)
+    async with get_db_manager().session() as sess:
+        service = ProjectService(sess)
+        return await service.list_projects(
+            tenant_id=principal.tenant_id,
+            limit=limit,
+            offset=offset,
+        )

ml_mcp/mcp/tools/results.py

@@ -0,0 +1,59 @@
+"""Result and artifact MCP tool handlers."""
+
+from typing import Any
+
+from ml_mcp.application.artifacts.service import ArtifactService
+from ml_mcp.domain.policies import Principal, Scope
+from ml_mcp.infrastructure.postgres.session import get_db_manager
+
+
+async def handle_get_experiment_metrics(
+    experiment_id: str,
+    principal: Principal,
+    split: str | None = None,
+) -> dict[str, Any]:
+    principal.enforce_permission(Scope.EXPERIMENTS_READ)
+    async with get_db_manager().session() as sess:
+        service = ArtifactService(sess)
+        return await service.get_metrics(principal.tenant_id, experiment_id, split)
+
+
+async def handle_get_experiment_predictions(
+    experiment_id: str,
+    principal: Principal,
+    limit: int = 5,
+) -> dict[str, Any]:
+    principal.enforce_permission(Scope.ARTIFACTS_READ)
+    async with get_db_manager().session() as sess:
+        service = ArtifactService(sess)
+        return await service.get_predictions(principal.tenant_id, experiment_id, limit)
+
+
+async def handle_list_experiment_artifacts(
+    experiment_id: str,
+    principal: Principal,
+) -> list[dict[str, Any]]:
+    principal.enforce_permission(Scope.ARTIFACTS_READ)
+    async with get_db_manager().session() as sess:
+        service = ArtifactService(sess)
+        return await service.list_artifacts(principal.tenant_id, experiment_id)
+
+
+async def handle_read_experiment_artifact(
+    artifact_id: str,
+    principal: Principal,
+) -> dict[str, Any]:
+    principal.enforce_permission(Scope.ARTIFACTS_READ)
+    async with get_db_manager().session() as sess:
+        service = ArtifactService(sess)
+        return await service.get_artifact(principal.tenant_id, artifact_id)
+
+
+async def handle_compare_experiments(
+    experiment_ids: list[str],
+    principal: Principal,
+) -> dict[str, Any]:
+    principal.enforce_permission(Scope.EXPERIMENTS_READ)
+    async with get_db_manager().session() as sess:
+        service = ArtifactService(sess)
+        return await service.compare_experiments(principal.tenant_id, experiment_ids)

ml_mcp/server/__init__.py

@@ -0,0 +1 @@
+"""Server package for ML-MCP."""