diff --git a/CHANGELOG.md b/CHANGELOG.md index c7572ad..4c97ad3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -50,6 +50,11 @@ under their entry. endpoint state (no endpoint configured, or endpoint not ready) instead of surfacing a misleading `FileNotFoundError` from the local fallback. Local paths and HEALPix specs still fall back silently. (#27) +- `plot_mesh`, `plot_variable`, and `plot_zonal_mean` HPC wrappers now + accept `session_id` + `dataset_handle` and resolve grid/data paths + from a registered session dataset, matching the handle pattern + already used by `subset_bbox` / `subset_polygon`. Direct paths still + work unchanged. (#25) ## 0.1.0 — initial scaffold diff --git a/src/uxarray_mcp/tools/remote_tools.py b/src/uxarray_mcp/tools/remote_tools.py index 171b680..ae263d4 100644 --- a/src/uxarray_mcp/tools/remote_tools.py +++ b/src/uxarray_mcp/tools/remote_tools.py @@ -377,12 +377,13 @@ def calculate_zonal_mean_hpc( def plot_mesh_hpc( - grid_path: str, + grid_path: str | None = None, width: int = 800, height: int = 400, use_remote: bool = False, endpoint: str | None = None, session_id: str | None = None, + dataset_handle: str | None = None, ) -> list[Any]: """Render a mesh wireframe PNG with optional HPC execution. @@ -392,15 +393,19 @@ def plot_mesh_hpc( Parameters ---------- - grid_path : str + grid_path : str | None Path to mesh file. Can be a local path or an HPC filesystem path - when use_remote=True. + when use_remote=True. Optional when ``session_id`` and + ``dataset_handle`` are provided. width : int Image width in pixels (default 800). height : int Image height in pixels (default 400). use_remote : bool If True and HPC is configured, render on the remote endpoint. + session_id, dataset_handle : str | None + When both are provided, the grid path is looked up from the + registered session dataset instead of ``grid_path``. Returns ------- @@ -415,13 +420,15 @@ def plot_mesh_hpc( >>> result = plot_mesh_hpc("/hpc/data/grid.nc", use_remote=True) >>> open("mesh.png", "wb").write(base64.b64decode(result["png_b64"])) """ - from .plotting import plot_mesh + from .plotting import _resolve_plot_paths, plot_mesh + + resolved_grid, _ = _resolve_plot_paths(grid_path, None, session_id, dataset_handle) def _local() -> Dict[str, Any]: import base64 import json - items = plot_mesh(grid_path, width=width, height=height) + items = plot_mesh(resolved_grid, width=width, height=height) # plot_mesh returns [ImageContent, TextContent]; extract png_b64 + metadata img = items[0] meta = json.loads(items[1].text) @@ -439,19 +446,19 @@ def _local() -> Dict[str, Any]: tool_name="plot_mesh_hpc", use_remote=use_remote, endpoint=endpoint, - path_hint=grid_path, + path_hint=resolved_grid, session_id=session_id, local_call=_local, remote_call=lambda agent: _run_sync( - lambda: agent.plot_mesh_remote(grid_path, width, height, use_remote) + lambda: agent.plot_mesh_remote(resolved_grid, width, height, use_remote) ), ) return _plot_result_to_mcp_contents(result) def plot_variable_hpc( - grid_path: str, - data_path: str, + grid_path: str | None = None, + data_path: str | None = None, variable_name: Optional[str] = None, width: int = 800, height: int = 400, @@ -462,15 +469,18 @@ def plot_variable_hpc( use_remote: bool = False, endpoint: str | None = None, session_id: str | None = None, + dataset_handle: str | None = None, ) -> list[Any]: """Render a face-centered variable as a filled-polygon PNG with optional HPC execution. Parameters ---------- - grid_path : str - Path to mesh grid file (local or HPC filesystem). - data_path : str - Path to data file (local or HPC filesystem). + grid_path : str | None + Path to mesh grid file (local or HPC filesystem). Optional when + ``session_id`` and ``dataset_handle`` are provided. + data_path : str | None + Path to data file (local or HPC filesystem). Optional when + ``session_id`` and ``dataset_handle`` are provided. variable_name : str | None Variable to plot. If None, the first face-centered variable is used. width : int @@ -487,6 +497,10 @@ def plot_variable_hpc( Custom plot title. use_remote : bool If True and HPC is configured, render on the remote endpoint. + session_id, dataset_handle : str | None + When both are provided, the grid and data paths are looked up + from the registered session dataset instead of ``grid_path`` / + ``data_path``. Returns ------- @@ -497,15 +511,19 @@ def plot_variable_hpc( - grid_info: {n_face, n_node, n_edge} - execution_venue: "local" or "hpc:" """ - from .plotting import plot_variable + from .plotting import _resolve_plot_paths, plot_variable + + resolved_grid, resolved_data = _resolve_plot_paths( + grid_path, data_path, session_id, dataset_handle + ) def _local() -> Dict[str, Any]: import base64 import json items = plot_variable( - grid_path, - data_path, + resolved_grid, + resolved_data, variable_name, width=width, height=height, @@ -531,13 +549,13 @@ def _local() -> Dict[str, Any]: tool_name="plot_variable_hpc", use_remote=use_remote, endpoint=endpoint, - path_hint=grid_path, + path_hint=resolved_grid, session_id=session_id, local_call=_local, remote_call=lambda agent: _run_sync( lambda: agent.plot_variable_remote( - grid_path, - data_path, + resolved_grid, + resolved_data, variable_name, width, height, @@ -553,9 +571,9 @@ def _local() -> Dict[str, Any]: def plot_zonal_mean_hpc( - grid_path: str, - data_path: str, - variable_name: str, + grid_path: str | None = None, + data_path: str | None = None, + variable_name: str | None = None, width: int = 800, height: int = 400, lat_spec: Optional[tuple | float | List] = None, @@ -565,15 +583,18 @@ def plot_zonal_mean_hpc( use_remote: bool = False, endpoint: str | None = None, session_id: str | None = None, + dataset_handle: str | None = None, ) -> list[Any]: """Render a zonal mean profile PNG with optional HPC execution. Parameters ---------- - grid_path : str - Path to mesh grid file (local or HPC filesystem). - data_path : str - Path to data file (local or HPC filesystem). + grid_path : str | None + Path to mesh grid file (local or HPC filesystem). Optional when + ``session_id`` and ``dataset_handle`` are provided. + data_path : str | None + Path to data file (local or HPC filesystem). Optional when + ``session_id`` and ``dataset_handle`` are provided. variable_name : str Face-centered variable to compute and plot. width : int @@ -590,6 +611,10 @@ def plot_zonal_mean_hpc( Custom plot title. use_remote : bool If True and HPC is configured, render on the remote endpoint. + session_id, dataset_handle : str | None + When both are provided, the grid and data paths are looked up + from the registered session dataset instead of ``grid_path`` / + ``data_path``. Returns ------- @@ -601,15 +626,19 @@ def plot_zonal_mean_hpc( - zonal_mean_values: List of zonal mean values - execution_venue: "local" or "hpc:" """ - from .plotting import plot_zonal_mean + from .plotting import _resolve_plot_paths, plot_zonal_mean + + resolved_grid, resolved_data = _resolve_plot_paths( + grid_path, data_path, session_id, dataset_handle + ) def _local() -> Dict[str, Any]: import base64 import json items = plot_zonal_mean( - grid_path, - data_path, + resolved_grid, + resolved_data, variable_name, width=width, height=height, @@ -636,13 +665,13 @@ def _local() -> Dict[str, Any]: tool_name="plot_zonal_mean_hpc", use_remote=use_remote, endpoint=endpoint, - path_hint=grid_path, + path_hint=resolved_grid, session_id=session_id, local_call=_local, remote_call=lambda agent: _run_sync( lambda: agent.plot_zonal_mean_remote( - grid_path, - data_path, + resolved_grid, + resolved_data, variable_name, width, height, diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 83f0a28..7f42c74 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -426,3 +426,110 @@ def test_plot_zonal_mean_hpc_returns_image_content(self, mock_run): metadata = json.loads(result[1].text) assert metadata["latitudes"] == [-90.0, 0.0, 90.0] assert "png_b64" not in metadata + + +class TestHpcPlotWrappersDatasetHandle: + """Issue #25: HPC plot wrappers must accept session_id + dataset_handle and + resolve the grid/data paths from the session before dispatching.""" + + @staticmethod + def _stub_run_result(extra: dict | None = None) -> dict: + return { + "png_b64": base64.b64encode(b"\x89PNG_fake").decode("utf-8"), + "image_size_bytes": 9, + "grid_info": {"n_face": 1}, + "_provenance": {"execution_venue": "local", "warnings": []}, + **(extra or {}), + } + + @patch("uxarray_mcp.tools.remote_tools._run_with_optional_hpc") + def test_plot_mesh_hpc_resolves_dataset_handle( + self, mock_run, synthetic_mesh_with_data + ): + from uxarray_mcp.tools import create_session, register_dataset + + grid_file, _ = synthetic_mesh_with_data + session = create_session("plot-mesh-hpc-handle") + registered = register_dataset( + session_id=session["session_id"], + grid_path=grid_file, + name="grid-only", + ) + + mock_run.return_value = self._stub_run_result() + + plot_mesh_hpc( + session_id=session["session_id"], + dataset_handle=registered["dataset_handle"], + ) + + kwargs = mock_run.call_args.kwargs + assert kwargs["path_hint"] == grid_file + + @patch("uxarray_mcp.tools.remote_tools._run_with_optional_hpc") + def test_plot_variable_hpc_resolves_dataset_handle( + self, mock_run, synthetic_mesh_with_data + ): + from uxarray_mcp.tools import create_session, register_dataset + + grid_file, data_file = synthetic_mesh_with_data + session = create_session("plot-variable-hpc-handle") + registered = register_dataset( + session_id=session["session_id"], + grid_path=grid_file, + data_path=data_file, + name="grid-and-data", + ) + + mock_run.return_value = self._stub_run_result({"variable_name": "temperature"}) + + plot_variable_hpc( + variable_name="temperature", + session_id=session["session_id"], + dataset_handle=registered["dataset_handle"], + ) + + kwargs = mock_run.call_args.kwargs + assert kwargs["path_hint"] == grid_file + + @patch("uxarray_mcp.tools.remote_tools._run_with_optional_hpc") + def test_plot_zonal_mean_hpc_resolves_dataset_handle( + self, mock_run, synthetic_mesh_with_data + ): + from uxarray_mcp.tools import create_session, register_dataset + + grid_file, data_file = synthetic_mesh_with_data + session = create_session("plot-zonal-hpc-handle") + registered = register_dataset( + session_id=session["session_id"], + grid_path=grid_file, + data_path=data_file, + name="zonal-handle", + ) + + mock_run.return_value = self._stub_run_result( + { + "variable_name": "temperature", + "latitudes": [-90.0, 0.0, 90.0], + "zonal_mean_values": [270.0, 300.0, 270.0], + } + ) + + plot_zonal_mean_hpc( + variable_name="temperature", + session_id=session["session_id"], + dataset_handle=registered["dataset_handle"], + ) + + kwargs = mock_run.call_args.kwargs + assert kwargs["path_hint"] == grid_file + + def test_plot_mesh_hpc_handle_without_session_raises(self): + """dataset_handle without session_id is a clear ValueError.""" + with pytest.raises(ValueError, match="session_id is required"): + plot_mesh_hpc(dataset_handle="some-handle") + + def test_plot_mesh_hpc_no_path_no_handle_raises(self): + """At least one of grid_path or dataset_handle must be provided.""" + with pytest.raises(ValueError, match="grid_path is required"): + plot_mesh_hpc()