Update xmitgcm/mcp_output/mcp_plugin/mcp_service.py

xmitgcm/mcp_output/mcp_plugin/mcp_service.py

@@ -1,146 +1,414 @@
+import os
+import sys
+from typing import List, Optional, Dict, Any
+
+# 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
-from xmitgcm import open_mdsdataset
 
-# Create the FastMCP service application
+# Import xmitgcm modules
+from xmitgcm import utils, file_utils
+from xmitgcm.llcreader import known_models
+
 mcp = FastMCP("xmitgcm_service")
 
-@mcp.tool(name="list_available_datasets", description="List all available MITgcm datasets")
-def list_available_datasets() -> dict:
+
+# ===================== LLC Model Tools =====================
+
+@mcp.tool(name="get_known_models", description="Retrieve known LLC models")
+def get_known_models_tool() -> dict:
     """
-    List all available MITgcm datasets.
+    Retrieve list of known LLC (Lat-Lon-Cap) models available in xmitgcm.
 
     Returns:
-    - dict: A dictionary with success status and list of available datasets.
+    - Dictionary containing available model classes and their descriptions.
     """
     try:
-        # Example implementation (replace with actual logic)
-        datasets = ["dataset1", "dataset2", "dataset3"]
-        return {"success": True, "datasets": datasets}
+        models_info = {
+            "LLC90Model": {
+                "nx": 90,
+                "nz": 50,
+                "delta_t": 3600,
+                "description": "LLC90 model configuration"
+            },
+            "LLC2160Model": {
+                "nx": 2160,
+                "nz": 90,
+                "delta_t": 45,
+                "description": "High-resolution LLC2160 model"
+            },
+            "LLC4320Model": {
+                "nx": 4320,
+                "nz": 90,
+                "delta_t": 25,
+                "description": "Ultra high-resolution LLC4320 model"
+            },
+            "ASTE270Model": {
+                "nx": 270,
+                "nz": 50,
+                "nface": 6,
+                "domain": "aste",
+                "description": "Arctic Subpolar gyre sTate Estimate model"
+            },
+            "ECCOPortalLLC2160Model": {
+                "description": "ECCO Portal LLC2160 model with remote data access"
+            },
+            "ECCOPortalLLC4320Model": {
+                "description": "ECCO Portal LLC4320 model with remote data access"
+            },
+            "CRIOSPortalASTE270Model": {
+                "description": "CRIOS Portal ASTE270 model with remote data access"
+            }
+        }
+        return {"success": True, "result": models_info, "error": None}
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
 
-@mcp.tool(name="get_dataset_info", description="Get information about a specific dataset")
-def get_dataset_info(dataset_name: str) -> dict:
+
+@mcp.tool(name="load_llc_model", description="Load LLC model data")
+def load_llc_model(model_name: str) -> dict:
     """
-    Get detailed information about a specific MITgcm dataset.
+    Get information about an LLC model configuration.
 
     Parameters:
-    - dataset_name: Name of the dataset
+    - model_name: Name of the LLC model (e.g., 'LLC90Model', 'LLC2160Model', 'LLC4320Model', 'ASTE270Model').
 
     Returns:
-    - dict: Dataset information including dimensions, variables, etc.
+    - Dictionary with model configuration details.
     """
     try:
-        # Example implementation (replace with actual logic)
-        info = {
-            "name": dataset_name,
-            "dimensions": ["X", "Y", "Z", "time"],
-            "variables": ["temperature", "salinity", "velocity"]
+        model_configs = {
+            "LLC90Model": {
+                "nx": 90, "nz": 50, "delta_t": 3600,
+                "time_units": "seconds since 1948-01-01 12:00:00",
+                "calendar": "gregorian"
+            },
+            "LLC2160Model": {
+                "nx": 2160, "nz": 90, "delta_t": 45,
+                "iter_start": 92160, "iter_stop": 1586401, "iter_step": 80,
+                "time_units": "seconds since 2011-01-17",
+                "calendar": "gregorian",
+                "varnames": ['Eta', 'KPPhbl', 'oceFWflx', 'oceQnet', 'oceQsw', 'oceSflux',
+                            'oceTAUX', 'oceTAUY', 'PhiBot', 'Salt', 'SIarea', 'SIheff',
+                            'SIhsalt', 'SIhsnow', 'SIuice', 'SIvice', 'Theta', 'U', 'V', 'W']
+            },
+            "LLC4320Model": {
+                "nx": 4320, "nz": 90, "delta_t": 25,
+                "iter_start": 10368, "iter_stop": 1495153, "iter_step": 144,
+                "time_units": "seconds since 2011-09-10",
+                "calendar": "gregorian",
+                "varnames": ['Eta', 'KPPhbl', 'oceFWflx', 'oceQnet', 'oceQsw', 'oceSflux',
+                            'oceTAUX', 'oceTAUY', 'PhiBot', 'Salt', 'SIarea', 'SIheff',
+                            'SIhsalt', 'SIhsnow', 'SIuice', 'SIvice', 'Theta', 'U', 'V', 'W']
+            },
+            "ASTE270Model": {
+                "nface": 6, "nx": 270, "nz": 50, "domain": "aste",
+                "delta_t": 600,
+                "time_units": "seconds since 2002-01-01",
+                "calendar": "gregorian"
+            }
         }
-        return {"success": True, "info": info}
+        
+        if model_name in model_configs:
+            return {"success": True, "result": model_configs[model_name], "error": None}
+        else:
+            return {
+                "success": False, 
+                "result": None, 
+                "error": f"Unknown model: {model_name}. Available models: {list(model_configs.keys())}"
+            }
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
+
 
-@mcp.tool(name="load_dataset", description="Load a MITgcm dataset")
-def load_dataset(dataset_path: str, grid_dir: str = None, iters: list = None) -> dict:
+# ===================== MDS File Utilities =====================
+
+@mcp.tool(name="parse_meta_file", description="Parse MITgcm .meta file to extract metadata")
+def parse_meta_file_tool(fname: str) -> dict:
     """
-    Load a MITgcm dataset.
+    Parse an MITgcm .meta file and extract metadata.
 
     Parameters:
-    - dataset_path: Path to the dataset directory
-    - grid_dir: Path to the grid directory (optional)
-    - iters: List of iterations to load (optional)
+    - fname: Path to the .meta file.
 
     Returns:
-    - dict: Information about the loaded dataset
+    - Dictionary containing parsed metadata (dimList, nDims, nrecords, dataprec, etc.).
     """
     try:
-        ds = open_mdsdataset(dataset_path, grid_dir=grid_dir, iters=iters)
-        return {"success": True, "dataset": str(ds)}
+        meta = utils.parse_meta_file(fname)
+        # Convert numpy types to Python native types for JSON serialization
+        result = {}
+        for key, value in meta.items():
+            if hasattr(value, 'tolist'):
+                result[key] = value.tolist()
+            elif hasattr(value, 'name'):  # dtype
+                result[key] = str(value)
+            else:
+                result[key] = value
+        return {"success": True, "result": result, "error": None}
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
+
 
-@mcp.tool(name="download_dataset", description="Download a MITgcm dataset")
-def download_dataset(dataset_url: str, output_dir: str) -> dict:
+# ===================== File Utilities =====================
+
+@mcp.tool(name="list_directory", description="List contents of a directory with caching")
+def list_directory_tool(path: str) -> dict:
     """
-    Download a MITgcm dataset from a given URL.
+    List contents of a directory (with caching for performance).
 
     Parameters:
-    - dataset_url: URL of the dataset to download
-    - output_dir: Directory to save the downloaded dataset
+    - path: Path to the directory.
+
+    Returns:
+    - List of files and directories in the path.
+    """
+    try:
+        files = file_utils.listdir(path)
+        return {"success": True, "result": files, "error": None}
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
+@mcp.tool(name="list_files_by_pattern", description="List files matching a pattern")
+def list_files_by_pattern(path: str, pattern: str, match_type: str = "fnmatch") -> dict:
+    """
+    List files in a directory matching a specific pattern.
+
+    Parameters:
+    - path: Path to the directory.
+    - pattern: Pattern to match.
+    - match_type: Type of matching - 'startswith', 'endswith', or 'fnmatch' (default).
+
+    Returns:
+    - List of matching files.
+    """
+    try:
+        if match_type == "startswith":
+            files = file_utils.listdir_startswith(path, pattern)
+        elif match_type == "endswith":
+            files = file_utils.listdir_endswith(path, pattern)
+        else:
+            files = file_utils.listdir_fnmatch(path, pattern)
+        return {"success": True, "result": files, "error": None}
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
+@mcp.tool(name="clear_file_cache", description="Clear the file listing cache")
+def clear_file_cache() -> dict:
+    """
+    Clear the cached file listings.
 
     Returns:
-    - dict: Status of the download
+    - Success status.
     """
     try:
-        # Example implementation (replace with actual logic)
-        # Download logic here
-        return {"success": True, "message": f"Dataset downloaded to {output_dir}"}
+        file_utils.clear_cache()
+        return {"success": True, "result": "Cache cleared successfully", "error": None}
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
 
-@mcp.tool(name="list_model_benchmarks", description="List available model benchmarks")
-def list_model_benchmarks() -> dict:
+
+# ===================== Grid Metrics Tools =====================
+
+@mcp.tool(name="calculate_grid_metrics", description="Calculate grid metrics")
+def calculate_grid_metrics(grid_dir: str) -> dict:
     """
-    List available model benchmarks.
+    Get information about grid variables available in MITgcm.
+
+    Parameters:
+    - grid_dir: Directory containing grid files.
 
     Returns:
-    - dict: A dictionary with success status and list of benchmarks.
+    - Dictionary with grid variable information.
     """
     try:
-        # Example implementation (replace with actual logic)
-        benchmarks = ["benchmark1", "benchmark2", "benchmark3"]
-        return {"success": True, "benchmarks": benchmarks}
+        # Return information about standard MITgcm grid variables
+        grid_info = {
+            "horizontal_grid_variables": [
+                "XC", "YC", "XG", "YG", "dxC", "dyC", "dxG", "dyG", 
+                "dxF", "dyF", "dxV", "dyU", "rA", "rAw", "rAs", "rAz"
+            ],
+            "vertical_grid_variables": [
+                "drC", "drF", "Z", "Zp1", "Zu", "Zl", "PHrefC", "PHrefF"
+            ],
+            "volume_grid_variables": [
+                "hFacC", "hFacW", "hFacS", "Depth"
+            ],
+            "description": "Standard MITgcm grid variables for staggered Arakawa C-grid"
+        }
+        
+        # Check if directory exists and list available grid files
+        if os.path.exists(grid_dir):
+            meta_files = [f for f in os.listdir(grid_dir) if f.endswith('.meta')]
+            data_files = [f for f in os.listdir(grid_dir) if f.endswith('.data')]
+            grid_info["available_meta_files"] = meta_files[:20]  # Limit to 20
+            grid_info["available_data_files"] = data_files[:20]
+            grid_info["grid_dir_exists"] = True
+        else:
+            grid_info["grid_dir_exists"] = False
+            grid_info["message"] = f"Directory {grid_dir} does not exist"
+            
+        return {"success": True, "result": grid_info, "error": None}
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "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:
+# ===================== Variable Metadata Tools =====================
+
+@mcp.tool(name="get_variable_info", description="Get metadata for MITgcm variables")
+def get_variable_info(variable_name: Optional[str] = None) -> dict:
     """
-    Get benchmark results for a specific model and dataset.
+    Get metadata information for MITgcm state variables.
 
     Parameters:
-    - model_name: Name of the model
-    - dataset_name: Name of the dataset
+    - variable_name: Optional specific variable name. If None, returns all available variables.
 
     Returns:
-    - dict: Benchmark results
+    - Dictionary with variable metadata.
     """
     try:
-        # Example implementation (replace with actual logic)
-        results = {
-            "model": model_name,
-            "dataset": dataset_name,
-            "accuracy": 0.95,
-            "runtime": "10 minutes"
+        from xmitgcm.variables import state_variables, package_state_variables
+        
+        all_vars = {}
+        all_vars.update(state_variables)
+        all_vars.update(package_state_variables)
+        
+        if variable_name:
+            if variable_name in all_vars:
+                var_info = all_vars[variable_name]
+                # Convert to JSON-serializable format
+                result = {
+                    "name": variable_name,
+                    "dims": var_info.get("dims", []),
+                    "attrs": var_info.get("attrs", {})
+                }
+                return {"success": True, "result": result, "error": None}
+            else:
+                return {
+                    "success": False, 
+                    "result": None, 
+                    "error": f"Variable '{variable_name}' not found. Use without variable_name to list all."
+                }
+        else:
+            # Return list of available variables
+            var_list = list(all_vars.keys())
+            return {"success": True, "result": {"available_variables": var_list, "count": len(var_list)}, "error": None}
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
+# ===================== Dimension and Coordinate Tools =====================
+
+@mcp.tool(name="get_dimension_info", description="Get MITgcm grid dimension information")
+def get_dimension_info() -> dict:
+    """
+    Get information about MITgcm grid dimensions and coordinates.
+
+    Returns:
+    - Dictionary with dimension metadata.
+    """
+    try:
+        from xmitgcm.variables import dimensions, vertical_coordinates
+        
+        dims_info = {}
+        for dim_name, dim_data in dimensions.items():
+            dims_info[dim_name] = {
+                "dims": dim_data.get("dims", []),
+                "attrs": dim_data.get("attrs", {})
+            }
+        
+        vert_info = {}
+        for coord_name, coord_data in vertical_coordinates.items():
+            vert_info[coord_name] = {
+                "dims": coord_data.get("dims", []),
+                "attrs": coord_data.get("attrs", {})
+            }
+        
+        result = {
+            "dimensions": dims_info,
+            "vertical_coordinates": vert_info,
+            "description": "MITgcm uses staggered Arakawa C-grid with i,j horizontal and k vertical indices"
+        }
+        return {"success": True, "result": result, "error": None}
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
+# ===================== Geometry Tools =====================
+
+@mcp.tool(name="get_geometry_info", description="Get information about supported grid geometries")
+def get_geometry_info() -> dict:
+    """
+    Get information about MITgcm grid geometries supported by xmitgcm.
+
+    Returns:
+    - Dictionary describing available geometries.
+    """
+    try:
+        geometries = {
+            "sphericalpolar": {
+                "description": "Standard lat-lon spherical polar grid",
+                "coordinates": ["XC", "YC", "XG", "YG"],
+                "use_case": "Regional and global ocean models"
+            },
+            "cartesian": {
+                "description": "Cartesian coordinate grid",
+                "coordinates": ["XC", "YC", "XG", "YG"],
+                "use_case": "Idealized simulations, process studies"
+            },
+            "llc": {
+                "description": "Lat-Lon-Cap grid (13 faces)",
+                "coordinates": ["XC", "YC", "XG", "YG", "AngleCS", "AngleSN"],
+                "use_case": "Global ECCO simulations (LLC90, LLC2160, LLC4320)",
+                "num_faces": 13
+            },
+            "curvilinear": {
+                "description": "Curvilinear orthogonal grid",
+                "coordinates": ["XC", "YC", "XG", "YG", "AngleCS", "AngleSN"],
+                "use_case": "Regional models with complex boundaries"
+            },
+            "cs": {
+                "description": "Cubed-sphere grid (6 faces)",
+                "coordinates": ["XC", "YC", "XG", "YG", "AngleCS", "AngleSN"],
+                "use_case": "Global atmospheric and ocean models",
+                "num_faces": 6
+            }
         }
-        return {"success": True, "results": results}
+        return {"success": True, "result": geometries, "error": None}
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
 
-@mcp.tool(name="export_dataset_to_yolo", description="Export MITgcm dataset to YOLO format")
-def export_dataset_to_yolo(dataset_name: str, output_dir: str) -> dict:
+
+# ===================== Extra Metadata Tools =====================
+
+@mcp.tool(name="get_extra_metadata", description="Get extra metadata for LLC/ASTE configurations")
+def get_extra_metadata_tool(geometry: str = "llc90") -> dict:
     """
-    Export a MITgcm dataset to YOLO format.
+    Get extra metadata needed for LLC and ASTE grid configurations.
 
     Parameters:
-    - dataset_name: Name of the dataset to export
-    - output_dir: Directory to save the YOLO formatted dataset
+    - geometry: Grid geometry type ('llc90', 'llc', 'aste', etc.)
 
     Returns:
-    - dict: Status of the export
+    - Dictionary with extra metadata for the specified geometry.
     """
     try:
-        # Example implementation (replace with actual logic)
-        return {"success": True, "message": f"Dataset {dataset_name} exported to {output_dir} in YOLO format."}
+        extra_meta = utils.get_extra_metadata(geometry)
+        return {"success": True, "result": extra_meta, "error": None}
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
+
 
 def create_app() -> FastMCP:
     """
     Create and return the FastMCP application instance.
 
     Returns:
-    - FastMCP: The FastMCP application instance.
+    - FastMCP instance.
     """
     return mcp