Update osmnx/mcp_output/mcp_plugin/mcp_service.py

osmnx/mcp_output/mcp_plugin/mcp_service.py

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