Update AgML/mcp_output/mcp_plugin/mcp_service.py

AgML/mcp_output/mcp_plugin/mcp_service.py

@@ -1,103 +1,268 @@
-import os
-import sys
-
-# Add the local source directory to sys.path
-source_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "source")
-if source_path not in sys.path:
-    sys.path.insert(0, source_path)
-
 from fastmcp import FastMCP
 
-# Import core modules
-from agml.data.loader import DataLoader
-from agml.models.classification import ClassificationModel
-from agml.models.detection import DetectionModel
-from agml.models.segmentation import SegmentationModel
-
 # Create the FastMCP service application
 mcp = FastMCP("agml_service")
 
-@mcp.tool(name="load_data", description="Load dataset using DataLoader.")
-def load_data(dataset_name: str) -> dict:
+
+@mcp.tool(name="list_available_datasets", description="List all available AgML public datasets")
+def list_available_datasets() -> dict:
+    """
+    List all available public datasets in AgML.
+
+    Returns:
+    - dict: A dictionary with success status and list of available datasets.
+    """
+    try:
+        from agml.data.public import public_data_sources
+        datasets = list(public_data_sources().keys())
+        return {
+            "success": True,
+            "datasets": datasets,
+            "count": len(datasets)
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+
+@mcp.tool(name="get_dataset_info", description="Get information about a specific dataset")
+def get_dataset_info(dataset_name: str) -> dict:
+    """
+    Get detailed information about a specific AgML dataset.
+
+    Parameters:
+    - dataset_name: Name of the dataset (e.g., 'bean_disease_uganda')
+
+    Returns:
+    - dict: Dataset information including task, location, classes, etc.
+    """
+    try:
+        from agml.data.public import public_data_sources
+        sources = public_data_sources()
+        
+        if dataset_name not in sources:
+            return {
+                "success": False,
+                "error": f"Dataset '{dataset_name}' not found. Use list_available_datasets to see available options."
+            }
+        
+        info = sources[dataset_name]
+        return {
+            "success": True,
+            "dataset_name": dataset_name,
+            "info": info
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+
+@mcp.tool(name="load_dataset", description="Load an AgML dataset")
+def load_dataset(dataset_name: str, batch_size: int = 8) -> dict:
+    """
+    Load an AgML dataset.
+
+    Parameters:
+    - dataset_name: Name of the dataset to load
+    - batch_size: Batch size for data loading (default: 8)
+
+    Returns:
+    - dict: Information about the loaded dataset
+    """
+    try:
+        from agml.data import AgMLDataLoader
+        
+        loader = AgMLDataLoader(dataset_name)
+        
+        return {
+            "success": True,
+            "dataset_name": dataset_name,
+            "task_type": str(loader.info['task_type']),
+            "num_images": loader.num_images,
+            "num_classes": getattr(loader, 'num_classes', None),
+            "classes": getattr(loader, 'classes', None)
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+
+@mcp.tool(name="download_dataset", description="Download a public AgML dataset")
+def download_dataset(dataset_name: str) -> dict:
     """
-    Load a dataset using the DataLoader.
+    Download a public dataset from AgML.
 
     Parameters:
-    - dataset_name (str): The name of the dataset to load.
+    - dataset_name: Name of the dataset to download
 
     Returns:
-    - dict: A dictionary containing success, result, or error fields.
+    - dict: Download status information
     """
     try:
-        loader = DataLoader(dataset_name)
-        data = loader.load()
-        return {"success": True, "result": data}
+        from agml.data.public import download_public_dataset
+        
+        download_public_dataset(dataset_name)
+        
+        return {
+            "success": True,
+            "message": f"Dataset '{dataset_name}' downloaded successfully"
+        }
     except Exception as e:
         return {"success": False, "error": str(e)}
 
-@mcp.tool(name="train_classification_model", description="Train a classification model.")
-def train_classification_model(dataset_name: str, model_name: str) -> dict:
+
+@mcp.tool(name="list_model_benchmarks", description="List available model benchmarks")
+def list_model_benchmarks() -> dict:
+    """
+    List all available model benchmarks in AgML.
+
+    Returns:
+    - dict: Available model benchmarks by task type
+    """
+    try:
+        from agml.models.benchmarks import list_model_benchmarks
+        
+        benchmarks = list_model_benchmarks()
+        
+        return {
+            "success": True,
+            "benchmarks": benchmarks
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+
+@mcp.tool(name="get_model_benchmark", description="Get benchmark results for a specific model")
+def get_model_benchmark(model_name: str, dataset_name: str) -> dict:
     """
-    Train a classification model on a specified dataset.
+    Get benchmark results for a specific model on a dataset.
 
     Parameters:
-    - dataset_name (str): The name of the dataset to use for training.
-    - model_name (str): The name of the classification model to train.
+    - model_name: Name of the model
+    - dataset_name: Name of the dataset
 
     Returns:
-    - dict: A dictionary containing success, result, or error fields.
+    - dict: Benchmark results
     """
     try:
-        model = ClassificationModel(model_name)
-        loader = DataLoader(dataset_name)
-        data = loader.load()
-        model.train(data)
-        return {"success": True, "result": "Model trained successfully"}
+        from agml.models.benchmarks import get_model_benchmark
+        
+        result = get_model_benchmark(model_name, dataset_name)
+        
+        return {
+            "success": True,
+            "model": model_name,
+            "dataset": dataset_name,
+            "benchmark": result
+        }
     except Exception as e:
         return {"success": False, "error": str(e)}
 
-@mcp.tool(name="train_detection_model", description="Train a detection model.")
-def train_detection_model(dataset_name: str, model_name: str) -> dict:
+
+@mcp.tool(name="create_classification_model", description="Create an AgML classification model")
+def create_classification_model(model_name: str, num_classes: int) -> dict:
     """
-    Train a detection model on a specified dataset.
+    Create an AgML classification model.
 
     Parameters:
-    - dataset_name (str): The name of the dataset to use for training.
-    - model_name (str): The name of the detection model to train.
+    - model_name: Name of the model architecture (e.g., 'resnet18', 'efficientnet_b0')
+    - num_classes: Number of classes for classification
 
     Returns:
-    - dict: A dictionary containing success, result, or error fields.
+    - dict: Model creation status and information
     """
     try:
-        model = DetectionModel(model_name)
-        loader = DataLoader(dataset_name)
-        data = loader.load()
-        model.train(data)
-        return {"success": True, "result": "Model trained successfully"}
+        from agml.models import ClassificationModel
+        
+        model = ClassificationModel(model_name=model_name, num_classes=num_classes)
+        
+        return {
+            "success": True,
+            "model_name": model_name,
+            "num_classes": num_classes,
+            "message": "Model created successfully"
+        }
     except Exception as e:
         return {"success": False, "error": str(e)}
 
-@mcp.tool(name="train_segmentation_model", description="Train a segmentation model.")
-def train_segmentation_model(dataset_name: str, model_name: str) -> dict:
+
+@mcp.tool(name="create_detection_model", description="Create an AgML object detection model")
+def create_detection_model(model_name: str, num_classes: int) -> dict:
     """
-    Train a segmentation model on a specified dataset.
+    Create an AgML object detection model.
 
     Parameters:
-    - dataset_name (str): The name of the dataset to use for training.
-    - model_name (str): The name of the segmentation model to train.
+    - model_name: Name of the model architecture (e.g., 'fasterrcnn_resnet50_fpn')
+    - num_classes: Number of classes for detection
 
     Returns:
-    - dict: A dictionary containing success, result, or error fields.
+    - dict: Model creation status and information
     """
     try:
-        model = SegmentationModel(model_name)
-        loader = DataLoader(dataset_name)
-        data = loader.load()
-        model.train(data)
-        return {"success": True, "result": "Model trained successfully"}
+        from agml.models import DetectionModel
+        
+        model = DetectionModel(model_name=model_name, num_classes=num_classes)
+        
+        return {
+            "success": True,
+            "model_name": model_name,
+            "num_classes": num_classes,
+            "message": "Model created successfully"
+        }
     except Exception as e:
         return {"success": False, "error": str(e)}
 
+
+@mcp.tool(name="create_segmentation_model", description="Create an AgML segmentation model")
+def create_segmentation_model(model_name: str, num_classes: int) -> dict:
+    """
+    Create an AgML segmentation model.
+
+    Parameters:
+    - model_name: Name of the model architecture (e.g., 'deeplabv3_resnet50')
+    - num_classes: Number of classes for segmentation
+
+    Returns:
+    - dict: Model creation status and information
+    """
+    try:
+        from agml.models import SegmentationModel
+        
+        model = SegmentationModel(model_name=model_name, num_classes=num_classes)
+        
+        return {
+            "success": True,
+            "model_name": model_name,
+            "num_classes": num_classes,
+            "message": "Model created successfully"
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+
+@mcp.tool(name="export_dataset_to_yolo", description="Export AgML dataset to YOLO format")
+def export_dataset_to_yolo(dataset_name: str, output_dir: str) -> dict:
+    """
+    Export an AgML dataset to YOLO format.
+
+    Parameters:
+    - dataset_name: Name of the dataset to export
+    - output_dir: Directory where to save the YOLO formatted data
+
+    Returns:
+    - dict: Export status
+    """
+    try:
+        from agml.data import AgMLDataLoader
+        
+        loader = AgMLDataLoader(dataset_name)
+        loader.export_to_yolo(output_dir)
+        
+        return {
+            "success": True,
+            "message": f"Dataset exported to YOLO format at {output_dir}"
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+
 def create_app() -> FastMCP:
     """
     Create and return the FastMCP application instance.