QuickLab-MCP 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (92) hide show
  1. ml_mcp/__init__.py +3 -0
  2. ml_mcp/application/__init__.py +1 -0
  3. ml_mcp/application/analysis/__init__.py +15 -0
  4. ml_mcp/application/analysis/diagnostics.py +122 -0
  5. ml_mcp/application/analysis/service.py +155 -0
  6. ml_mcp/application/artifacts/__init__.py +5 -0
  7. ml_mcp/application/artifacts/service.py +149 -0
  8. ml_mcp/application/datasets/__init__.py +6 -0
  9. ml_mcp/application/datasets/service.py +178 -0
  10. ml_mcp/application/datasets/validator.py +156 -0
  11. ml_mcp/application/experiments/__init__.py +5 -0
  12. ml_mcp/application/experiments/service.py +254 -0
  13. ml_mcp/application/models/__init__.py +6 -0
  14. ml_mcp/application/models/catalog_seed.py +131 -0
  15. ml_mcp/application/models/service.py +177 -0
  16. ml_mcp/application/projects/__init__.py +5 -0
  17. ml_mcp/application/projects/service.py +93 -0
  18. ml_mcp/config/__init__.py +5 -0
  19. ml_mcp/config/settings.py +176 -0
  20. ml_mcp/domain/__init__.py +1 -0
  21. ml_mcp/domain/errors/__init__.py +39 -0
  22. ml_mcp/domain/errors/base.py +155 -0
  23. ml_mcp/domain/errors/codes.py +23 -0
  24. ml_mcp/domain/policies/__init__.py +16 -0
  25. ml_mcp/domain/policies/rbac.py +126 -0
  26. ml_mcp/domain/policies/tenant_isolation.py +34 -0
  27. ml_mcp/domain/policies/tool_policy.py +185 -0
  28. ml_mcp/domain/value_objects/__init__.py +41 -0
  29. ml_mcp/domain/value_objects/dataset.py +60 -0
  30. ml_mcp/domain/value_objects/experiment.py +131 -0
  31. ml_mcp/domain/value_objects/model.py +48 -0
  32. ml_mcp/infrastructure/__init__.py +1 -0
  33. ml_mcp/infrastructure/object_storage/__init__.py +5 -0
  34. ml_mcp/infrastructure/object_storage/s3.py +272 -0
  35. ml_mcp/infrastructure/postgres/__init__.py +14 -0
  36. ml_mcp/infrastructure/postgres/base.py +40 -0
  37. ml_mcp/infrastructure/postgres/models.py +356 -0
  38. ml_mcp/infrastructure/postgres/repositories/__init__.py +21 -0
  39. ml_mcp/infrastructure/postgres/repositories/artifacts.py +34 -0
  40. ml_mcp/infrastructure/postgres/repositories/audit.py +34 -0
  41. ml_mcp/infrastructure/postgres/repositories/datasets.py +85 -0
  42. ml_mcp/infrastructure/postgres/repositories/experiments.py +134 -0
  43. ml_mcp/infrastructure/postgres/repositories/metrics.py +61 -0
  44. ml_mcp/infrastructure/postgres/repositories/models.py +61 -0
  45. ml_mcp/infrastructure/postgres/repositories/outbox.py +47 -0
  46. ml_mcp/infrastructure/postgres/repositories/projects.py +83 -0
  47. ml_mcp/infrastructure/postgres/session.py +117 -0
  48. ml_mcp/infrastructure/queue/__init__.py +6 -0
  49. ml_mcp/infrastructure/queue/outbox_processor.py +67 -0
  50. ml_mcp/infrastructure/queue/task_queue.py +98 -0
  51. ml_mcp/infrastructure/redis/__init__.py +11 -0
  52. ml_mcp/infrastructure/redis/client.py +64 -0
  53. ml_mcp/infrastructure/redis/rate_limiter.py +129 -0
  54. ml_mcp/infrastructure/telemetry/__init__.py +23 -0
  55. ml_mcp/infrastructure/telemetry/logging.py +155 -0
  56. ml_mcp/infrastructure/telemetry/otel.py +141 -0
  57. ml_mcp/mcp/__init__.py +1 -0
  58. ml_mcp/mcp/prompts.py +38 -0
  59. ml_mcp/mcp/resources.py +35 -0
  60. ml_mcp/mcp/schemas/__init__.py +1 -0
  61. ml_mcp/mcp/schemas/tools.py +124 -0
  62. ml_mcp/mcp/tools/__init__.py +63 -0
  63. ml_mcp/mcp/tools/analysis.py +37 -0
  64. ml_mcp/mcp/tools/datasets.py +82 -0
  65. ml_mcp/mcp/tools/discovery.py +82 -0
  66. ml_mcp/mcp/tools/experiments.py +96 -0
  67. ml_mcp/mcp/tools/projects.py +39 -0
  68. ml_mcp/mcp/tools/results.py +59 -0
  69. ml_mcp/server/__init__.py +1 -0
  70. ml_mcp/server/app.py +488 -0
  71. ml_mcp/server/auth.py +146 -0
  72. ml_mcp/server/context.py +59 -0
  73. ml_mcp/server/http.py +183 -0
  74. ml_mcp/server/middleware.py +121 -0
  75. ml_mcp/server/stdio.py +42 -0
  76. ml_mcp/workers/__init__.py +1 -0
  77. ml_mcp/workers/daemon.py +130 -0
  78. ml_mcp/workers/metrics.py +110 -0
  79. ml_mcp/workers/preprocessing.py +170 -0
  80. ml_mcp/workers/runner.py +250 -0
  81. ml_mcp/workers/sandbox.py +35 -0
  82. ml_mcp/workers/trainers/__init__.py +52 -0
  83. ml_mcp/workers/trainers/base.py +52 -0
  84. ml_mcp/workers/trainers/catboost_model.py +85 -0
  85. ml_mcp/workers/trainers/lightgbm_model.py +81 -0
  86. ml_mcp/workers/trainers/sklearn_models.py +266 -0
  87. ml_mcp/workers/trainers/xgboost_model.py +82 -0
  88. quicklab_mcp-0.1.0.dist-info/METADATA +540 -0
  89. quicklab_mcp-0.1.0.dist-info/RECORD +92 -0
  90. quicklab_mcp-0.1.0.dist-info/WHEEL +5 -0
  91. quicklab_mcp-0.1.0.dist-info/entry_points.txt +4 -0
  92. quicklab_mcp-0.1.0.dist-info/top_level.txt +1 -0
ml_mcp/__init__.py ADDED
@@ -0,0 +1,3 @@
1
+ """Agentic ML MCP Server (ModelLab)."""
2
+
3
+ __version__ = "0.1.0"
@@ -0,0 +1 @@
1
+ """Application layer package."""
@@ -0,0 +1,15 @@
1
+ """Structured analysis package."""
2
+
3
+ from ml_mcp.application.analysis.diagnostics import (
4
+ DataLeakageDetector,
5
+ ErrorPatternAnalyzer,
6
+ OverfittingDetector,
7
+ )
8
+ from ml_mcp.application.analysis.service import AnalysisService
9
+
10
+ __all__ = [
11
+ "OverfittingDetector",
12
+ "DataLeakageDetector",
13
+ "ErrorPatternAnalyzer",
14
+ "AnalysisService",
15
+ ]
@@ -0,0 +1,122 @@
1
+ """Diagnostic evaluators for overfitting, data leakage, and error patterns."""
2
+
3
+ from typing import Any
4
+
5
+ import numpy as np
6
+
7
+
8
+ class OverfittingDetector:
9
+ """Diagnoses generalization gaps between training and validation splits."""
10
+
11
+ @staticmethod
12
+ def evaluate(
13
+ train_metrics: dict[str, float],
14
+ val_metrics: dict[str, float],
15
+ threshold: float = 0.15,
16
+ ) -> dict[str, Any]:
17
+ gap_details: dict[str, float] = {}
18
+ overfitting_detected = False
19
+
20
+ common_keys = set(train_metrics.keys()).intersection(val_metrics.keys())
21
+ for k in common_keys:
22
+ tr_v = train_metrics[k]
23
+ val_v = val_metrics[k]
24
+ if (
25
+ isinstance(tr_v, (int, float))
26
+ and isinstance(val_v, (int, float))
27
+ and k in ("accuracy", "roc_auc", "pr_auc", "f1", "f1_macro", "r2")
28
+ ):
29
+ gap = tr_v - val_v
30
+ gap_details[k] = round(gap, 4)
31
+ if gap >= threshold:
32
+ overfitting_detected = True
33
+
34
+ return {
35
+ "overfitting_detected": overfitting_detected,
36
+ "metric_gaps": gap_details,
37
+ "severity": "HIGH"
38
+ if any(g > 0.25 for g in gap_details.values())
39
+ else ("MEDIUM" if overfitting_detected else "LOW"),
40
+ "observation": (
41
+ f"Validation performance lags training metrics by {max(gap_details.values()):.2%}, indicating model memorization."
42
+ if overfitting_detected and gap_details
43
+ else "Generalization gap is within acceptable operational tolerance."
44
+ ),
45
+ }
46
+
47
+
48
+ class DataLeakageDetector:
49
+ """Detects suspicious statistical patterns indicative of label or feature leakage."""
50
+
51
+ @staticmethod
52
+ def evaluate(
53
+ train_metrics: dict[str, float],
54
+ val_metrics: dict[str, float],
55
+ ) -> dict[str, Any]:
56
+ signals: list[str] = []
57
+ suspicious = False
58
+
59
+ # Suspiciously perfect validation performance
60
+ for k in ("roc_auc", "accuracy", "f1"):
61
+ if val_metrics.get(k) is not None and val_metrics[k] >= 0.999:
62
+ signals.append(
63
+ f"Metric '{k}' is suspiciously perfect ({val_metrics[k]:.4f}) on validation split."
64
+ )
65
+ suspicious = True
66
+
67
+ # Validation strictly outperforms training by a large margin (often indicates wrong split or inverted leakage)
68
+ for k in ("accuracy", "f1", "roc_auc"):
69
+ if k in train_metrics and k in val_metrics and val_metrics[k] - train_metrics[k] > 0.20:
70
+ signals.append(
71
+ f"Validation '{k}' ({val_metrics[k]:.2f}) significantly exceeds training ({train_metrics[k]:.2f})."
72
+ )
73
+ suspicious = True
74
+
75
+ return {
76
+ "leakage_risk_detected": suspicious,
77
+ "signals": signals,
78
+ "confidence": "HIGH" if len(signals) > 1 else ("MEDIUM" if suspicious else "NONE"),
79
+ }
80
+
81
+
82
+ class ErrorPatternAnalyzer:
83
+ """Analyzes confusion matrices and residual error distributions."""
84
+
85
+ @staticmethod
86
+ def evaluate_confusion_matrix(cm: list[list[int]]) -> dict[str, Any]:
87
+ """Analyze false positives and false negatives in binary or multiclass matrix."""
88
+ cm_arr = np.array(cm)
89
+ total_samples = int(cm_arr.sum())
90
+ if total_samples == 0:
91
+ return {"status": "EMPTY"}
92
+
93
+ if cm_arr.shape == (2, 2):
94
+ tn, fp, fn, tp = cm_arr.ravel()
95
+ fp_rate = float(fp / (fp + tn)) if (fp + tn) > 0 else 0.0
96
+ fn_rate = float(fn / (fn + tp)) if (fn + tp) > 0 else 0.0
97
+
98
+ dominant_error = (
99
+ "FALSE_POSITIVES" if fp > fn else ("FALSE_NEGATIVES" if fn > fp else "BALANCED")
100
+ )
101
+
102
+ return {
103
+ "total_samples": total_samples,
104
+ "true_positives": int(tp),
105
+ "true_negatives": int(tn),
106
+ "false_positives": int(fp),
107
+ "false_negatives": int(fn),
108
+ "false_positive_rate": round(fp_rate, 4),
109
+ "false_negative_rate": round(fn_rate, 4),
110
+ "dominant_error_pattern": dominant_error,
111
+ }
112
+ else:
113
+ # Multiclass
114
+ diagonal = np.diag(cm_arr)
115
+ per_class_accuracy = diagonal / np.maximum(cm_arr.sum(axis=1), 1)
116
+ lowest_class = int(np.argmin(per_class_accuracy))
117
+ return {
118
+ "total_samples": total_samples,
119
+ "classes_count": int(cm_arr.shape[0]),
120
+ "lowest_performing_class_index": lowest_class,
121
+ "lowest_class_accuracy": round(float(per_class_accuracy[lowest_class]), 4),
122
+ }
@@ -0,0 +1,155 @@
1
+ """Structured experiment analysis application service."""
2
+
3
+ import logging
4
+ from typing import Any
5
+
6
+ from sqlalchemy import select
7
+ from sqlalchemy.ext.asyncio import AsyncSession
8
+
9
+ from ml_mcp.application.analysis.diagnostics import (
10
+ DataLeakageDetector,
11
+ OverfittingDetector,
12
+ )
13
+ from ml_mcp.domain.errors import InvalidInputError, ResourceNotFoundError
14
+ from ml_mcp.infrastructure.postgres.base import generate_uuid7
15
+ from ml_mcp.infrastructure.postgres.models import AnalysisRunOrm, ExperimentOrm, MetricOrm
16
+
17
+ logger = logging.getLogger(__name__)
18
+
19
+
20
+ class AnalysisService:
21
+ """Service providing structured, read-only analytical evaluation of ML experiment results."""
22
+
23
+ def __init__(self, session: AsyncSession) -> None:
24
+ self.session = session
25
+
26
+ async def analyze_experiment(self, tenant_id: str, experiment_id: str) -> dict[str, Any]:
27
+ """Perform comprehensive statistical analysis of an experiment's metrics and predictions."""
28
+ # 1. Fetch experiment
29
+ stmt = select(ExperimentOrm).where(
30
+ ExperimentOrm.tenant_id == tenant_id, ExperimentOrm.id == experiment_id
31
+ )
32
+ result = await self.session.execute(stmt)
33
+ exp = result.scalar_one_or_none()
34
+ if not exp:
35
+ raise ResourceNotFoundError("Experiment", experiment_id)
36
+
37
+ if exp.status != "SUCCEEDED":
38
+ raise InvalidInputError(
39
+ f"Experiment '{experiment_id}' has not completed successfully (current status: {exp.status}). "
40
+ "Only completed experiments with recorded metrics can be analyzed."
41
+ )
42
+
43
+ # 2. Fetch metrics
44
+ m_stmt = select(MetricOrm).where(MetricOrm.experiment_id == experiment_id)
45
+ m_result = await self.session.execute(m_stmt)
46
+ all_metrics = m_result.scalars().all()
47
+ if not all_metrics:
48
+ raise InvalidInputError(
49
+ f"Experiment '{experiment_id}' has no recorded evaluation metrics to analyze."
50
+ )
51
+
52
+ train_metrics: dict[str, float] = {}
53
+ val_metrics: dict[str, float] = {}
54
+
55
+ for m in all_metrics:
56
+ if m.split == "train":
57
+ train_metrics[m.metric_name] = m.metric_value
58
+ elif m.split == "validation":
59
+ val_metrics[m.metric_name] = m.metric_value
60
+
61
+ # Check for confusion matrix in artifacts or metrics
62
+ # 3. Run diagnostic engines
63
+ overfitting = OverfittingDetector.evaluate(train_metrics, val_metrics)
64
+ leakage = DataLeakageDetector.evaluate(train_metrics, val_metrics)
65
+
66
+ # 4. Synthesize recommendations
67
+ recommendations: list[str] = []
68
+ if overfitting["overfitting_detected"]:
69
+ recommendations.append(
70
+ "Increase regularization (e.g. increase alpha, reduce max_depth, or increase min_samples_split)."
71
+ )
72
+ recommendations.append(
73
+ "Apply cross-validation (K-Fold) to evaluate variance across folds."
74
+ )
75
+ if leakage["leakage_risk_detected"]:
76
+ recommendations.append(
77
+ "Audit dataset columns for target leakage or deterministic identifiers."
78
+ )
79
+ if not recommendations:
80
+ recommendations.append(
81
+ "Model demonstrates healthy generalization. Consider testing next candidate architecture or fine-tuning learning rate."
82
+ )
83
+
84
+ # 5. Format according to Architecture.md Section 21
85
+ structured_findings: dict[str, Any] = {
86
+ "summary": f"Experiment {experiment_id} evaluated with {len(val_metrics)} validation metrics.",
87
+ "metric_comparison": {
88
+ "train": train_metrics,
89
+ "validation": val_metrics,
90
+ },
91
+ "statistical_observations": {
92
+ "primary_metric": exp.spec_json.get("evaluation_config", {}).get("primary_metric"),
93
+ "primary_metric_validation_value": val_metrics.get(
94
+ exp.spec_json.get("evaluation_config", {}).get("primary_metric", "")
95
+ ),
96
+ },
97
+ "error_patterns": {
98
+ "observed_anomalies": len(leakage["signals"]),
99
+ },
100
+ "data_quality_findings": {
101
+ "target_column": exp.spec_json.get("target_column"),
102
+ "task_type": exp.spec_json.get("task_type"),
103
+ },
104
+ "overfitting_signals": overfitting,
105
+ "leakage_signals": leakage,
106
+ "limitations": [
107
+ "Analysis performed strictly on tabular evaluation split.",
108
+ "Hyperparameter space explored only for single trial run.",
109
+ ],
110
+ "recommended_next_experiments": recommendations,
111
+ "confidence": "HIGH",
112
+ }
113
+
114
+ # 6. Upsert read-only analysis run record
115
+ analysis_run_stmt = select(AnalysisRunOrm).where(
116
+ AnalysisRunOrm.experiment_id == experiment_id
117
+ )
118
+ existing_run = (await self.session.execute(analysis_run_stmt)).scalar_one_or_none()
119
+
120
+ if not existing_run:
121
+ analysis_orm = AnalysisRunOrm(
122
+ id=generate_uuid7(),
123
+ experiment_id=experiment_id,
124
+ tenant_id=tenant_id,
125
+ summary_json={"summary": structured_findings["summary"]},
126
+ diagnostics_json=structured_findings,
127
+ )
128
+ self.session.add(analysis_orm)
129
+ await self.session.flush()
130
+
131
+ return structured_findings
132
+
133
+ async def analyze_model_errors(self, tenant_id: str, experiment_id: str) -> dict[str, Any]:
134
+ """Deep dive into false positives and false negatives."""
135
+ analysis = await self.analyze_experiment(tenant_id, experiment_id)
136
+ return {
137
+ "experiment_id": experiment_id,
138
+ "error_patterns": analysis["error_patterns"],
139
+ "overfitting_diagnosis": analysis["overfitting_signals"],
140
+ }
141
+
142
+ async def check_experiment_validity(self, tenant_id: str, experiment_id: str) -> dict[str, Any]:
143
+ """Validate whether an experiment's metrics can be trusted or should be flagged for review."""
144
+ analysis = await self.analyze_experiment(tenant_id, experiment_id)
145
+ is_valid = not (
146
+ analysis["leakage_signals"]["leakage_risk_detected"]
147
+ or analysis["overfitting_signals"]["severity"] == "HIGH"
148
+ )
149
+ return {
150
+ "experiment_id": experiment_id,
151
+ "is_valid": is_valid,
152
+ "leakage_risk": analysis["leakage_signals"]["leakage_risk_detected"],
153
+ "overfitting_severity": analysis["overfitting_signals"]["severity"],
154
+ "recommendations": analysis["recommended_next_experiments"],
155
+ }
@@ -0,0 +1,5 @@
1
+ """Artifacts application module."""
2
+
3
+ from ml_mcp.application.artifacts.service import ArtifactService
4
+
5
+ __all__ = ["ArtifactService"]
@@ -0,0 +1,149 @@
1
+ """Artifact and results application service managing metrics, predictions, and model artifacts."""
2
+
3
+ import contextlib
4
+ import io
5
+ from typing import Any
6
+
7
+ import pandas as pd
8
+ from sqlalchemy.ext.asyncio import AsyncSession
9
+
10
+ from ml_mcp.domain.errors import ResourceNotFoundError
11
+ from ml_mcp.infrastructure.object_storage.s3 import S3StorageService, get_storage_service
12
+ from ml_mcp.infrastructure.postgres.repositories.artifacts import ArtifactRepository
13
+ from ml_mcp.infrastructure.postgres.repositories.experiments import ExperimentRepository
14
+ from ml_mcp.infrastructure.postgres.repositories.metrics import MetricRepository
15
+
16
+
17
+ class ArtifactService:
18
+ """Encapsulates querying of experiment results, metrics, prediction samples, and stored artifacts."""
19
+
20
+ def __init__(
21
+ self,
22
+ session: AsyncSession,
23
+ storage_service: S3StorageService | None = None,
24
+ ) -> None:
25
+ self.session = session
26
+ self.storage = storage_service or get_storage_service()
27
+ self.exp_repo = ExperimentRepository(session)
28
+ self.metric_repo = MetricRepository(session)
29
+ self.artifact_repo = ArtifactRepository(session)
30
+
31
+ async def get_metrics(
32
+ self,
33
+ tenant_id: str,
34
+ experiment_id: str,
35
+ split: str | None = None,
36
+ ) -> dict[str, Any]:
37
+ """Fetch metrics for an experiment with tenant isolation and anti-enumeration enforcement."""
38
+ exp = await self.exp_repo.get_experiment(tenant_id, experiment_id)
39
+ if not exp:
40
+ raise ResourceNotFoundError("Experiment", experiment_id)
41
+
42
+ metrics = await self.metric_repo.get_metrics(tenant_id, experiment_id, split)
43
+
44
+ formatted: dict[str, Any] = {}
45
+ for m in metrics:
46
+ key = f"{m.split}_{m.metric_name}" if not split else m.metric_name
47
+ formatted[key] = m.metric_value
48
+
49
+ return {
50
+ "experiment_id": experiment_id,
51
+ "split": split or "all",
52
+ "metrics": formatted,
53
+ }
54
+
55
+ async def get_predictions(
56
+ self,
57
+ tenant_id: str,
58
+ experiment_id: str,
59
+ limit: int = 5,
60
+ ) -> dict[str, Any]:
61
+ """Fetch sample predictions and generate presigned download URL with tenant verification."""
62
+ exp = await self.exp_repo.get_experiment(tenant_id, experiment_id)
63
+ if not exp:
64
+ raise ResourceNotFoundError("Experiment", experiment_id)
65
+
66
+ key = self.storage.get_experiment_artifact_key(
67
+ tenant_id, experiment_id, "predictions", "predictions.parquet"
68
+ )
69
+
70
+ download_url = self.storage.generate_presigned_get_url(key, expires_in=900)
71
+
72
+ sample_preview = []
73
+ with contextlib.suppress(Exception):
74
+ data = self.storage.get_object(key)
75
+ df = pd.read_parquet(io.BytesIO(data))
76
+ sample_preview = df.head(limit).to_dict(orient="records")
77
+
78
+ return {
79
+ "experiment_id": experiment_id,
80
+ "sample_preview": sample_preview,
81
+ "download_url": download_url,
82
+ "expires_in_seconds": 900,
83
+ }
84
+
85
+ async def list_artifacts(
86
+ self,
87
+ tenant_id: str,
88
+ experiment_id: str,
89
+ ) -> list[dict[str, Any]]:
90
+ """List stored artifacts for an experiment after tenant ownership check."""
91
+ exp = await self.exp_repo.get_experiment(tenant_id, experiment_id)
92
+ if not exp:
93
+ raise ResourceNotFoundError("Experiment", experiment_id)
94
+
95
+ artifacts = await self.artifact_repo.list_artifacts(tenant_id, experiment_id)
96
+ return [
97
+ {
98
+ "artifact_id": a.id,
99
+ "artifact_type": a.artifact_type,
100
+ "size_bytes": a.size_bytes,
101
+ "content_hash": a.content_hash,
102
+ "created_at": a.created_at.isoformat(),
103
+ }
104
+ for a in artifacts
105
+ ]
106
+
107
+ async def get_artifact(
108
+ self,
109
+ tenant_id: str,
110
+ artifact_id: str,
111
+ ) -> dict[str, Any]:
112
+ """Fetch artifact metadata and presigned URL by artifact ID."""
113
+ art = await self.artifact_repo.get_artifact(tenant_id, artifact_id)
114
+ if not art:
115
+ raise ResourceNotFoundError("Artifact", artifact_id)
116
+
117
+ prefix = f"s3://{self.storage.bucket_name}/"
118
+ key = (
119
+ art.storage_uri[len(prefix) :]
120
+ if art.storage_uri.startswith(prefix)
121
+ else art.storage_uri
122
+ )
123
+ download_url = self.storage.generate_presigned_get_url(key, expires_in=900)
124
+
125
+ return {
126
+ "artifact_id": art.id,
127
+ "artifact_type": art.artifact_type,
128
+ "size_bytes": art.size_bytes,
129
+ "content_hash": art.content_hash,
130
+ "download_url": download_url,
131
+ "metadata": art.metadata_json,
132
+ }
133
+
134
+ async def compare_experiments(
135
+ self,
136
+ tenant_id: str,
137
+ experiment_ids: list[str],
138
+ ) -> dict[str, Any]:
139
+ """Compare validation metrics across experiments ensuring caller owns every experiment."""
140
+ for exp_id in experiment_ids:
141
+ exp = await self.exp_repo.get_experiment(tenant_id, exp_id)
142
+ if not exp:
143
+ raise ResourceNotFoundError("Experiment", exp_id)
144
+
145
+ comparison = await self.metric_repo.compare_metrics(tenant_id, experiment_ids)
146
+ return {
147
+ "compared_experiment_ids": experiment_ids,
148
+ "validation_metrics_by_experiment": comparison,
149
+ }
@@ -0,0 +1,6 @@
1
+ """Dataset management and ingestion package."""
2
+
3
+ from ml_mcp.application.datasets.service import DatasetService
4
+ from ml_mcp.application.datasets.validator import DatasetValidator
5
+
6
+ __all__ = ["DatasetValidator", "DatasetService"]
@@ -0,0 +1,178 @@
1
+ """Dataset application service managing registration, validation, and inspection."""
2
+
3
+ import io
4
+ from typing import Any
5
+
6
+ import pandas as pd
7
+ from sqlalchemy.ext.asyncio import AsyncSession
8
+
9
+ from ml_mcp.application.datasets.validator import DatasetValidator
10
+ from ml_mcp.domain.errors import ResourceNotFoundError
11
+ from ml_mcp.domain.value_objects import DatasetFormat
12
+ from ml_mcp.infrastructure.object_storage.s3 import S3StorageService, get_storage_service
13
+ from ml_mcp.infrastructure.postgres.base import generate_uuid7
14
+ from ml_mcp.infrastructure.postgres.models import DatasetOrm, DatasetVersionOrm
15
+ from ml_mcp.infrastructure.postgres.repositories.datasets import DatasetRepository
16
+
17
+
18
+ class DatasetService:
19
+ """Application service for dataset lifecycle and profiling."""
20
+
21
+ def __init__(
22
+ self,
23
+ session: AsyncSession,
24
+ storage_service: S3StorageService | None = None,
25
+ validator: DatasetValidator | None = None,
26
+ ) -> None:
27
+ self.session = session
28
+ self.repo = DatasetRepository(session)
29
+ self.storage = storage_service or get_storage_service()
30
+ self.validator = validator or DatasetValidator()
31
+
32
+ async def register_dataset(
33
+ self,
34
+ tenant_id: str,
35
+ project_id: str,
36
+ name: str,
37
+ description: str,
38
+ data_format: DatasetFormat | str,
39
+ raw_bytes: bytes,
40
+ version: str = "1.0",
41
+ ) -> dict[str, Any]:
42
+ """Validate, store to S3, and register a new dataset version."""
43
+ # 1. Validate & extract schema
44
+ df, parquet_bytes, schema = self.validator.process(raw_bytes, data_format)
45
+
46
+ # 2. Check or create parent Dataset record
47
+ dataset_id = generate_uuid7()
48
+ version_id = generate_uuid7()
49
+
50
+ dataset_key = self.storage.get_dataset_key(tenant_id, dataset_id, version, "data.parquet")
51
+ storage_uri, _ = self.storage.put_object(
52
+ key=dataset_key,
53
+ data=parquet_bytes,
54
+ content_type="application/vnd.apache.parquet",
55
+ metadata={"row_count": str(schema.row_count), "column_count": str(schema.column_count)},
56
+ )
57
+
58
+ # 3. Persist metadata to DB
59
+ fmt = (
60
+ data_format.value
61
+ if isinstance(data_format, DatasetFormat)
62
+ else str(data_format).lower()
63
+ )
64
+ dataset_orm = DatasetOrm(
65
+ id=dataset_id,
66
+ tenant_id=tenant_id,
67
+ project_id=project_id,
68
+ name=name,
69
+ description=description,
70
+ format=fmt,
71
+ )
72
+ await self.repo.create_dataset(dataset_orm)
73
+
74
+ version_orm = DatasetVersionOrm(
75
+ id=version_id,
76
+ dataset_id=dataset_id,
77
+ version=version,
78
+ schema_json=schema.model_dump(),
79
+ row_count=schema.row_count,
80
+ column_count=schema.column_count,
81
+ size_bytes=schema.size_bytes,
82
+ content_hash=schema.content_hash,
83
+ storage_uri=storage_uri,
84
+ validation_status="VALIDATED",
85
+ )
86
+ await self.repo.create_version(version_orm)
87
+
88
+ return {
89
+ "dataset_id": dataset_id,
90
+ "version_id": version_id,
91
+ "version": version,
92
+ "name": name,
93
+ "format": fmt,
94
+ "row_count": schema.row_count,
95
+ "column_count": schema.column_count,
96
+ "content_hash": schema.content_hash,
97
+ "storage_uri": storage_uri,
98
+ }
99
+
100
+ async def validate_only(
101
+ self,
102
+ raw_bytes: bytes,
103
+ data_format: DatasetFormat | str,
104
+ ) -> dict[str, Any]:
105
+ """Validate dataset and extract schema without persisting."""
106
+ _, _, schema = self.validator.process(raw_bytes, data_format)
107
+ return schema.model_dump()
108
+
109
+ async def get_dataset(self, tenant_id: str, dataset_id: str) -> dict[str, Any]:
110
+ ds = await self.repo.get_dataset(tenant_id, dataset_id)
111
+ if not ds:
112
+ raise ResourceNotFoundError("Dataset", dataset_id)
113
+
114
+ return {
115
+ "dataset_id": ds.id,
116
+ "name": ds.name,
117
+ "description": ds.description,
118
+ "format": ds.format,
119
+ "versions": [
120
+ {
121
+ "version_id": v.id,
122
+ "version": v.version,
123
+ "row_count": v.row_count,
124
+ "column_count": v.column_count,
125
+ "content_hash": v.content_hash,
126
+ "created_at": v.created_at.isoformat(),
127
+ }
128
+ for v in ds.versions
129
+ ],
130
+ }
131
+
132
+ async def list_datasets(self, tenant_id: str, project_id: str | None = None) -> list[dict[str, Any]]:
133
+ datasets = await self.repo.list_datasets(tenant_id, project_id)
134
+ return [
135
+ {
136
+ "dataset_id": ds.id,
137
+ "project_id": ds.project_id,
138
+ "name": ds.name,
139
+ "format": ds.format,
140
+ "latest_version": ds.versions[0].version if ds.versions else None,
141
+ }
142
+ for ds in datasets
143
+ ]
144
+
145
+ async def inspect_dataset(
146
+ self,
147
+ tenant_id: str,
148
+ dataset_id: str,
149
+ version: str = "1.0",
150
+ sample_rows: int = 5,
151
+ ) -> dict[str, Any]:
152
+ """Retrieve dataset schema, summary statistics, and head sample rows."""
153
+ dsv = await self.repo.get_version(tenant_id, dataset_id, version)
154
+ if not dsv:
155
+ raise ResourceNotFoundError("DatasetVersion", f"{dataset_id}:{version}")
156
+
157
+ # Download parquet from storage and load sample rows
158
+ key = self.storage.get_dataset_key(tenant_id, dataset_id, version, "data.parquet")
159
+ data = self.storage.get_object(key)
160
+ df = pd.read_parquet(io.BytesIO(data))
161
+
162
+ head_df = df.head(sample_rows)
163
+ # Convert to records format with string/native conversion
164
+ sample_data = head_df.to_dict(orient="records")
165
+
166
+ # Generate presigned download URL for full dataset access
167
+ download_url = self.storage.generate_presigned_get_url(key, expires_in=900)
168
+
169
+ return {
170
+ "dataset_id": dataset_id,
171
+ "version": version,
172
+ "schema": dsv.schema_json,
173
+ "row_count": dsv.row_count,
174
+ "column_count": dsv.column_count,
175
+ "size_bytes": dsv.size_bytes,
176
+ "sample_rows": sample_data,
177
+ "download_url": download_url,
178
+ }