Update yt/mcp_output/mcp_plugin/mcp_service.py

yt/mcp_output/mcp_plugin/mcp_service.py

@@ -1,72 +1,404 @@
 import os
 import sys
-
-# Path settings to include the local source directory
-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 typing import List, Optional, Dict, Any
 
 from fastmcp import FastMCP
-from yt.loaders import load
-from yt.data_objects.static_output import Dataset
-from yt.visualization.plot_window import PlotWindow
-from yt.fields.field_info_container import FieldInfoContainer
+
+# Import yt from PyPI (will be installed via requirements.txt)
+import yt
+from yt import load
+from yt.units import dimensions
 
 # Create the FastMCP service application
 mcp = FastMCP("yt_service")
 
-@mcp.tool(name="load_dataset", description="Load a dataset from a file.")
-def load_dataset(file_path: str) -> dict:
+
+@mcp.tool(name="get_yt_version", description="Get yt library version and configuration")
+def get_yt_version() -> dict:
+    """
+    Get the yt library version and configuration.
+
+    Returns:
+    - dict: Version and configuration information.
+    """
+    try:
+        return {
+            "success": True,
+            "result": {
+                "version": yt.__version__,
+                "description": "yt is a toolkit for analyzing and visualizing volumetric data"
+            },
+            "error": None
+        }
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
+@mcp.tool(name="list_sample_datasets", description="List available sample datasets in yt")
+def list_sample_datasets() -> dict:
+    """
+    List available sample datasets that can be loaded with yt.load_sample().
+
+    Returns:
+    - dict: List of sample datasets.
+    """
+    try:
+        # Common sample datasets in yt
+        datasets = {
+            "IsolatedGalaxy": "Isolated galaxy simulation (ENZO)",
+            "enzo_tiny_cosmology": "Tiny cosmology dataset (ENZO)",
+            "output_00080": "RAMSES output",
+            "GasSloshingLowRes": "Gas sloshing simulation",
+            "WindTunnel": "Wind tunnel test problem",
+            "DD0010": "ENZO data dump",
+            "cluster_merger_rockstar_halos": "Cluster merger with Rockstar halos",
+            "KelvinHelmholtz_hdf5_chk_0100": "Kelvin-Helmholtz instability (FLASH)",
+        }
+        return {"success": True, "result": datasets, "error": None}
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
+@mcp.tool(name="list_frontends", description="List available data frontends in yt")
+def list_frontends() -> dict:
+    """
+    List available data frontends (simulation code formats) supported by yt.
+
+    Returns:
+    - dict: Dictionary of frontends and descriptions.
+    """
+    try:
+        frontends = {
+            "enzo": "Enzo AMR cosmology code",
+            "flash": "FLASH multiphysics code",
+            "ramses": "RAMSES AMR code",
+            "athena": "Athena MHD code",
+            "athena_pp": "Athena++ MHD code",
+            "gadget": "Gadget N-body/SPH code",
+            "tipsy": "Tipsy/ChaNGa SPH code",
+            "arepo": "Arepo moving-mesh code",
+            "gizmo": "GIZMO meshless code",
+            "swift": "SWIFT SPH code",
+            "art": "ART cosmology code",
+            "chombo": "Chombo AMR library",
+            "boxlib": "BoxLib/AMReX framework",
+            "fits": "FITS astronomical format",
+            "hdf5": "Generic HDF5",
+            "stream": "In-memory particle/grid data"
+        }
+        return {"success": True, "result": frontends, "error": None}
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
+@mcp.tool(name="list_field_types", description="List common field types in yt")
+def list_field_types() -> dict:
+    """
+    List common field types available in yt datasets.
+
+    Returns:
+    - dict: Field categories and examples.
     """
-    Load a dataset from a file.
+    try:
+        field_types = {
+            "gas": {
+                "description": "Gas/fluid fields",
+                "examples": ["density", "temperature", "pressure", "velocity_x", "velocity_y", "velocity_z", "entropy", "specific_thermal_energy"]
+            },
+            "particle": {
+                "description": "Particle fields (dark matter, stars)",
+                "examples": ["particle_mass", "particle_position", "particle_velocity"]
+            },
+            "derived": {
+                "description": "Derived/computed fields",
+                "examples": ["cell_volume", "cell_mass", "sound_speed", "mach_number", "magnetic_field_strength"]
+            },
+            "index": {
+                "description": "Grid index fields",
+                "examples": ["x", "y", "z", "dx", "dy", "dz", "radius"]
+            }
+        }
+        return {"success": True, "result": field_types, "error": None}
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
 
-    :param file_path: Path to the dataset file.
-    :return: Dictionary with success status and dataset object or error message.
+@mcp.tool(name="list_plot_types", description="List available plot types in yt")
+def list_plot_types() -> dict:
+    """
+    List available visualization/plot types in yt.
+
+    Returns:
+    - dict: Plot types and descriptions.
     """
     try:
-        ds = load(file_path)
-        return {"success": True, "result": ds}
+        plot_types = {
+            "SlicePlot": {
+                "description": "2D slice through data at a specific location",
+                "usage": "yt.SlicePlot(ds, 'z', 'density')"
+            },
+            "ProjectionPlot": {
+                "description": "2D projection (integral) along an axis",
+                "usage": "yt.ProjectionPlot(ds, 'z', 'density')"
+            },
+            "PhasePlot": {
+                "description": "2D histogram of two fields colored by third",
+                "usage": "yt.PhasePlot(ad, 'density', 'temperature', 'cell_mass')"
+            },
+            "ProfilePlot": {
+                "description": "1D binned profile of fields",
+                "usage": "yt.ProfilePlot(ad, 'radius', 'density')"
+            },
+            "ParticlePlot": {
+                "description": "Scatter plot of particle positions",
+                "usage": "yt.ParticlePlot(ds, 'particle_position_x', 'particle_position_y')"
+            },
+            "OffAxisSlicePlot": {
+                "description": "Slice along arbitrary vector",
+                "usage": "yt.OffAxisSlicePlot(ds, normal, 'density')"
+            },
+            "OffAxisProjectionPlot": {
+                "description": "Projection along arbitrary vector",
+                "usage": "yt.OffAxisProjectionPlot(ds, normal, 'density')"
+            }
+        }
+        return {"success": True, "result": plot_types, "error": None}
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
+
 
-@mcp.tool(name="create_plot_window", description="Create a plot window for a dataset.")
-def create_plot_window(ds: Dataset, field: str, axis: int) -> dict:
+@mcp.tool(name="list_data_objects", description="List available data selection objects in yt")
+def list_data_objects() -> dict:
     """
-    Create a plot window for a dataset.
+    List available data selection objects for extracting regions of data.
 
-    :param ds: The dataset object.
-    :param field: The field to plot.
-    :param axis: The axis to plot along.
-    :return: Dictionary with success status and plot window object or error message.
+    Returns:
+    - dict: Data object types and descriptions.
     """
     try:
-        pw = PlotWindow(ds, axis, fields=[field])
-        return {"success": True, "result": pw}
+        data_objects = {
+            "all_data": {
+                "description": "Select entire domain",
+                "usage": "ds.all_data()"
+            },
+            "sphere": {
+                "description": "Spherical region",
+                "usage": "ds.sphere(center, radius)"
+            },
+            "region": {
+                "description": "Rectangular box region",
+                "usage": "ds.region(center, left_edge, right_edge)"
+            },
+            "disk": {
+                "description": "Cylindrical disk region",
+                "usage": "ds.disk(center, normal, radius, height)"
+            },
+            "ray": {
+                "description": "1D ray through domain",
+                "usage": "ds.ray(start_point, end_point)"
+            },
+            "slice": {
+                "description": "2D slice at fixed coordinate",
+                "usage": "ds.slice(axis, coordinate)"
+            },
+            "covering_grid": {
+                "description": "Uniformly spaced grid covering region",
+                "usage": "ds.covering_grid(level, left_edge, dims)"
+            },
+            "arbitrary_grid": {
+                "description": "Grid with arbitrary spacing",
+                "usage": "ds.arbitrary_grid(left_edge, right_edge, dims)"
+            },
+            "cut_region": {
+                "description": "Region defined by conditional expression",
+                "usage": "ad.cut_region(['obj[\"temperature\"] > 1e6'])"
+            }
+        }
+        return {"success": True, "result": data_objects, "error": None}
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
 
-@mcp.tool(name="get_field_info", description="Get field information from a dataset.")
-def get_field_info(ds: Dataset, field: str) -> dict:
+
+@mcp.tool(name="list_units", description="List common units and unit systems in yt")
+def list_units() -> dict:
     """
-    Get field information from a dataset.
+    List common units and unit systems available in yt.
 
-    :param ds: The dataset object.
-    :param field: The field to get information about.
-    :return: Dictionary with success status and field info or error message.
+    Returns:
+    - dict: Unit categories and examples.
     """
     try:
-        field_info = ds.field_info[field]
-        return {"success": True, "result": field_info}
+        units = {
+            "length": ["cm", "m", "km", "pc", "kpc", "Mpc", "AU", "ly", "Rsun"],
+            "mass": ["g", "kg", "Msun", "Mjup", "Mearth"],
+            "time": ["s", "yr", "Myr", "Gyr"],
+            "velocity": ["cm/s", "m/s", "km/s", "km/h"],
+            "density": ["g/cm**3", "kg/m**3", "Msun/kpc**3"],
+            "temperature": ["K", "keV"],
+            "energy": ["erg", "J", "eV", "keV"],
+            "magnetic": ["gauss", "T"],
+            "cgs_units": "Base CGS system (cm, g, s)",
+            "mks_units": "Base MKS/SI system (m, kg, s)",
+            "galactic_units": "Galactic units (kpc, Msun, Myr)"
+        }
+        return {"success": True, "result": units, "error": None}
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "result": None, "error": str(e)}
+
+
+@mcp.tool(name="list_colormaps", description="List recommended colormaps for yt")
+def list_colormaps() -> dict:
+    """
+    List recommended colormaps for scientific visualization in yt.
+
+    Returns:
+    - dict: Colormap categories and names.
+    """
+    try:
+        colormaps = {
+            "yt_native": ["algae", "kamae", "arbre", "octarine", "kelp", "dusk", "B-W LINEAR"],
+            "sequential": ["viridis", "plasma", "inferno", "magma", "cividis"],
+            "diverging": ["RdBu", "RdYlBu", "coolwarm", "seismic", "bwr"],
+            "perceptually_uniform": ["viridis", "plasma", "inferno", "magma", "cividis"],
+            "notes": {
+                "default": "arbre is the default yt colormap",
+                "recommendation": "Use perceptually uniform colormaps for quantitative data",
+                "diverging_use": "Diverging colormaps are good for data with a meaningful midpoint"
+            }
+        }
+        return {"success": True, "result": colormaps, "error": None}
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+
+@mcp.tool(name="get_yt_code_example", description="Get example code for common yt operations")
+def get_yt_code_example(operation: str) -> dict:
+    """
+    Get example code for common yt operations.
+
+    Parameters:
+    - operation: The operation to get example code for.
+      Options: "load", "slice", "projection", "profile", "phase", "sphere", "volume_render"
+
+    Returns:
+    - dict: Example code and explanation.
+    """
+    try:
+        examples = {
+            "load": {
+                "code": """
+import yt
+# Load a dataset
+ds = yt.load("path/to/dataset")
+# Print basic info
+print(ds)
+print(ds.field_list)
+print(ds.derived_field_list)
+""",
+                "description": "Load a dataset and inspect its properties"
+            },
+            "slice": {
+                "code": """
+import yt
+ds = yt.load("path/to/dataset")
+# Create a slice plot along z-axis
+slc = yt.SlicePlot(ds, 'z', 'density')
+slc.set_cmap('density', 'viridis')
+slc.annotate_grids()  # Show AMR grid structure
+slc.save('slice.png')
+""",
+                "description": "Create a 2D slice plot of density"
+            },
+            "projection": {
+                "code": """
+import yt
+ds = yt.load("path/to/dataset")
+# Create a projection (column density)
+prj = yt.ProjectionPlot(ds, 'z', 'density', weight_field='density')
+prj.set_unit('density', 'g/cm**2')
+prj.save('projection.png')
+""",
+                "description": "Create a 2D projection plot"
+            },
+            "profile": {
+                "code": """
+import yt
+ds = yt.load("path/to/dataset")
+ad = ds.all_data()
+# Create 1D radial profile
+profile = yt.create_profile(ad, 'radius', 'density', weight_field='cell_mass')
+# Plot it
+plot = yt.ProfilePlot.from_profiles(profile)
+plot.save('profile.png')
+""",
+                "description": "Create a 1D radial profile"
+            },
+            "phase": {
+                "code": """
+import yt
+ds = yt.load("path/to/dataset")
+ad = ds.all_data()
+# Create 2D phase plot
+phase = yt.PhasePlot(ad, 'density', 'temperature', 'cell_mass', weight_field=None)
+phase.set_unit('density', 'g/cm**3')
+phase.set_unit('temperature', 'K')
+phase.save('phase.png')
+""",
+                "description": "Create a 2D phase diagram"
+            },
+            "sphere": {
+                "code": """
+import yt
+ds = yt.load("path/to/dataset")
+# Create spherical region around center
+center = ds.domain_center
+sp = ds.sphere(center, (100, 'kpc'))
+# Get data from sphere
+density = sp['gas', 'density']
+temperature = sp['gas', 'temperature']
+# Compute quantities
+total_mass = sp.quantities.total_mass()
+""",
+                "description": "Extract data from a spherical region"
+            },
+            "volume_render": {
+                "code": """
+import yt
+ds = yt.load("path/to/dataset")
+# Create volume rendering
+sc = yt.create_scene(ds, field='density')
+# Customize transfer function
+source = sc[0]
+source.tfh.set_bounds((1e-30, 1e-23))
+source.tfh.set_log(True)
+# Render and save
+sc.save('volume_render.png', sigma_clip=4)
+""",
+                "description": "Create a 3D volume rendering"
+            }
+        }
+        
+        if operation.lower() in examples:
+            return {"success": True, "result": examples[operation.lower()], "error": None}
+        else:
+            return {
+                "success": True,
+                "result": {
+                    "available_operations": list(examples.keys()),
+                    "message": f"Operation '{operation}' not found. Available: {list(examples.keys())}"
+                },
+                "error": None
+            }
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
 
 def create_app() -> FastMCP:
     """
     Create and return the FastMCP application instance.
 
-    :return: FastMCP instance.
+    Returns:
+    - FastMCP: The FastMCP application instance.
     """
     return mcp
-
-# Ensure the application is created when the module is imported
-app = create_app()
+