Update osmnx/mcp_output/mcp_plugin/mcp_service.py

osmnx/mcp_output/mcp_plugin/mcp_service.py

@@ -1,45 +1,273 @@
 from fastmcp import FastMCP
 
 # Create the FastMCP service application
-mcp = FastMCP("osmnx_service")
+mcp = FastMCP("agml_service")
 
-@mcp.tool(name="get_graph_from_place", description="Retrieve a graph from a place name using OSMnx")
-def get_graph_from_place(place_name: str, network_type: str = "drive") -> 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:
+    """
+    Download a public dataset from AgML.
+
+    Parameters:
+    - dataset_name: Name of the dataset to download
+
+    Returns:
+    - dict: Download status information
+    """
+    try:
+        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="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:
+    """
+    Get benchmark results for a specific model on a dataset.
+
+    Parameters:
+    - model_name: Name of the model
+    - dataset_name: Name of the dataset
+
+    Returns:
+    - dict: Benchmark results
+    """
+    try:
+        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="create_classification_model", description="Create an AgML classification model")
+def create_classification_model(model_name: str, num_classes: int) -> dict:
+    """
+    Create an AgML classification model.
+
+    Parameters:
+    - model_name: Name of the model architecture (e.g., 'resnet18', 'efficientnet_b0')
+    - num_classes: Number of classes for classification
+
+    Returns:
+    - dict: Model creation status and information
+    """
+    try:
+        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="create_detection_model", description="Create an AgML object detection model")
+def create_detection_model(model_name: str, num_classes: int) -> dict:
+    """
+    Create an AgML object detection model.
+
+    Parameters:
+    - model_name: Name of the model architecture (e.g., 'fasterrcnn_resnet50_fpn')
+    - num_classes: Number of classes for detection
+
+    Returns:
+    - dict: Model creation status and information
+    """
+    try:
+        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:
     """
-    Retrieve a graph representation of a place using OSMnx.
+    Create an AgML segmentation model.
 
     Parameters:
-    - place_name: Name of the place to retrieve the graph for
-    - network_type: Type of network to retrieve (e.g., 'drive', 'walk', 'bike')
+    - model_name: Name of the model architecture (e.g., 'deeplabv3_resnet50')
+    - num_classes: Number of classes for segmentation
 
     Returns:
-    - dict: Information about the retrieved graph
+    - dict: Model creation status and information
     """
     try:
-        import osmnx as ox
-        graph = ox.graph_from_place(place_name, network_type=network_type)
+        from agml.models import SegmentationModel
+        
+        model = SegmentationModel(model_name=model_name, num_classes=num_classes)
+        
         return {
             "success": True,
-            "graph_info": str(graph)
+            "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="plot_graph", description="Plot a graph using OSMnx")
-def plot_graph(graph_data: dict) -> dict:
+
+@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:
     """
-    Plot a graph using OSMnx.
+    Export an AgML dataset to YOLO format.
 
     Parameters:
-    - graph_data: Serialized graph data to plot
+    - dataset_name: Name of the dataset to export
+    - output_dir: Directory where to save the YOLO formatted data
 
     Returns:
-    - dict: Status of the plotting operation
+    - dict: Export status
     """
     try:
-        import osmnx as ox
-        graph = ox.io.load_graphml(graph_data)
-        ox.plot_graph(graph)
-        return {"success": True}
+        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)}
+        return {"success": False, "error": str(e)}
+
+
+def create_app() -> FastMCP:
+    """
+    Create and return the FastMCP application instance.
+
+    Returns:
+    - FastMCP: The FastMCP application instance.
+    """
+    return mcp