diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index f9451b1..6582654 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -203,7 +203,7 @@ jobs: URL: ${{ steps.deploy.outputs.url }} # Defensive lower bound; the load-bearing assertions below are # the explicit name-absence and name-presence checks. - EXPECTED_MIN_TOOLS: '11' + EXPECTED_MIN_TOOLS: '9' run: | set -euo pipefail python <<'PY' @@ -222,13 +222,17 @@ jobs: "create_plotly_chart_from_parquet_output_file", "build_hydrofabric_feature_map_config", "query_hydrofabric_parquet_file", + # Query-cluster consolidation: 4 tools collapsed into query_files_by_selector + "query_output_file", + "query_output_file_from_output_selector", + "query_output_files_from_output_selector", + "resolve_output_file", ) - # Names that must exist (renamed-and-reshaped surface from - # the data-only cleanup in v0.2.0, plus multi-file query in v0.4.1). + # Names that must exist on the deployed server. REQUIRED = ( "lookup_hydrofabric_feature", - "query_output_files_from_output_selector", + "query_files_by_selector", ) # Prompts (slash-command templates) that must exist on the diff --git a/CHANGELOG.md b/CHANGELOG.md index 78dfebc..61eccf3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,128 @@ Image tags follow [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +## [0.5.0] - 2026-05-20 + +**Breaking release.** Major reshape of the query/resolve tool cluster and +removal of NetCDF support. External clients hardcoded to the deleted tool +names or NetCDF queries will break at the deploy. See "Migration" below +for the rename map. + +### Breaking changes + +- **4 query/resolve tools deleted** (consolidated into the new + `query_files_by_selector`): + - `query_output_file` (direct s3_url + SQL) + - `query_output_file_from_output_selector` (single-file with silent + `index=0` default — the diagnosed mis-selection source) + - `query_output_files_from_output_selector` (all-files cross-file query) + - `resolve_output_file` (selector → s3_url; redundant since + `list_available_output_files` already populates the `path` field per + entry) + +- **3 slash prompts deleted** alongside their target tools: + - `query_by_url` (targeted `query_output_file`) + - `resolve_file_by_index` (targeted `resolve_output_file`) + - `resolve_file_by_name` (targeted `resolve_output_file`) + +- **NetCDF (`.nc`, `.nc4`) support removed.** This server is parquet-only + as of v0.5.0. NetCDF outputs are still available in S3; users query + them locally with netCDF-aware tooling (xarray, h5netcdf, etc.). + +### Added + +- **New MCP tool `query_files_by_selector`** — the canonical query tool for + NRDS parquet outputs. Queries one OR many files in a single call: + - Omit `file_name` and `index` → query all parquet files for the + selector as a unioned dataset. Result rows carry `filename` + `source_path` + provenance columns (same as the legacy plural tool). + - Set `file_name` → exact-name lookup against the parquet-filtered list. + - Set `index` → 0-based index into the parquet-filtered list. **NetCDF + files do not consume index slots** — `index=N` always refers to the + N-th parquet file, even when NetCDF files exist in the same directory. + - `file_name` and `index` are mutually exclusive (XOR enforced at the + logic layer). + - Selector args (`model`, `date`, `forecast`, `cycle`, `vpu`, `ensemble`) + are required-together via the existing `_require(model, forecast, vpu)` + pattern. + +- **New `_excluded_netcdf_count` field** on result envelopes for + mixed-format selectors. When a selector contains both parquet and + NetCDF files and the caller didn't supply `file_name`/`index`, the + query silently filters to parquet AND surfaces the dropped-NetCDF + count via this optional field (omitted when zero). Preserves R2 ("no + silent default") parity with the deleted `index=0` fall-through. + +- **Two new error envelope classes**: + - `unsupported_format:` — fires when `file_name` points at a `.nc` + or `.nc4` file. Server-side check happens BEFORE any S3 I/O. Envelope + carries `format_detected="netcdf"`, the offending `file_name`, and a + `fix_hint` pointing the LLM at netCDF-aware tooling for local query. + - `no_supported_files:` — fires when the selector resolves to N files + none of which are parquet. Envelope carries `files_found`, + `netcdf_files`, `parquet_files=0`. + +- **Description-contract test** (`test_mcp/test_tool_descriptions.py`) + enforces the lockstep rule. Asserts positive invariants (parquet-only + mention, error class names, provenance columns) AND negative + invariants (no concrete `s3://` URLs, no example filenames, no inline + SQL). + +### Removed + +- Python deps `xarray` and `h5netcdf` — drops ~30MB+ from the Docker + image (xarray alone is large). +- Dead helpers in `nextgen_mcp/utils_rest.py`: + - `_duckdb_query_parquet` (single-file; consumer was `query_output_file`) + - `_detect_output_file_kind` (parquet-vs-netcdf branch) + - `_validate_nrds_output_file_url` (URL guard for arbitrary external input) + - `_normalize_output_file_url` (s3:// → https:// translator) + - `_get_troute_df` (t-route NetCDF crosswalk loader) + - `_duckdb_query_netcdf` (in-memory pandas → DuckDB) +- Dead helper in `nextgen_mcp/_io_config.py`: + - `open_fsspec_file` (fsspec wrapper used only by `_get_troute_df`) + +### Changed + +- `release.yml` smoke-gate: + - `EXPECTED_MIN_TOOLS` lowered 11 → 9 to match the post-merge catalog. + - `REMOVED` tuple extended with the 4 deleted query-cluster tool names. + - `REQUIRED` tuple: `query_output_files_from_output_selector` dropped; + `query_files_by_selector` added. `lookup_hydrofabric_feature` stays. + - `REQUIRED_PROMPTS` unchanged. + +- `plot_timeseries` slash prompt retargeted from + `query_output_file_from_output_selector` to `query_files_by_selector`. + Prompt-surface args dropped `index` — the new tool defaults to "query + all parquet files for the selector"; the SQL `WHERE feature_id = ...` + in the rendered prompt filters to one feature across the full union. + Other prompt-surface args (`variable`, `feature_id`, `model`, `forecast`, + `date`, `cycle`, `vpu`) unchanged. + +- `nextgen_mcp/README.md` tool list updated to reflect the 9-tool surface. + +### Migration + +If you were calling: + +| Old tool | Replacement | +|---|---| +| `query_output_file(s3_url, query)` | No direct replacement. Use `query_files_by_selector` with the selector args; for ad-hoc parquet URLs outside the NRDS selector hierarchy, query locally with DuckDB. | +| `query_output_file_from_output_selector(model, date, …, file_name, index)` | `query_files_by_selector(model, date, …, file_name=…, index=…)`. NetCDF support dropped; pass parquet `file_name` or `index` only. | +| `query_output_files_from_output_selector(model, date, …)` | `query_files_by_selector(model, date, …)` with `file_name`/`index` both omitted. Same UNION-ALL semantics, same provenance columns. | +| `resolve_output_file(model, date, …, file_name=…, index=…)` | `list_available_output_files(model, date, …)` and read the `path` field from each entry. The list response already carries full S3 URLs. | + +NetCDF queries: this server no longer supports them. Download the file +from S3 and query locally with `xarray` / `h5netcdf` / your tool of choice. + +External consumers using `:latest` Docker tag will hit this breaking +change at the next pull. Pin to a specific minor tag (e.g., +`ghcr.io/aquaveo/nrds-mcps:v0.4`) if you need to defer adoption — confirm +tag availability in the GHCR registry before relying on it. + +Plan: `docs/plans/2026-05-20-001-refactor-nrds-mcps-query-consolidation-plan.md` +Brainstorm: `docs/brainstorms/2026-05-20-nrds-mcps-query-cluster-consolidation-requirements.md` + ## [0.4.1] - 2026-05-18 ### Added diff --git a/nextgen_mcp/README.md b/nextgen_mcp/README.md index 038f0d4..61e8b86 100644 --- a/nextgen_mcp/README.md +++ b/nextgen_mcp/README.md @@ -72,15 +72,14 @@ when no credentials are present. ## Tools The MCP tool surface is discovered automatically by clients via `tools/list`. -The current 11 tools cover: +The current 9 tools (v0.5.0+) cover: - Discovery: `list_available_models`, `list_available_dates`, `list_available_forecasts`, `list_available_cycles`, `list_available_vpus`, `list_available_output_files` -- Resolution: `resolve_output_file` -- Query: `query_output_file`, `query_output_file_from_output_selector`, - `query_output_files_from_output_selector` -- Hydrofabric: `lookup_hydrofabric_feature` +- Query: `query_files_by_selector` (parquet-only; `file_name` or `index` + filters to one file, omit both to query all parquet files for the selector) +- Hydrofabric: `lookup_hydrofabric_feature`, `get_hydrofabric_pmtiles_layers` All tools return data only - no Plotly figure JSON or map config blobs. Charts and maps are the host's responsibility. diff --git a/nextgen_mcp/_io_config.py b/nextgen_mcp/_io_config.py index cbbfbea..61a39b3 100644 --- a/nextgen_mcp/_io_config.py +++ b/nextgen_mcp/_io_config.py @@ -1,10 +1,10 @@ """Centralized IO configuration for nrds_mcps - timeouts and helper factories. -All outbound network IO (S3 via fsspec, DuckDB httpfs, NetCDF via xarray) -goes through helpers in this module so a single ``NRDS_HTTP_TIMEOUT_SECONDS`` -env var controls every IO layer's per-request budget. +All outbound network IO (S3 via fsspec, DuckDB httpfs) goes through helpers +in this module so a single ``NRDS_HTTP_TIMEOUT_SECONDS`` env var controls +every IO layer's per-request budget. -Per the 2026-05-10-004 error-handling efficiency plan: +Conventions: - Default 60s, intentionally permissive ("don't break working flows"). - ``<=0`` env values fall back to default with a warning (defends against IaC pipelines inheriting misconfigured parent envs). @@ -119,21 +119,3 @@ def duckdb_connect_with_httpfs(database: str = ":memory:") -> duckdb.DuckDBPyCon return con -def open_fsspec_file(url: str, mode: str = "rb"): - """Open a remote file with the timeout-configured fsspec client. - - Returns the context manager from ``fsspec.open(...)``. Use as: - - with open_fsspec_file(s3_nc_url) as f: - ds = xarray.open_dataset(f, engine="h5netcdf") - - ``fsspec.open()`` returns a lazy ``OpenFile`` wrapper; passing it - directly to ``xarray.open_dataset`` does not work. The ``with`` block - ensures the underlying file handle uses the timeout-configured client. - """ - return fsspec.open( - url, - mode=mode, - anon=True, - config_kwargs=_s3fs_config_kwargs(), - ) diff --git a/nextgen_mcp/_tool_descriptions.py b/nextgen_mcp/_tool_descriptions.py new file mode 100644 index 0000000..91fd511 --- /dev/null +++ b/nextgen_mcp/_tool_descriptions.py @@ -0,0 +1,28 @@ +"""Module-level tool description constants. + +Tool description prose is extracted here so ``test_mcp/test_tool_descriptions.py`` +can lock the content against drift via positive + negative substring assertions. + +Description content must obey: + - Positive: name the load-bearing constraints (parquet-only, the error + classes the LLM may receive, provenance columns). + - Negative: no concrete example values. LLMs copy concrete examples + verbatim into tool calls. No ``s3://`` URLs, no example filenames, + no inline SQL. +""" + +QUERY_FILES_BY_SELECTOR_DESCRIPTION = ( + "Query NRDS parquet outputs by selector. Queries one OR many files " + "in a single call: pass `file_name` or `index` to filter to one file; " + "omit both to query the full selector as a unioned dataset. " + "Parquet only - NetCDF files return an `unsupported_format:` envelope. " + "Result rows always carry `filename` and `source_path` provenance " + "columns so SQL can group or filter by source. " + "Mixed-format selectors (parquet plus NetCDF) silently filter to " + "parquet and surface the exclusion count as `_excluded_netcdf_count` " + "on the result envelope when non-zero. " + "The SQL must be a single read-only SELECT or WITH...SELECT against " + "table `output`. For data extraction, prefer WHERE filtering over " + "LIMIT - LIMIT silently drops rows and breaks ordered time series. " + "Use aggregates (COUNT, SUM, AVG, MAX, MIN) for summary statistics." +) diff --git a/nextgen_mcp/logic.py b/nextgen_mcp/logic.py index 531d5f2..ebd9c32 100644 --- a/nextgen_mcp/logic.py +++ b/nextgen_mcp/logic.py @@ -19,13 +19,11 @@ from typing import Dict, List, Any, Optional from .validation import OutputsFilesQuery from pydantic import ValidationError +from .utils import _require from .utils_rest import ( _extract_yyyymmdd_from_date_folder, _label_from_id, _normalize_date_folder, - _duckdb_query_parquet, - _duckdb_query_netcdf, - _get_troute_df, _duckdb_lookup_hydrofabric_feature, _normalize_record, _get_feature_center, @@ -34,9 +32,6 @@ _success_payload, _error_payload, _list_payload, - _validate_nrds_output_file_url, - _detect_output_file_kind, - _normalize_output_file_url, _classify_io_error, _is_duckdb_programmer_error, _classify_llm_sql_error, @@ -162,7 +157,7 @@ def get_output_file(model, date, forecast, cycle, vpu, file_name=None, index=Non fs = s3_filesystem() files = fs.ls(s3_dir, detail=False) - files = [f for f in files if f.lower().endswith(".parquet") or f.lower().endswith(".nc")] + files = [f for f in files if f.lower().endswith(".parquet")] files = sorted(files) items = [{"name": f.split("/")[-1], "path": _ensure_full_s3_url(f)} for f in files] @@ -361,249 +356,65 @@ def list_available_models() -> Dict: return _error_payload(code, msg, fix_hint=fix_hint, path=s3_url) -def query_output_file(s3_url, query) -> Dict: - """Run a read-only DuckDB query against one NRDS output file in S3 (parquet or netcdf).""" - raw_url = str(s3_url or "").strip() - kind = _detect_output_file_kind(raw_url) - if kind == "parquet": - err = _validate_nrds_output_file_url(BUCKET, raw_url, (".parquet",)) - elif kind == "netcdf": - err = _validate_nrds_output_file_url(BUCKET, raw_url, (".nc", ".nc4")) - else: - err = "s3_url must point to one .parquet, .nc, or .nc4 NRDS output file" - - if err: - return _error_payload( - "validation_error", - err, - file=raw_url, - query=query, - ) - - file_url = _normalize_output_file_url(raw_url) - logger.info("Received query request for %s file: %s with query: %s", kind, file_url, query) - - try: - query = validate_output_sql(query) - except ValueError as e: - logger.error("Invalid SQL query: %s", e) - return _error_payload( - "validation_error", - str(e), - file=file_url, - query=query, - ) - - try: - if kind == "parquet": - df = _duckdb_query_parquet(file_url, query) - else: - initial_df = _get_troute_df(file_url) - logger.info( - "Initial NetCDF DataFrame loaded with %s rows and columns: %s", - len(initial_df), - initial_df.columns.tolist(), - ) - df = _duckdb_query_netcdf(initial_df, query) +_NETCDF_EXTS = (".nc", ".nc4") - if "time" in df.columns: - df["time"] = pd.to_datetime(df["time"], errors="coerce").dt.strftime("%Y-%m-%dT%H:%M:%S.%fZ") - logger.info("Query returned %s rows and columns: %s", len(df), df.columns.tolist()) - return _success_payload( - file=file_url, - file_type=kind, - query=query, - columns=list(df.columns), - rows=int(len(df)), - data=df.to_dict(orient="records"), - ) - - except (duckdb.BinderException, duckdb.ParserException, duckdb.CatalogException) as e: - # LLM-supplied SQL programmer error. Unlike the hardcoded-SQL - # paths (which re-raise via _is_duckdb_programmer_error), these - # errors are RECOVERABLE if the LLM gets a structured envelope - # with the available column list as fix_hint - same shape as the - # input-validation middleware's invalid_args response. Observed - # 2026-05-10: qwen stalled on a BinderException when this re-raised - # as a protocol error; with the structured envelope below the - # LLM has the actual column candidates and can retry in one turn. - code, msg, fix_hint, available_columns = _classify_llm_sql_error( - e, file_url, query - ) - logger.warning( - "LLM SQL error %s on %s file %s: %s", - type(e).__name__, - kind, - file_url, - e, - ) - return _error_payload( - code, - msg, - fix_hint=fix_hint, - file=file_url, - file_type=kind, - query=query, - available_columns=available_columns, - ) - except (OSError, ClientError, duckdb.Error) as e: - if _is_duckdb_programmer_error(e): - raise - code, msg, fix_hint = _classify_io_error(e) - logger.error( - "IO error %s querying %s file %s: %s", - type(e).__name__, - kind, - file_url, - e, - ) - # not_found preserves the empty-result shape that callers expect - if code == "not_found": - return _error_payload( - code, - msg, - fix_hint=fix_hint, - file=file_url, - file_type=kind, - query=query, - columns=[], - rows=0, - data=[], - ) - return _error_payload( - code, - msg, - fix_hint=fix_hint, - file=file_url, - file_type=kind, - query=query, - ) - - -def query_output_file_from_output_selector( +def _resolve_parquet_files_for_query( model, date, forecast, cycle, vpu, - query, ensemble: Optional[str] = None, file_name: Optional[str] = None, - index: Optional[int] = 0, -) -> Dict: - """Resolve an output file by selector and run a raw query against the selected parquet or netcdf file.""" - - logger.info( - "Received request to query output file from selector with " - "model=%s date=%s forecast=%s cycle=%s vpu=%s ensemble=%s file_name=%s index=%s query=%s", - model, - date, - forecast, - cycle, - vpu, - ensemble, - file_name, - index, - query, - ) - - resolved = get_output_file( - model=model, - date=date, - forecast=forecast, - cycle=cycle, - vpu=vpu, - file_name=file_name, - index=None if file_name is not None else (0 if index is None else index), - ensemble=ensemble, - ) - - if not isinstance(resolved, dict): - return _error_payload( - "execution_error", - "Unexpected response while resolving output file.", - ) - - if resolved.get("ok") is False: - return resolved - - selected = resolved.get("selected") - if not selected: - return _error_payload( - "not_found", - "No output file matched the selector.", - dir=resolved.get("dir"), - count=resolved.get("count", 0), - selected=None, - ) - - selected_path = str((selected or {}).get("path") or "").strip() - if not selected_path: - return _error_payload( - "not_found", - "Resolved output file does not include a path.", - dir=resolved.get("dir"), - count=resolved.get("count", 0), - selected=selected, - ) - - query_result = query_output_file( - s3_url=selected_path, - query=query, - ) - - if isinstance(query_result, dict): - query_result.setdefault("dir", resolved.get("dir")) - query_result.setdefault("count", resolved.get("count")) - query_result.setdefault("selected", selected) - - return query_result - + index: Optional[int] = None, +) -> Dict[str, Any]: + """Resolve a list of parquet S3 URLs from selector args. -def query_output_files_from_output_selector( - model, - date, - forecast, - cycle, - vpu, - query, - ensemble: Optional[str] = None, -) -> Dict: - """Run a single DuckDB query across **all parquet** output files for a selector. + Three branches: + - (file_name=None, index=None): list all parquet files for the selector. + - (file_name set, index=None): filter to one file by exact name. + - (file_name=None, index set): filter to one file by 0-based index + into the parquet-only, sorted list. NetCDF files do NOT consume + index slots — `index=N` always refers to the N-th parquet file. - Mirrors :func:`query_output_file_from_output_selector` but skips the - "pick one file" branch. The resulting envelope swaps the singular - ``selected`` for plural ``files`` plus a ``file_count`` shortcut. + The (file_name, index) BOTH-set case is rejected by the caller's + Pydantic model_validator before reaching here. We don't re-check. - NetCDF outputs in the same directory are intentionally ignored — combining - netCDF files requires pandas concat and has no clean DuckDB primitive. - Callers needing single-file netCDF queries should use - :func:`query_output_file_from_output_selector`. + Returns one of: + {"ok": True, "urls": [...], "s3_dir": ..., "excluded_netcdf_count": int} + or an error envelope (ok=False) ready to return to the caller. """ - from .utils_rest import _duckdb_query_parquets - - logger.info( - "Received request to query output files from selector with " - "model=%s date=%s forecast=%s cycle=%s vpu=%s ensemble=%s query=%s", - model, - date, - forecast, - cycle, - vpu, - ensemble, - query, - ) - date_folder = _normalize_date_folder(date) - s3_dir = f"s3://{BUCKET}/{OUTPUTS_DIR}/{model}/{PREFIX_HYDROFABRIC}/{date_folder}/{forecast}/{cycle}" + s3_dir = ( + f"s3://{BUCKET}/{OUTPUTS_DIR}/{model}/{PREFIX_HYDROFABRIC}/" + f"{date_folder}/{forecast}/{cycle}" + ) if forecast == "medium_range": ens = ensemble or "1" s3_dir += f"/{ens}/{vpu}/{NGEN_RUN_PREFIX}" else: s3_dir += f"/{vpu}/{NGEN_RUN_PREFIX}" + # file_name path: check extension locally BEFORE any S3 I/O. + # Cheap short-circuit for the common LLM mistake of pointing at .nc/.nc4. + if file_name is not None: + lower = file_name.lower() + if lower.endswith(_NETCDF_EXTS): + return _error_payload( + "unsupported_format", + "NetCDF (.nc/.nc4) files are not supported by this server.", + fix_hint=( + "This server queries parquet files only. NetCDF outputs " + "are available in S3 - download with netCDF-aware tooling " + "(xarray, h5netcdf) and query locally." + ), + format_detected="netcdf", + file_name=file_name, + ) + + # List S3 for both no-filter and index/file_name paths. try: fs = s3_filesystem() listing = fs.ls(s3_dir, detail=False) @@ -613,56 +424,213 @@ def query_output_files_from_output_selector( "No output files matched the selector.", dir=s3_dir, count=0, - files=[], - file_count=0, - query=query, ) except (OSError, ClientError, duckdb.Error) as e: if _is_duckdb_programmer_error(e): raise code, msg, fix_hint = _classify_io_error(e) logger.error("IO error %s listing %s: %s", type(e).__name__, s3_dir, e) - return _error_payload(code, msg, fix_hint=fix_hint, dir=s3_dir, query=query) + return _error_payload(code, msg, fix_hint=fix_hint, dir=s3_dir) listing_sorted = sorted(listing) items_all = [ {"name": f.split("/")[-1], "path": _ensure_full_s3_url(f)} for f in listing_sorted ] - items = [it for it in items_all if it["name"].lower().endswith(".parquet")] + parquet_items = [ + it for it in items_all if it["name"].lower().endswith(".parquet") + ] + netcdf_items = [ + it for it in items_all + if it["name"].lower().endswith(_NETCDF_EXTS) + ] - if not items: + # no_supported_files: selector resolved to N files, none parquet. + if not parquet_items: return _error_payload( - "not_found", - "No parquet output files matched the selector.", + "no_supported_files", + ( + f"Selector resolved to {len(items_all)} files, none parquet. " + "This server queries parquet only." + ), + fix_hint=( + "Check whether parquet outputs exist for this selector. If " + "only NetCDF is available, download from S3 with netCDF-aware " + "tooling." + ), dir=s3_dir, - count=len(items_all), - files=[], - file_count=0, - query=query, + files_found=len(items_all), + netcdf_files=len(netcdf_items), + parquet_files=0, + ) + + # file_name branch: exact lookup against the parquet-filtered list. + # (We already short-circuited .nc/.nc4 above; getting here means the + # caller wants a parquet file. If the name doesn't match anything, it's + # genuinely not found.) + if file_name is not None: + match = next( + (it for it in parquet_items if it["name"] == file_name), + None, + ) + if match is None: + return _error_payload( + "not_found", + f"file_name not found in selector: {file_name}", + fix_hint=( + "Call list_available_output_files with the same selector " + "to see valid file names." + ), + dir=s3_dir, + file_name=file_name, + files_found=len(items_all), + parquet_files=len(parquet_items), + ) + return { + "ok": True, + "urls": [match["path"]], + "s3_dir": s3_dir, + "excluded_netcdf_count": len(netcdf_items), + } + + # index branch: parquet-only semantics. + if index is not None: + if index >= len(parquet_items): + return _error_payload( + "invalid_args", + f"index {index} out of range; selector has {len(parquet_items)} parquet files.", + fix_hint=( + "Use an index in [0, parquet_files - 1] or call " + "list_available_output_files to see the file list." + ), + dir=s3_dir, + index=index, + files_found=len(parquet_items), + parquet_files=len(parquet_items), + ) + return { + "ok": True, + "urls": [parquet_items[index]["path"]], + "s3_dir": s3_dir, + "excluded_netcdf_count": len(netcdf_items), + } + + # No-filter branch: query all parquet files. + return { + "ok": True, + "urls": [it["path"] for it in parquet_items], + "s3_dir": s3_dir, + "excluded_netcdf_count": len(netcdf_items), + } + + +def query_files_by_selector( + model, + date, + forecast, + cycle, + vpu, + query, + ensemble: Optional[str] = None, + file_name: Optional[str] = None, + index: Optional[int] = None, +) -> Dict[str, Any]: + """Unified query tool for NRDS parquet outputs by selector. + + Single-file filter (via file_name or index) and no-filter (all parquet + files for the selector) share one resolution path; both end up calling + _duckdb_query_parquets with a list of 1+ URLs and a unified result- + envelope shape. + + file_name XOR index: this logic-layer function normalizes/strips + file_name and enforces the XOR so the contract holds even when invoked + outside the MCP tool wrapper. + """ + from .utils_rest import _duckdb_query_parquets + + # Normalize file_name: strip whitespace; treat empty-after-strip as None + # so we behave the same as "no file_name set" rather than a vacuous match. + if isinstance(file_name, str): + file_name = file_name.strip() + if not file_name: + return _error_payload( + "invalid_args", + "file_name must be a non-empty string or omitted.", + fix_hint=( + "Omit file_name to query all parquet files for the " + "selector, or provide an exact filename." + ), + ) + + # Negative-index defense-in-depth (Pydantic ge=0 also catches this). + if index is not None and index < 0: + return _error_payload( + "invalid_args", + f"index must be >= 0; got {index}.", + fix_hint="Use a non-negative index, or omit index to query all files.", ) + # XOR — both file_name and index set is invalid. + if file_name is not None and index is not None: + return _error_payload( + "invalid_args", + "Provide file_name OR index, not both.", + fix_hint=( + "Pick one filter mode: file_name for an exact match, or " + "index for the N-th parquet file (0-based)." + ), + ) + + err = _require(model=model, forecast=forecast, vpu=vpu) + if err: + # _require returns a non-standard {"error": ""} shape; wrap + # into the standard _error_payload envelope so callers get a + # consistent invalid_args: response. + return _error_payload( + "invalid_args", + err["error"], + fix_hint=( + "Call list_available_models / list_available_forecasts / " + "list_available_vpus to discover valid selector values." + ), + ) + + logger.info( + "Received request to query_files_by_selector with " + "model=%s date=%s forecast=%s cycle=%s vpu=%s ensemble=%s " + "file_name=%s index=%s query=%s", + model, date, forecast, cycle, vpu, ensemble, file_name, index, query, + ) + + resolved = _resolve_parquet_files_for_query( + model=model, + date=date, + forecast=forecast, + cycle=cycle, + vpu=vpu, + ensemble=ensemble, + file_name=file_name, + index=index, + ) + + if not resolved.get("ok"): + return resolved + + file_urls = resolved["urls"] + s3_dir = resolved["s3_dir"] + excluded_netcdf_count = resolved["excluded_netcdf_count"] + try: query = validate_output_sql(query) except ValueError as e: - logger.error("Invalid SQL query: %s", e) return _error_payload( "validation_error", str(e), dir=s3_dir, - files=items, - file_count=len(items), + file_count=len(file_urls), query=query, ) - file_urls = [it["path"] for it in items] - logger.info( - "Querying %s parquet files in %s with: %s", - len(file_urls), - s3_dir, - query, - ) - try: df = _duckdb_query_parquets(file_urls, query) @@ -671,15 +639,9 @@ def query_output_files_from_output_selector( "%Y-%m-%dT%H:%M:%S.%fZ" ) - logger.info( - "Query returned %s rows and columns: %s", - len(df), - df.columns.tolist(), - ) - return _success_payload( + payload = _success_payload( dir=s3_dir, - files=items, - file_count=len(items), + file_count=len(file_urls), file_type="parquet", query=query, columns=list(df.columns), @@ -687,29 +649,27 @@ def query_output_files_from_output_selector( data=df.to_dict(orient="records"), ) + # Surface the exclusion count when non-zero so the user/LLM has a + # signal that NetCDF files were dropped. Omit when zero to keep the + # common-case envelope lean. + if excluded_netcdf_count > 0: + payload["_excluded_netcdf_count"] = excluded_netcdf_count + + return payload + except (duckdb.BinderException, duckdb.ParserException, duckdb.CatalogException) as e: - # LLM-supplied SQL programmer error — recoverable. Surface the same - # structured envelope as the singular-file path so the LLM gets - # available_columns + fix_hint and can retry in one turn. - # _classify_llm_sql_error expects a representative file URL it can - # introspect for columns; use the first one. code, msg, fix_hint, available_columns = _classify_llm_sql_error( e, file_urls[0], query ) logger.warning( "LLM SQL error %s across %s parquet files in %s: %s", - type(e).__name__, - len(file_urls), - s3_dir, - e, + type(e).__name__, len(file_urls), s3_dir, e, ) return _error_payload( - code, - msg, + code, msg, fix_hint=fix_hint, dir=s3_dir, - files=items, - file_count=len(items), + file_count=len(file_urls), file_type="parquet", query=query, available_columns=available_columns, @@ -720,32 +680,13 @@ def query_output_files_from_output_selector( code, msg, fix_hint = _classify_io_error(e) logger.error( "IO error %s querying %s parquet files in %s: %s", - type(e).__name__, - len(file_urls), - s3_dir, - e, + type(e).__name__, len(file_urls), s3_dir, e, ) - if code == "not_found": - return _error_payload( - code, - msg, - fix_hint=fix_hint, - dir=s3_dir, - files=items, - file_count=len(items), - file_type="parquet", - query=query, - columns=[], - rows=0, - data=[], - ) return _error_payload( - code, - msg, + code, msg, fix_hint=fix_hint, dir=s3_dir, - files=items, - file_count=len(items), + file_count=len(file_urls), file_type="parquet", query=query, ) diff --git a/nextgen_mcp/middleware/_input_validation_middleware.py b/nextgen_mcp/middleware/_input_validation_middleware.py index b5e00f7..766017d 100644 --- a/nextgen_mcp/middleware/_input_validation_middleware.py +++ b/nextgen_mcp/middleware/_input_validation_middleware.py @@ -5,9 +5,8 @@ Without this middleware, a hallucinated kwarg on an NRDS tool call causes pydantic's TypeAdapter to raise ``ValidationError`` inside ``Tool._run`` (FastMCP 3.2.x), which propagates out as an MCP-protocol -error and produces a server traceback that any MCP-compatible client -(chatbox-core, Claude Desktop, Cursor, Cline, etc.) cannot recover from -cleanly. +error and produces a server traceback that MCP clients cannot recover +from cleanly. With this middleware in the FastMCP server's middleware stack, the same call returns a structured envelope as a normal tool result with diff --git a/nextgen_mcp/prompts.py b/nextgen_mcp/prompts.py index 118c599..3b56afe 100644 --- a/nextgen_mcp/prompts.py +++ b/nextgen_mcp/prompts.py @@ -21,16 +21,15 @@ def plot_timeseries( date: Annotated[str, Field(description=DATE_HINT)], cycle: Annotated[str, Field(description=CYCLE_HINT)], vpu: Annotated[str, Field(description=VPU_HINT)], - index: Annotated[str, Field(description="0-based output index, e.g., 0")], ) -> str: """Plot a NRDS output-file timeseries as a line chart. Renders a natural-language request that drives a downstream - ``query_output_file_from_output_selector`` invocation followed by a - line-chart visualization of the resulting time series. Designed for - the chatbox slash-command surface. + ``query_files_by_selector`` invocation followed by a line-chart + visualization of the resulting time series. Designed for the + chatbox slash-command surface. - All 8 arguments are ``required: true`` with no Python-level + All 7 arguments are ``required: true`` with no Python-level defaults; each carries a ``Field(description=...)`` advertising the valid format or enum (e.g., ``cfe_nom / lstm / routing_only``, ``yyyy-mm-dd``). Calling ``prompts/get(name, {})`` with empty args @@ -44,25 +43,31 @@ def plot_timeseries( ``[bracket]`` tokens for the user to replace. Hints are derived from the validation types on - ``query_output_file_from_output_selector`` and the NRDS Literal - types in ``validation.py`` (``MODELS``, ``FORECASTS``, - ``DATE_PATTERN``). When NRDS adds a new model, forecast, or vpu, - update the description string here in lockstep. - - Argument names ``model``, ``forecast``, ``date``, ``cycle``, ``vpu``, - and ``index`` align with the selector args of - ``query_output_file_from_output_selector``. ``variable`` and - ``feature_id`` are narrative-only - they help the LLM build the - DuckDB ``query`` value but have no first-class counterpart in the - selector tool's schema. + ``query_files_by_selector`` and the NRDS Literal types in + ``validation.py`` (``MODELS``, ``FORECASTS``, ``DATE_PATTERN``). + When NRDS adds a new model, forecast, or vpu, update the + description string here in lockstep. + + Argument names ``model``, ``forecast``, ``date``, ``cycle``, and + ``vpu`` align with the selector args of ``query_files_by_selector``. + ``variable`` and ``feature_id`` are narrative-only - they help the + LLM build the DuckDB ``query`` value but have no first-class + counterpart in the selector tool's schema. + + The rendered prompt instructs the LLM to omit ``file_name`` and + ``index`` so the query unions ALL parquet files for the selector; + ``WHERE feature_id = ...`` in the SQL filters to the single feature + across the full time series. """ return ( f"Retrieve a line chart plotting the {variable} time series " - f"for feature id {feature_id} for output index {index} for the " + f"for feature id {feature_id} for the " f"{forecast} forecast on {model} model and date {date}, " f"cycle {cycle}, and vpu {vpu}. " + f"Use query_files_by_selector with no file_name or index so " + f"all parquet files for the selector are unioned. " f"Use a query like: SELECT time, {variable} FROM output " - f"WHERE feature_id = {feature_id}" + f"WHERE feature_id = {feature_id} ORDER BY time" ) @@ -190,8 +195,7 @@ def list_output_files( # --------------------------------------------------------------------------- -# Query/lookup prompt templates - one per query/lookup tool plus -# a second variant for resolve_output_file's XOR. +# Lookup prompt templates - one per lookup tool. # # Pattern mirrors the discovery prompts above: # - Argument names mirror the underlying tool's argument names exactly. @@ -200,9 +204,11 @@ def list_output_files( # LOCKSTEP RULE: when the tool's description changes, update both the # tool and the @mcp.prompt arg description here. # - Prose is imperative declarative. -# - The two resolve_file_* variants split the file_name XOR index -# constraint so the user picks intent at the slash level - see each -# variant's "do NOT also supply" instruction. +# +# v0.5.0 deletion note: query_by_url, resolve_file_by_index, and +# resolve_file_by_name lived here previously and targeted query_output_file +# / resolve_output_file. Both target tools were deleted alongside the +# query-cluster consolidation; the prompts were removed in lockstep. # --------------------------------------------------------------------------- @@ -226,88 +232,3 @@ def lookup_feature( return f"Look up the hydrofabric feature with id {hydrofabric_id}." -@mcp.prompt -def query_by_url( - s3_url: Annotated[ - str, - Field( - description=( - "Full URL to ONE parquet or netcdf output file " - "(s3://... or https://...)" - ) - ), - ], - query: Annotated[ - str, Field(description="DuckDB SQL query against table output") - ], -) -> str: - """Run a DuckDB SQL query against a single NRDS output file. - - Drives the ``query_output_file`` tool. The arg name ``s3_url`` is - surfaced as-is on the prompt (it mirrors the underlying tool's arg - name per the parity contract); the description is the user-friendly - explanation. - """ - return ( - f"Run the DuckDB SQL query {query} against the output file at " - f"{s3_url}." - ) - - -@mcp.prompt -def resolve_file_by_index( - model: Annotated[str, Field(description=MODEL_HINT)], - date: Annotated[str, Field(description=DATE_HINT)], - forecast: Annotated[ - str, - Field(description=FORECAST_HINT), - ], - cycle: Annotated[str, Field(description=CYCLE_HINT)], - vpu: Annotated[str, Field(description=VPU_HINT)], - index: Annotated[str, Field(description="0-based output index, e.g., 0")], -) -> str: - """Resolve a single output file by index in the sorted output-file list. - - Drives the ``resolve_output_file`` tool. The XOR constraint on - ``resolve_output_file`` (file_name XOR index) is resolved at the - slash level - this variant supplies ``index`` and the LLM should NOT - also supply ``file_name``. ``index`` defaults to 0 on the underlying - tool but is surfaced as required on the prompt so editors are - explicit about routing. - """ - return ( - f"Resolve the output file by index {index} for the {model} model " - f"on {date}, {forecast} forecast, cycle {cycle}, vpu {vpu}. " - f"Do NOT also supply file_name." - ) - - -@mcp.prompt -def resolve_file_by_name( - model: Annotated[str, Field(description=MODEL_HINT)], - date: Annotated[str, Field(description=DATE_HINT)], - forecast: Annotated[ - str, - Field(description=FORECAST_HINT), - ], - cycle: Annotated[str, Field(description=CYCLE_HINT)], - vpu: Annotated[str, Field(description=VPU_HINT)], - file_name: Annotated[ - str, Field(description="Exact filename (e.g. troute_output_...parquet)") - ], -) -> str: - """Resolve a single output file by exact filename. - - Drives the ``resolve_output_file`` tool. The XOR constraint on - ``resolve_output_file`` (file_name XOR index) is resolved at the - slash level - this variant supplies ``file_name`` and the LLM should - NOT also supply ``index``. The underlying tool's ``index`` defaults - to 0 (not None), so explicitly passing both file_name and index would - fail the XOR check; the docstring instruction tells the LLM to omit - index in this variant. - """ - return ( - f"Resolve the output file by exact filename {file_name} for the " - f"{model} model on {date}, {forecast} forecast, cycle {cycle}, " - f"vpu {vpu}. Do NOT also supply index." - ) diff --git a/nextgen_mcp/requirements.lock b/nextgen_mcp/requirements.lock index 61bcdc3..80ccb4c 100644 --- a/nextgen_mcp/requirements.lock +++ b/nextgen_mcp/requirements.lock @@ -30,7 +30,6 @@ frozenlist==1.8.0 fsspec==2026.4.0 griffelib==2.0.2 h11==0.16.0 -h5netcdf==1.8.1 httpcore==1.0.9 httpx==0.28.1 httpx-sse==0.4.3 @@ -94,6 +93,5 @@ uvicorn==0.46.0 watchfiles==1.1.1 websockets==16.0 wrapt==2.1.2 -xarray==2026.4.0 yarl==1.23.0 zipp==3.23.1 diff --git a/nextgen_mcp/requirements.txt b/nextgen_mcp/requirements.txt index 0b50285..6a38b8c 100644 --- a/nextgen_mcp/requirements.txt +++ b/nextgen_mcp/requirements.txt @@ -1,7 +1,6 @@ duckdb fastmcp>=3.2.3,<4 fsspec -h5netcdf numpy<2 pandas pyarrow @@ -9,4 +8,3 @@ pydantic s3fs starlette typing_extensions -xarray diff --git a/nextgen_mcp/tools.py b/nextgen_mcp/tools.py index 6feabca..7c15336 100644 --- a/nextgen_mcp/tools.py +++ b/nextgen_mcp/tools.py @@ -4,6 +4,7 @@ from typing import Optional, Dict, Any from typing_extensions import Annotated from pydantic import Field +from pydantic.functional_validators import BeforeValidator from .utils import ( _prefer_id_objects, _as_id, @@ -15,6 +16,7 @@ _validate_date_bounds, _preview_text, _summarize_tool_result, + _coerce_none_string, ) from .validation import ( DATE_PATTERN, @@ -28,13 +30,11 @@ list_available_cycles, list_available_vpus, list_available_output_files, - get_output_file, - query_output_file, - query_output_file_from_output_selector, - query_output_files_from_output_selector, + query_files_by_selector, lookup_hydrofabric_feature as _lookup_hydrofabric_feature, get_hydrofabric_pmtiles_layers ) +from ._tool_descriptions import QUERY_FILES_BY_SELECTOR_DESCRIPTION from .middleware._input_validation_middleware import InvalidLLMInputError @@ -314,7 +314,15 @@ def list_available_output_files_tool( ), ] = None, ensemble: Annotated[ - Optional[str], Field(description="Optional ensemble member (1 or 16)", pattern=r"^(?:1|16)$") + Optional[str], + BeforeValidator(_coerce_none_string), + Field( + description=( + "Ensemble member for medium_range forecast. Current data uses " + "ensemble 1; defaults to 1 when omitted. Ignored for " + "short_range and analysis_assim_extend (no ensemble dimension)." + ), + ), ] = None, ) -> Dict[str, Any]: err = _require(model=model, forecast=forecast, vpu=vpu) @@ -355,113 +363,22 @@ def list_available_output_files_tool( @mcp.tool( - name="resolve_output_file", - description="Resolve a single output file path for model/date/forecast/cycle/vpu. Provide exactly one of file_name or index.", + name="query_files_by_selector", + description=QUERY_FILES_BY_SELECTOR_DESCRIPTION, ) -def resolve_output_file_tool( - model: Annotated[MODELS, Field(description="Model id - call list_available_models to discover valid values")] = None, - date: Annotated[ - Optional[str], - Field(description="YYYY-MM-DD or YYYY/MM/DD", pattern=DATE_PATTERN), - ] = None, - forecast: Annotated[FORECASTS, Field(description="Forecast id - call list_available_forecasts to discover valid values")] = None, - cycle: Annotated[str, Field(description="Cycle", pattern=r"^(?:[01]\d|2[0-3])$")] = "00", - vpu: Annotated[ - str, - Field( - description="VPU identifier - call list_available_vpus to discover valid values. Accepts formats like '06', 'VPU_06', or '3W'" - ), - ] = None, - ensemble: Annotated[ - Optional[str], - Field(description="Ensemble (medium_range)", pattern=r"^\d+$"), +def query_files_by_selector_tool( + model: Annotated[ + MODELS, + Field(description="Model id - call list_available_models to discover valid values"), ] = None, - file_name: Annotated[ - Optional[str], - Field(description="Exact filename (e.g. troute_output_...parquet)"), - ] = None, - index: Annotated[ - Optional[int], - Field(description="0-based index into sorted file list", ge=0), - ] = 0, -) -> Dict[str, Any]: - err = _require(model=model, forecast=forecast, vpu=vpu) - if err: - return err - LOGGER.info( - "Tool resolve_output_file called model=%s date=%s forecast=%s cycle=%s vpu=%s ensemble=%s file_name=%s index=%s", - model, - date, - _as_id(forecast), - cycle, - _as_id(vpu), - ensemble, - file_name, - index, - ) - if (file_name is None) == (index is None): - LOGGER.warning( - "Invalid resolve_output_file call: exactly one of file_name or index is required" - ) - raise InvalidLLMInputError( - "Provide exactly one of 'file_name' or 'index'." - ) - - end_date = _parse_date_or_today(date, "date") - params: Dict[str, Any] = { - "model": model, - "date": end_date.isoformat(), - "forecast": _as_id(forecast), - "cycle": cycle, - "vpu": _as_id(vpu), - } - if ensemble is not None: - params["ensemble"] = ensemble - if file_name is not None: - params["file_name"] = file_name - if index is not None: - params["index"] = index - - result = get_output_file( - model=params["model"], - date=params["date"], - forecast=params["forecast"], - cycle=params["cycle"], - vpu=params["vpu"], - file_name=params.get("file_name"), - index=params.get("index"), - ensemble=params.get("ensemble"), - ) - LOGGER.info( - "Tool resolve_output_file completed model=%s date=%s forecast=%s cycle=%s vpu=%s", - model, - end_date.isoformat(), - params["forecast"], - cycle, - params["vpu"], - ) - return result - - -@mcp.tool( - name="query_output_file_from_output_selector", - description=( - "Resolve one NRDS output file from model/date/forecast/cycle/vpu and run a read-only " - "DuckDB SQL query against it in one step. " - "Supports parquet (.parquet) and netcdf (.nc, .nc4). " - "Use this when you know model/date/forecast/cycle/vpu instead of a direct s3_url. " - "If file_name is provided it is used; otherwise index is used and defaults to 0 " - "(the first sorted output file). " - "The SQL query must be a single read-only SELECT or WITH...SELECT statement and must read FROM output." - ), -) -def query_output_file_from_output_selector_tool( - model: Annotated[MODELS, Field(description="Model id - call list_available_models to discover valid values")] = None, date: Annotated[ Optional[str], Field(description="YYYY-MM-DD or YYYY/MM/DD", pattern=DATE_PATTERN), ] = None, - forecast: Annotated[FORECASTS, Field(description="Forecast id - call list_available_forecasts to discover valid values")] = None, + forecast: Annotated[ + FORECASTS, + Field(description="Forecast id - call list_available_forecasts to discover valid values"), + ] = None, cycle: Annotated[ str, Field( @@ -472,47 +389,59 @@ def query_output_file_from_output_selector_tool( vpu: Annotated[ str, Field( - description="VPU identifier - call list_available_vpus to discover valid values. Accepts formats like '06', 'VPU_06', or '3W'" + description=( + "VPU identifier - call list_available_vpus to discover valid values. " + "Accepts formats like '06', 'VPU_06', or '3W'" + ), ), ] = None, query: Annotated[ str, Field( description=( - "DuckDB SQL query against table `output`. " - "Single read-only SELECT or WITH...SELECT statement only. Must read FROM output." + "DuckDB SQL query against table `output` (parquet files unioned with " + "filename + source_path provenance columns). Single read-only SELECT " + "or WITH...SELECT statement only. Must read FROM output. Prefer WHERE " + "filtering over LIMIT for data extraction; LIMIT silently drops rows." ), pattern=r"(?is)^\s*(?:WITH\b.*?\bSELECT\b|SELECT\b).*$", ), - ] = "SELECT * FROM output LIMIT 10", + ] = "SELECT filename, COUNT(*) AS rows_per_file FROM output GROUP BY filename ORDER BY filename", ensemble: Annotated[ Optional[str], - Field(description="Optional ensemble member for medium_range.", pattern=r"^\d+$"), + BeforeValidator(_coerce_none_string), + Field( + description=( + "Ensemble member for medium_range forecast. Current data uses " + "ensemble 1; defaults to 1 when omitted. Ignored for " + "short_range and analysis_assim_extend (no ensemble dimension)." + ), + ), ] = None, file_name: Annotated[ Optional[str], Field( description=( - "Exact filename to query. If provided, it is used and index is ignored." - ) + "Optional filter to one file by exact name. Mutually exclusive with " + "index. Omit both to query all parquet files for the selector." + ), + min_length=1, ), ] = None, index: Annotated[ Optional[int], Field( description=( - "0-based index into the sorted output file list. " - "Used only when file_name is not provided. Defaults to 0 (first file)." + "Optional filter to one file by 0-based index into the parquet-only " + "sorted file list. NetCDF files do not consume index slots. Mutually " + "exclusive with file_name." ), ge=0, ), - ] = 0, + ] = None, ) -> Dict[str, Any]: - err = _require(model=model, forecast=forecast, vpu=vpu) - if err: - return err LOGGER.info( - "Tool query_output_file_from_output_selector called model=%s date=%s forecast=%s cycle=%s " + "Tool query_files_by_selector called model=%s date=%s forecast=%s cycle=%s " "vpu=%s ensemble=%s file_name=%s index=%s query_preview=%s", model, date, @@ -526,194 +455,31 @@ def query_output_file_from_output_selector_tool( ) end_date = _parse_date_or_today(date, "date") - params: Dict[str, Any] = { - "model": model, - "date": end_date.isoformat(), - "forecast": _as_id(forecast), - "cycle": cycle, - "vpu": _as_id(vpu), - "query": query, - } - - if ensemble is not None: - params["ensemble"] = ensemble - - if file_name is not None: - params["file_name"] = file_name - else: - params["index"] = 0 if index is None else index - - result = query_output_file_from_output_selector( - model=params["model"], - date=params["date"], - forecast=params["forecast"], - cycle=params["cycle"], - vpu=params["vpu"], - query=params["query"], - ensemble=params.get("ensemble"), - file_name=params.get("file_name"), - index=params.get("index"), + result = query_files_by_selector( + model=model, + date=end_date.isoformat(), + forecast=_as_id(forecast), + cycle=cycle, + vpu=_as_id(vpu), + query=query, + ensemble=ensemble, + file_name=file_name, + index=index, ) LOGGER.info( - "Tool query_output_file_from_output_selector completed model=%s date=%s forecast=%s cycle=%s vpu=%s result=%s", + "Tool query_files_by_selector completed model=%s date=%s forecast=%s " + "cycle=%s vpu=%s result=%s", model, end_date.isoformat(), - params["forecast"], - cycle, - params["vpu"], - _summarize_tool_result(result), - ) - return result - - -@mcp.tool( - name="query_output_files_from_output_selector", - description=( - "Run ONE read-only DuckDB SQL query across ALL parquet output files for a " - "selector (model/date/forecast/cycle/vpu, plus ensemble for medium_range) " - "as a single combined dataset. " - "Use this when the question spans the whole output bundle (aggregations, " - "ranking across files, max/min/avg) — the database does the merge in one " - "S3-streaming pass instead of one call per file. " - "Every result row carries two provenance columns: `filename` is the file " - "basename only (no path or s3:// prefix), and `source_path` is the full " - "S3 URL. Use `filename` when extracting time portions or labels from the " - "name; use `source_path` when full provenance is needed. " - "Parquet only — use `query_output_file_from_output_selector` for single-file " - "queries or for NetCDF outputs. " - "The SQL must be a single read-only SELECT or WITH...SELECT and must read " - "FROM output. " - "For data extraction (e.g. a feature's time series, a single VPU's flow), " - "use a WHERE clause to filter rows by feature_id, time range, or similar — " - "this bounds response size while preserving the full time series for the " - "subset of interest. Avoid LIMIT in extraction queries; LIMIT silently " - "drops rows from the tail and breaks ordered time series. Use LIMIT only " - "when exploring schema or sampling. Use aggregates (COUNT, SUM, AVG, MAX, " - "MIN) when the question is a summary statistic, not a row list." - ), -) -def query_output_files_from_output_selector_tool( - model: Annotated[MODELS, Field(description="Model id - call list_available_models to discover valid values")] = None, - date: Annotated[ - Optional[str], - Field(description="YYYY-MM-DD or YYYY/MM/DD", pattern=DATE_PATTERN), - ] = None, - forecast: Annotated[FORECASTS, Field(description="Forecast id - call list_available_forecasts to discover valid values")] = None, - cycle: Annotated[ - str, - Field( - description="Cycle (00-23)", - pattern=r"^(?:[01]\d|2[0-3])$", - ), - ] = "00", - vpu: Annotated[ - str, - Field( - description="VPU identifier - call list_available_vpus to discover valid values. Accepts formats like '06', 'VPU_06', or '3W'" - ), - ] = None, - query: Annotated[ - str, - Field( - description=( - "DuckDB SQL query against table `output` — the union of all parquet files in " - "the selector. Two provenance columns are added: `filename` (basename only) " - "and `source_path` (full S3 URL). Single read-only SELECT or WITH...SELECT " - "only. Must read FROM output. For data extraction, prefer WHERE filtering " - "over LIMIT — LIMIT silently drops rows and breaks time series." - ), - pattern=r"(?is)^\s*(?:WITH\b.*?\bSELECT\b|SELECT\b).*$", - ), - ] = "SELECT filename, COUNT(*) AS rows_per_file FROM output GROUP BY filename ORDER BY filename", - ensemble: Annotated[ - Optional[str], - Field(description="Optional ensemble member for medium_range.", pattern=r"^\d+$"), - ] = None, -) -> Dict[str, Any]: - err = _require(model=model, forecast=forecast, vpu=vpu) - if err: - return err - LOGGER.info( - "Tool query_output_files_from_output_selector called model=%s date=%s forecast=%s cycle=%s " - "vpu=%s ensemble=%s query_preview=%s", - model, - date, _as_id(forecast), cycle, _as_id(vpu), - ensemble, - _preview_text(query), - ) - - end_date = _parse_date_or_today(date, "date") - params: Dict[str, Any] = { - "model": model, - "date": end_date.isoformat(), - "forecast": _as_id(forecast), - "cycle": cycle, - "vpu": _as_id(vpu), - "query": query, - } - - if ensemble is not None: - params["ensemble"] = ensemble - - result = query_output_files_from_output_selector( - model=params["model"], - date=params["date"], - forecast=params["forecast"], - cycle=params["cycle"], - vpu=params["vpu"], - query=params["query"], - ensemble=params.get("ensemble"), - ) - - LOGGER.info( - "Tool query_output_files_from_output_selector completed model=%s date=%s forecast=%s cycle=%s vpu=%s result=%s", - model, - end_date.isoformat(), - params["forecast"], - cycle, - params["vpu"], _summarize_tool_result(result), ) return result -@mcp.tool( - name="query_output_file", - description=( - "Run a read-only DuckDB SQL query against ONE NRDS output file in S3. " - "Supports parquet (.parquet) and netcdf (.nc, .nc4). " - "The file is exposed as table `output`." - ), -) -def query_output_file_tool( - s3_url: Annotated[ - str, - Field( - description="Full URL to ONE parquet or netcdf output file (s3://... or https://...)", - pattern=r"^(?:https://|s3://).+\.(?:parquet|nc|nc4)$", - ), - ], - query: Annotated[ - str, - Field( - description="DuckDB SQL query against table `output`.", - pattern=r"(?is)^\s*(?:WITH\b.*?\bSELECT\b|SELECT\b).*$", - ), - ], -) -> Dict[str, Any]: - LOGGER.info( - "Tool query_output_file called s3_url=%s query_preview=%s", - s3_url, - _preview_text(query), - ) - result = query_output_file(s3_url=s3_url, query=query) - LOGGER.info("Tool query_output_file completed s3_url=%s", s3_url) - return result - @mcp.tool( name="lookup_hydrofabric_feature", diff --git a/nextgen_mcp/utils.py b/nextgen_mcp/utils.py index a551f12..c048110 100644 --- a/nextgen_mcp/utils.py +++ b/nextgen_mcp/utils.py @@ -217,6 +217,27 @@ def _parse_date_or_today(date_str: Optional[str], field_name: str): return validated +_NULL_LITERALS = frozenset(("none", "null", "nil", "", "undefined", "")) + + +def _coerce_none_string(v): + """Coerce LLM-emitted null-literals to actual ``None``. + + Workshop-class small models sometimes emit string literals like + ``""`` / ``"None"`` / ``"null"`` when they want to pass ``None`` + to an ``Optional[str]`` arg. Pydantic's pattern matcher rejects these, + the validator-envelope middleware surfaces ``invalid_args:``, and the + LLM retries — recovery works but costs a round-trip. + + Apply as a ``BeforeValidator`` to strip common null-literals to actual + ``None`` before the pattern check fires. Recovery via the validator + envelope still handles novel null-literals outside this allowlist. + """ + if isinstance(v, str) and v.strip().lower() in _NULL_LITERALS: + return None + return v + + def _require(**kwargs): """Validate that required parameters are not None. Returns error dict or None.""" missing = [k for k, v in kwargs.items() if v is None] diff --git a/nextgen_mcp/utils_rest.py b/nextgen_mcp/utils_rest.py index 2830a3d..fe29727 100644 --- a/nextgen_mcp/utils_rest.py +++ b/nextgen_mcp/utils_rest.py @@ -5,9 +5,8 @@ import re import pandas as pd import duckdb -import xarray as xr -from ._io_config import HYDROFABRIC_INDEX_URL, duckdb_connect_with_httpfs, open_fsspec_file +from ._io_config import HYDROFABRIC_INDEX_URL, duckdb_connect_with_httpfs # Per-code sanitized message + fix_hint. NEVER use str(exc) directly - @@ -141,10 +140,11 @@ def _is_duckdb_programmer_error(exc: BaseException) -> bool: not normalized to a polite envelope. CRITICAL: this guard applies ONLY to hardcoded-SQL call sites (e.g. _duckdb_lookup_hydrofabric_feature). - For LLM-supplied-SQL call sites (query_output_file's `query` arg), use - ``_classify_llm_sql_error`` instead - the LLM CAN recover from these - if given a structured envelope with the column list as fix_hint, the - same pattern InputValidationEnvelopeMiddleware uses for kwarg errors. + For LLM-supplied-SQL call sites (query_files_by_selector's `query` + arg), use ``_classify_llm_sql_error`` instead - the LLM CAN recover + from these if given a structured envelope with the column list as + fix_hint, the same pattern InputValidationEnvelopeMiddleware uses + for kwarg errors. """ return isinstance( exc, @@ -187,11 +187,11 @@ def _classify_llm_sql_error( The LLM-facing envelope from this classifier is the SQL analogue of InputValidationEnvelopeMiddleware's `invalid_args` envelope: it gives the LLM a structured way to recover in one retry instead of stalling - in a thinking loop (as observed with qwen on 2026-05-10). + in a thinking loop. - Callers must use this AT LLM-supplied-SQL call sites only (currently - ``query_output_file`` in rest.py). For hardcoded-SQL paths, use the - existing ``_is_duckdb_programmer_error`` re-raise guard instead. + Use this AT LLM-supplied-SQL call sites only (``query_files_by_selector``). + For hardcoded-SQL paths, use ``_is_duckdb_programmer_error`` re-raise + instead. """ exc_msg = str(exc) if isinstance(exc, duckdb.BinderException): @@ -246,45 +246,6 @@ def _classify_llm_sql_error( code, msg, fix_hint = _classify_io_error(exc) return code, msg, fix_hint, [] -def _normalize_output_file_url(s3_url: str) -> str: - file_url = str(s3_url or "").strip() - if file_url.startswith("s3://ciroh-community-ngen-datastream"): - file_url = file_url.replace( - "s3://ciroh-community-ngen-datastream", - "https://ciroh-community-ngen-datastream.s3.us-east-1.amazonaws.com", - ) - return file_url - -def _detect_output_file_kind(file_url: str) -> Optional[str]: - lower = str(file_url or "").lower() - if lower.endswith(".parquet"): - return "parquet" - if lower.endswith(".nc") or lower.endswith(".nc4"): - return "netcdf" - return None - -def _validate_nrds_output_file_url(bucket: str, file_url: str, allowed_exts: tuple[str, ...]) -> Optional[str]: - url = str(file_url or "").strip() - if not url: - return "Missing required query param: s3_url" - - lower = url.lower() - - if not lower.endswith(allowed_exts): - return f"s3_url must point to one file ending in {', '.join(allowed_exts)}" - - allowed_prefixes = ( - f"s3://{bucket.lower()}/", - f"https://{bucket.lower()}.s3.us-east-1.amazonaws.com/", - ) - if not any(lower.startswith(prefix) for prefix in allowed_prefixes): - return f"s3_url must point to bucket {bucket}" - - if "/outputs/" not in lower: - return "s3_url must point to an NRDS outputs file under /outputs/" - - return None - def _success_payload(**kwargs) -> Dict[str, Any]: return { "ok": True, @@ -424,37 +385,6 @@ def _duckdb_lookup_hydrofabric_feature( except Exception: pass -def _get_troute_df(s3_nc_url: str) -> pd.DataFrame: - """Load the t-route crosswalk DataFrame. - - Uses ``open_fsspec_file`` so the timeout-configured fsspec client - reaches the underlying h5netcdf transport. ``xarray.open_dataset`` - cannot be called directly on a URL with a custom fsspec config - the - OpenFile context manager handles that. - """ - - with open_fsspec_file(s3_nc_url) as f: - nc_xarray = xr.open_dataset(f, engine="h5netcdf") - nc_df = nc_xarray.to_dataframe() - nc_df = nc_df.reset_index() - - return nc_df - -def _duckdb_query_parquet(file_url: str, query: str) -> pd.DataFrame: - """Execute an arbitrary DuckDB query against a parquet file exposed as temp view `output`.""" - safe_file_url = file_url.replace("'", "''") - - con = duckdb_connect_with_httpfs() - try: - con.execute(f"CREATE OR REPLACE TEMP VIEW output AS SELECT * FROM read_parquet('{safe_file_url}')") - return con.sql(query).df() - finally: - try: - con.close() - except Exception: - pass - - def _duckdb_query_parquets(file_urls: List[str], query: str) -> pd.DataFrame: """Execute an arbitrary DuckDB query across multiple parquet files exposed as one temp view `output`. @@ -492,20 +422,6 @@ def _duckdb_query_parquets(file_urls: List[str], query: str) -> pd.DataFrame: pass -def _duckdb_query_netcdf(df: pd.DataFrame , query: str) -> pd.DataFrame: - """Execute an arbitrary DuckDB query against a netcdf file exposed as temp view `output`.""" - - con = duckdb.connect(database=":memory:") - con.register('tmp_table_nc', df) - try: - con.execute(f"CREATE OR REPLACE TEMP VIEW output AS SELECT * FROM tmp_table_nc") - return con.sql(query).df() - finally: - try: - con.close() - except Exception: - pass - def _normalize_date_yyyymmdd(date_str: str | None) -> str | None: """Normalize a date string to YYYYMMDD. diff --git a/test_mcp/test_exception_handling.py b/test_mcp/test_exception_handling.py index fc867e4..de716d7 100644 --- a/test_mcp/test_exception_handling.py +++ b/test_mcp/test_exception_handling.py @@ -389,29 +389,41 @@ def test_classify_llm_sql_error_falls_back_for_non_programmer_errors(): assert cols == [] -def test_query_output_file_returns_envelope_on_binder_exception(monkeypatch): +def test_query_files_by_selector_returns_envelope_on_binder_exception(monkeypatch): """Integration: when DuckDB raises BinderException against an LLM-supplied query, the tool returns a structured envelope with available_columns + fix_hint instead of letting the exception bubble. - This is the recovery path qwen needed in the 2026-05-10 bug - the LLM - gets the actual column list and can rewrite its query in one retry. + The LLM gets the actual column list and rewrites its query in one retry + rather than stalling on an unrecoverable traceback. """ import duckdb from nextgen_mcp import utils_rest - def _raise_binder(file_url, query): + # Mock S3 listing so the resolver returns a non-empty parquet list. + class _MockFs: + def ls(self, *_args, **_kwargs): + return [ + "ciroh-community-ngen-datastream/outputs/cfe_nom/v2.2_hydrofabric/" + "20260501/short_range/00/06/ngen-run/outputs/troute/file_a.parquet", + ] + monkeypatch.setattr(logic, "s3_filesystem", lambda: _MockFs()) + + def _raise_binder(file_urls, query): raise duckdb.BinderException( 'Referenced column "variable" not found in FROM clause!\n' 'Candidate bindings: "output.feature_id", "output.velocity"' ) - monkeypatch.setattr(utils_rest, "_duckdb_query_parquet", _raise_binder) - monkeypatch.setattr(logic, "_duckdb_query_parquet", _raise_binder) + monkeypatch.setattr(utils_rest, "_duckdb_query_parquets", _raise_binder) - result = logic.query_output_file( - s3_url="s3://ciroh-community-ngen-datastream/outputs/x/y.parquet", + result = logic.query_files_by_selector( + model="cfe_nom", + date="2026-05-01", + forecast="short_range", + cycle="00", + vpu="06", query="SELECT * FROM output WHERE variable = 'velocity'", ) assert result.get("ok") is False diff --git a/test_mcp/test_middleware.py b/test_mcp/test_middleware.py index 5939d8b..2deb54d 100644 --- a/test_mcp/test_middleware.py +++ b/test_mcp/test_middleware.py @@ -154,11 +154,9 @@ def test_pattern_mismatch_fix_hint_includes_pattern_and_field(): """When a tool input fails a Pydantic regex pattern, the envelope's fix_hint must include the expected pattern AND name the field. - Observed 2026-05-10: qwen passed date='20260510' (no separators), - Pydantic raised string_pattern_mismatch on the - ^(?:\\d{4}-\\d{2}-\\d{2}|\\d{4}/\\d{2}/\\d{2})$ regex. The - middleware's old fix_hint just said 'Fix the type / value errors in - details' - the LLM had no clue what pattern to satisfy. + Without the field+pattern in the hint, the LLM sees only a generic + "Fix the type / value errors in details" and has to guess what pattern + to satisfy. """ async def go(): @@ -196,18 +194,12 @@ async def go(): def test_pattern_mismatch_envelope_surfaces_field_description(): - """Bug 2b regression: pattern-mismatch envelopes include the field's - Pydantic ``Field(description=...)`` text alongside the regex. - - Observed 2026-05-18 against the deployed server: an LLM passed - date='ngen.20250929' (extracted from an S3 path) to - query_output_files_from_output_selector. The fix_hint at the time - was "date must match pattern '^(?:\\\\d{4}-\\\\d{2}-\\\\d{2}|\\\\d{4}/\\\\d{2}/\\\\d{2})$'." - — actionable but required the LLM to mentally parse the regex. - - The Pydantic Field for the same arg already carries - ``description="YYYY-MM-DD or YYYY/MM/DD"``. The middleware now - surfaces that description into both: + """Pattern-mismatch envelopes include the field's Pydantic + ``Field(description=...)`` text alongside the regex. + + The Pydantic Field for the date arg carries + ``description="YYYY-MM-DD or YYYY/MM/DD"``. The middleware surfaces + that description into both: - the ``details`` entry (so structured-data consumers see it), and - the ``fix_hint`` prose (so LLMs reading the natural-language instruction get the format hint inline with the regex). @@ -216,7 +208,7 @@ def test_pattern_mismatch_envelope_surfaces_field_description(): async def go(): async with Client(mcp) as c: return await c.call_tool( - "query_output_files_from_output_selector", + "query_files_by_selector", { "model": "cfe_nom", "forecast": "medium_range", @@ -258,11 +250,10 @@ async def go(): def test_date_out_of_bounds_returns_envelope_not_raise(): """ValueError from _validate_date_bounds in a tool body becomes an envelope. - Observed 2026-05-10: passing date=2023-10-01 to list_available_forecasts - raised ValueError, which FastMCP wrapped in ToolError, which bubbled up - as an MCP protocol error. The LLM saw an unrecoverable traceback. Now - the middleware catches the ToolError-wrapped ValueError and returns an - `invalid_args:` envelope carrying the original prescriptive message. + Without middleware conversion, the ValueError → FastMCP ToolError → + MCP-protocol error path produces an unrecoverable traceback for the LLM. + The middleware catches the ToolError-wrapped ValueError and returns an + ``invalid_args:`` envelope carrying the original prescriptive message. """ async def go(): @@ -283,38 +274,46 @@ async def go(): def test_xor_violation_returns_envelope_not_raise(): - """resolve_output_file file_name XOR index ValueError becomes an envelope. - - Tool body raises InvalidLLMInputError("Provide exactly one of - 'file_name' or 'index'.") when both or neither are supplied. - Middleware converts to invalid_args envelope so the LLM can fix the - call. + """query_files_by_selector file_name XOR index returns an envelope. + + Originally exercised resolve_output_file's tool-body raise of + InvalidLLMInputError; that tool was deleted in v0.5.0. The new + query_files_by_selector enforces file_name XOR index directly via + the logic-layer function (returns _error_payload rather than + raising), so the middleware-convert-from-raise path is no longer + exercised by a public tool. The end-state behavior — caller gets + an invalid_args envelope naming both fields — is unchanged. + + See ``test_query_files_by_selector.py::test_both_file_name_and_index_rejected`` + for the same assertion at the logic-function entry point. """ async def go(): async with Client(mcp) as c: # Pass both file_name AND index - triggers the XOR check. return await c.call_tool( - "resolve_output_file", + "query_files_by_selector", { "model": "cfe_nom", "forecast": "short_range", "vpu": "06", "file_name": "fake.parquet", "index": 0, + "query": "SELECT * FROM output", }, ) result = _run(go()) payload = result.structured_content assert isinstance(payload, dict) - assert payload.get("error", "").startswith("invalid_args:") - # Both field names appear so the LLM knows exactly which two are in - # conflict, and fix_hint is non-empty with the prescriptive message. - assert "file_name" in payload["error"] - assert "index" in payload["error"] + error_obj = payload.get("error") + # _error_payload shape: error is a dict with code + message + assert isinstance(error_obj, dict), f"expected dict error envelope; got {payload!r}" + assert error_obj.get("code") == "invalid_args" + assert "file_name" in error_obj.get("message", "") + assert "index" in error_obj.get("message", "") fix_hint = payload.get("fix_hint") or "" - assert "exactly one" in fix_hint + assert "file_name" in fix_hint.lower() or "pick one" in fix_hint.lower() def test_incidental_value_error_is_not_enveloped(monkeypatch): diff --git a/test_mcp/test_prompts.py b/test_mcp/test_prompts.py index 780e645..da753fc 100644 --- a/test_mcp/test_prompts.py +++ b/test_mcp/test_prompts.py @@ -7,7 +7,7 @@ v1 ships a single prompt - ``plot_timeseries`` - driving the timeseries-chart workflow against -``query_output_file_from_output_selector``. These tests lock the prompt +``query_files_by_selector``. These tests lock the prompt shape, the placeholder-default convention (K4), substitution semantics, the argument-name parity contract with the underlying selector tool, the intentionally narrative-only ``variable`` / ``feature_id`` args, @@ -61,14 +61,13 @@ def _concat_text(messages) -> str: "date", "cycle", "vpu", - "index", ) # Hint-bearing argument descriptions. Each description is the # user-facing format hint advertised via `Field(description=...)` # on the @mcp.prompt arg - derived from the NRDS validation types # (`MODELS`, `FORECASTS`, `DATE_PATTERN` in `nextgen_mcp/validations.py` -# / `nextgen_mcp/utils.py`) and the `query_output_file_from_output_selector` +# / `nextgen_mcp/utils.py`) and the `query_files_by_selector` # field descriptions. When NRDS adds a new model/forecast/vpu, update # both the `Field(description=...)` in `mcp_server.py` and this dict # in lockstep. @@ -85,7 +84,6 @@ def _concat_text(messages) -> str: "date": "yyyy-mm-dd", "cycle": "00-23, e.g., 00", "vpu": "06, VPU_06, or 3W", - "index": "0-based output index, e.g., 0", } @@ -103,8 +101,10 @@ def _strip_fastmcp_schema_note(desc: str) -> str: return "" return desc.split("\n\nProvide as a JSON string")[0].strip() -# Args shared with query_output_file_from_output_selector (lock parity). -OVERLAPPING_ARG_NAMES = ("model", "date", "forecast", "cycle", "vpu", "index") +# Args shared with query_files_by_selector (lock parity). `index` is +# intentionally absent — plot_timeseries always queries the full +# selector and filters via WHERE feature_id in the SQL. +OVERLAPPING_ARG_NAMES = ("model", "date", "forecast", "cycle", "vpu") # Args intentionally narrative-only - must NOT appear in the selector # tool's schema. Locks the partial-alignment design. @@ -258,8 +258,8 @@ async def go(): f"substituted; got: {text!r}" ) - # The remaining 6 synthesized hint brackets survive. - for name in ("model", "forecast", "date", "cycle", "vpu", "index"): + # The remaining 5 synthesized hint brackets survive. + for name in ("model", "forecast", "date", "cycle", "vpu"): hint_bracket = f"[{PLOT_TIMESERIES_DESCRIPTIONS[name]}]" assert hint_bracket in text, ( f"expected unsubstituted hint bracket {hint_bracket!r} for " @@ -271,7 +271,6 @@ async def go(): # Small-model phrasing - plot_timeseries prose must give an unambiguous # SQL hint so small Ollama models (qwen, gemma) don't hallucinate a # column named "variable" from the phrase "for variable {variable}". -# Bug observed 2026-05-10 on qwen running the full template. # --------------------------------------------------------------------------- @@ -346,7 +345,7 @@ async def go(): # --------------------------------------------------------------------------- -# Argument-name parity with query_output_file_from_output_selector +# Argument-name parity with query_files_by_selector # --------------------------------------------------------------------------- @@ -359,11 +358,11 @@ async def go(): tools = _run(go()) selector = next( - (t for t in tools if t.name == "query_output_file_from_output_selector"), + (t for t in tools if t.name == "query_files_by_selector"), None, ) assert selector is not None, ( - "query_output_file_from_output_selector missing from tools/list - " + "query_files_by_selector missing from tools/list - " "the parity contract cannot be evaluated" ) schema = getattr(selector, "inputSchema", None) or {} @@ -384,7 +383,7 @@ async def go(): def test_overlapping_arg_names_present_on_both_surfaces(arg_name): """Each of the 6 overlapping arg names exists on ``plot_timeseries`` AND on - ``query_output_file_from_output_selector``'s schema. Locks the + ``query_files_by_selector``'s schema. Locks the contract one arg at a time so a future rename trips a precise test. """ prompt_args = _plot_timeseries_arg_names() @@ -394,7 +393,7 @@ def test_overlapping_arg_names_present_on_both_surfaces(arg_name): f"{arg_name!r} expected on plot_timeseries; got {prompt_args}" ) assert arg_name in selector_args, ( - f"{arg_name!r} expected on query_output_file_from_output_selector; " + f"{arg_name!r} expected on query_files_by_selector; " f"got {selector_args}" ) @@ -404,7 +403,7 @@ def test_narrative_only_args_present_on_prompt_absent_on_selector(arg_name): """``variable`` and ``feature_id`` are intentionally narrative-only: present on ``plot_timeseries`` (they help the LLM build the SQL ``query`` value), absent from - ``query_output_file_from_output_selector``'s schema. Locks the + ``query_files_by_selector``'s schema. Locks the intentional partial-alignment so a future refactor doesn't silently drop the narrative args or accidentally promote them. """ @@ -417,7 +416,7 @@ def test_narrative_only_args_present_on_prompt_absent_on_selector(arg_name): ) assert arg_name not in selector_args, ( f"{arg_name!r} unexpectedly present on " - f"query_output_file_from_output_selector - narrative-only args " + f"query_files_by_selector - narrative-only args " f"must not be promoted to selector tool args without review" ) @@ -687,7 +686,7 @@ def test_discovery_prompt_arg_name_parity_with_underlying_tool( ): """Each prompt argument name exists on the underlying list_available_* tool's input schema. Catches arg-name drift between prompt and tool - - the #1 risk in this plan (per feedback_input_output_name_alignment.md). + a recurring risk class for slash-prompt surfaces. """ tool_name = DISCOVERY_PROMPT_TO_TOOL[prompt_name] tool_args = _tool_schema_properties(tool_name) @@ -698,13 +697,12 @@ def test_discovery_prompt_arg_name_parity_with_underlying_tool( # --------------------------------------------------------------------------- -# Query/lookup prompts (Phase 2b) - lookup_feature, query_by_url, -# resolve_file_by_index, resolve_file_by_name +# Query/lookup prompts (Phase 2b) - lookup_feature # -# These are query/lookup-archetype prompts (one per query/lookup tool plus -# the XOR-driven extra variant for resolve_output_file). Tests mirror the -# Phase 2a discovery harness shape: same five-test pattern parametrized -# over the prompt set, plus a parallel arg-name parity block. +# Originally this batch covered lookup_feature, query_by_url, +# resolve_file_by_index, and resolve_file_by_name. The latter three were +# deleted in v0.5.0 alongside the query-cluster consolidation (their +# target tools were removed) — only lookup_feature remains. # --------------------------------------------------------------------------- @@ -716,44 +714,14 @@ def test_discovery_prompt_arg_name_parity_with_underlying_tool( "hydrofabric_id": ( "Hydrofabric identifier to search in columns id and divide_id" ), - "s3_url": ( - "Full URL to ONE parquet or netcdf output file " - "(s3://... or https://...)" - ), - "query": "DuckDB SQL query against table output", - "index": "0-based output index, e.g., 0", - "file_name": "Exact filename (e.g. troute_output_...parquet)", } QUERY_LOOKUP_PROMPTS = { "lookup_feature": ("hydrofabric_id",), - "query_by_url": ("s3_url", "query"), - "resolve_file_by_index": ( - "model", - "date", - "forecast", - "cycle", - "vpu", - "index", - ), - "resolve_file_by_name": ( - "model", - "date", - "forecast", - "cycle", - "vpu", - "file_name", - ), } -# Both resolve_file_* variants target the same underlying tool - -# resolve_output_file's input schema contains both file_name and index, -# so the parity test passes for either variant. QUERY_LOOKUP_PROMPT_TO_TOOL = { "lookup_feature": "lookup_hydrofabric_feature", - "query_by_url": "query_output_file", - "resolve_file_by_index": "resolve_output_file", - "resolve_file_by_name": "resolve_output_file", } @@ -884,11 +852,7 @@ def test_query_lookup_prompt_arg_name_parity_with_underlying_tool( ): """Each prompt argument name exists on the underlying query/lookup tool's input schema. Catches arg-name drift between prompt and tool - - the #1 risk in this plan (per feedback_input_output_name_alignment.md). - - Both resolve_file_by_index and resolve_file_by_name target - resolve_output_file; index and file_name both exist on its schema, so - parity holds for either variant. + a recurring risk class for slash-prompt surfaces. """ tool_name = QUERY_LOOKUP_PROMPT_TO_TOOL[prompt_name] tool_args = _tool_schema_properties(tool_name) diff --git a/test_mcp/test_query_files_by_selector.py b/test_mcp/test_query_files_by_selector.py new file mode 100644 index 0000000..6f4259b --- /dev/null +++ b/test_mcp/test_query_files_by_selector.py @@ -0,0 +1,461 @@ +"""Tests for ``query_files_by_selector`` — the unified parquet-query MCP tool +introduced in v0.5.0. + +Replaces the legacy 4-tool cluster (``query_output_file``, +``query_output_file_from_output_selector``, +``query_output_files_from_output_selector``, ``resolve_output_file``) +with one tool that handles single-file filter (via ``file_name`` or +``index``) AND no-filter ("query all parquet files for the selector") in +one signature. + +Test coverage spans: +- Happy paths: no-filter, file_name-set, index-set +- XOR validator: both file_name + index set +- Pydantic edge cases: empty string, whitespace, negative index, oob index +- unsupported_format: envelope (file_name points at .nc/.nc4) +- no_supported_files: envelope (selector resolves to only NetCDF files) +- Mixed-format selector: surfaces ``_excluded_netcdf_count`` +- Index parquet-only semantics: NetCDF files do not consume index slots + +Following the monkeypatch-the-helpers pattern from +``test_query_output_files.py`` — we never hit live S3 or DuckDB. +""" + +from __future__ import annotations + +from typing import Any, List + +import pandas as pd +import pytest +from pydantic import ValidationError + +from nextgen_mcp import logic + + +SELECTOR = { + "model": "cfe_nom", + "date": "2026-05-01", + "forecast": "short_range", + "cycle": "00", + "vpu": "06", +} + + +class _MockFs: + """Minimal fsspec filesystem stand-in returning a fixed listing.""" + + def __init__(self, listing: List[str]): + self._listing = listing + + def ls(self, *_args: Any, **_kwargs: Any) -> List[str]: + return list(self._listing) + + +def _install_fs(monkeypatch: pytest.MonkeyPatch, listing: List[str]) -> None: + monkeypatch.setattr(logic, "s3_filesystem", lambda: _MockFs(listing)) + + +def _install_fake_parquets_query( + monkeypatch: pytest.MonkeyPatch, df_factory +) -> dict: + """Patch ``_duckdb_query_parquets`` and capture the URLs + query passed.""" + captured: dict = {} + + def _fake(file_urls: List[str], query: str) -> pd.DataFrame: + captured["urls"] = list(file_urls) + captured["query"] = query + return df_factory(file_urls, query) + + from nextgen_mcp import utils_rest + monkeypatch.setattr(utils_rest, "_duckdb_query_parquets", _fake) + return captured + + +_S3_DIR = ( + "ciroh-community-ngen-datastream/outputs/cfe_nom/v2.2_hydrofabric/" + "20260501/short_range/00/06/ngen-run/outputs/troute" +) + + +def _full(name: str) -> str: + return f"{_S3_DIR}/{name}" + + +# --------------------------------------------------------------------------- +# Happy paths +# --------------------------------------------------------------------------- + + +def test_no_filter_returns_all_parquet_files(monkeypatch): + listing = [_full(n) for n in ("a.parquet", "b.parquet", "c.parquet")] + _install_fs(monkeypatch, listing) + captured = _install_fake_parquets_query( + monkeypatch, + lambda urls, q: pd.DataFrame( + [{"filename": f"file_{i}.parquet", "feature_id": i} for i in range(len(urls))] + ), + ) + + result = logic.query_files_by_selector( + **SELECTOR, + query="SELECT * FROM output", + ) + + assert result.get("ok") is True + assert len(captured["urls"]) == 3 + assert result["file_count"] == 3 + assert result["rows"] == 3 + # No exclusion field when no NetCDF was filtered out + assert "_excluded_netcdf_count" not in result + + +def test_file_name_filter_returns_single_file(monkeypatch): + listing = [_full(n) for n in ("a.parquet", "b.parquet", "c.parquet")] + _install_fs(monkeypatch, listing) + captured = _install_fake_parquets_query( + monkeypatch, + lambda urls, q: pd.DataFrame([{"filename": "b.parquet", "feature_id": 42}]), + ) + + result = logic.query_files_by_selector( + **SELECTOR, + query="SELECT * FROM output", + file_name="b.parquet", + ) + + assert result.get("ok") is True + assert len(captured["urls"]) == 1 + assert captured["urls"][0].endswith("b.parquet") + assert result["file_count"] == 1 + + +def test_index_filter_returns_single_file_at_index(monkeypatch): + listing = [_full(n) for n in ("a.parquet", "b.parquet", "c.parquet")] + _install_fs(monkeypatch, listing) + captured = _install_fake_parquets_query( + monkeypatch, + lambda urls, q: pd.DataFrame([{"filename": "b.parquet", "feature_id": 99}]), + ) + + result = logic.query_files_by_selector( + **SELECTOR, + query="SELECT * FROM output", + index=1, + ) + + assert result.get("ok") is True + assert len(captured["urls"]) == 1 + assert captured["urls"][0].endswith("b.parquet") + + +# --------------------------------------------------------------------------- +# XOR + Pydantic edge cases +# --------------------------------------------------------------------------- + + +def test_both_file_name_and_index_rejected(monkeypatch): + listing = [_full("a.parquet")] + _install_fs(monkeypatch, listing) + + result = logic.query_files_by_selector( + **SELECTOR, + query="SELECT * FROM output", + file_name="a.parquet", + index=0, + ) + + assert result.get("ok") is False + assert "invalid_args" in result["error"]["code"] + + +def test_empty_string_file_name_rejected(monkeypatch): + listing = [_full("a.parquet")] + _install_fs(monkeypatch, listing) + + result = logic.query_files_by_selector( + **SELECTOR, + query="SELECT * FROM output", + file_name="", + ) + + assert result.get("ok") is False + assert "invalid_args" in result["error"]["code"] + + +def test_whitespace_file_name_rejected(monkeypatch): + listing = [_full("a.parquet")] + _install_fs(monkeypatch, listing) + + result = logic.query_files_by_selector( + **SELECTOR, + query="SELECT * FROM output", + file_name=" ", + ) + + assert result.get("ok") is False + assert "invalid_args" in result["error"]["code"] + + +def test_negative_index_rejected(monkeypatch): + listing = [_full("a.parquet")] + _install_fs(monkeypatch, listing) + + result = logic.query_files_by_selector( + **SELECTOR, + query="SELECT * FROM output", + index=-1, + ) + + assert result.get("ok") is False + assert "invalid_args" in result["error"]["code"] + + +def test_out_of_range_index_returns_invalid_args(monkeypatch): + listing = [_full(n) for n in ("a.parquet", "b.parquet")] + _install_fs(monkeypatch, listing) + + result = logic.query_files_by_selector( + **SELECTOR, + query="SELECT * FROM output", + index=999999, + ) + + assert result.get("ok") is False + assert "invalid_args" in result["error"]["code"] + # Surface the actual count so the LLM can retry with a valid index + assert result.get("files_found") == 2 + + +def test_file_name_not_found_returns_not_found(monkeypatch): + listing = [_full(n) for n in ("a.parquet", "b.parquet")] + _install_fs(monkeypatch, listing) + + result = logic.query_files_by_selector( + **SELECTOR, + query="SELECT * FROM output", + file_name="missing.parquet", + ) + + assert result.get("ok") is False + assert "not_found" in result["error"]["code"] + + +# --------------------------------------------------------------------------- +# unsupported_format: envelope +# --------------------------------------------------------------------------- + + +def test_file_name_pointing_at_netcdf_returns_unsupported_format(monkeypatch): + """No S3 I/O required - cheap local extension check on file_name.""" + # We intentionally don't install fs because the check must short-circuit + # before fs.ls is called. If implementation tries to list, the next + # call would raise AttributeError on the missing patch. + + result = logic.query_files_by_selector( + **SELECTOR, + query="SELECT * FROM output", + file_name="something.nc", + ) + + assert result.get("ok") is False + assert result["error"]["code"] == "unsupported_format" + assert result.get("format_detected") == "netcdf" + assert result.get("file_name") == "something.nc" + assert "fix_hint" in result + + +def test_file_name_pointing_at_nc4_returns_unsupported_format(monkeypatch): + result = logic.query_files_by_selector( + **SELECTOR, + query="SELECT * FROM output", + file_name="something.nc4", + ) + + assert result.get("ok") is False + assert result["error"]["code"] == "unsupported_format" + assert result.get("format_detected") == "netcdf" + + +# --------------------------------------------------------------------------- +# no_supported_files: envelope +# --------------------------------------------------------------------------- + + +def test_selector_only_netcdf_returns_no_supported_files(monkeypatch): + listing = [_full(n) for n in ("a.nc", "b.nc", "c.nc4")] + _install_fs(monkeypatch, listing) + + result = logic.query_files_by_selector( + **SELECTOR, + query="SELECT * FROM output", + ) + + assert result.get("ok") is False + assert result["error"]["code"] == "no_supported_files" + assert result["files_found"] == 3 + assert result["netcdf_files"] == 3 + assert result["parquet_files"] == 0 + + +# --------------------------------------------------------------------------- +# Mixed-format selector +# --------------------------------------------------------------------------- + + +def test_mixed_format_selector_surfaces_excluded_netcdf_count(monkeypatch): + listing = [ + _full("a.parquet"), + _full("b.nc"), + _full("c.parquet"), + _full("d.parquet"), + _full("e.nc4"), + ] + _install_fs(monkeypatch, listing) + _install_fake_parquets_query( + monkeypatch, + lambda urls, q: pd.DataFrame( + [{"filename": "x.parquet", "feature_id": i} for i in range(len(urls))] + ), + ) + + result = logic.query_files_by_selector( + **SELECTOR, + query="SELECT * FROM output", + ) + + assert result.get("ok") is True + assert result["file_count"] == 3 # only parquet + assert result["_excluded_netcdf_count"] == 2 + + +# --------------------------------------------------------------------------- +# Index parquet-only semantics — NetCDF files do not consume index slots +# --------------------------------------------------------------------------- + + +def test_index_skips_netcdf_files_in_mixed_listing(monkeypatch): + """When mixed [a.parquet, b.nc, c.parquet], index=1 → c.parquet, NOT b.nc.""" + listing = [_full(n) for n in ("a.parquet", "b.nc", "c.parquet")] + _install_fs(monkeypatch, listing) + captured = _install_fake_parquets_query( + monkeypatch, + lambda urls, q: pd.DataFrame([{"filename": "c.parquet", "feature_id": 1}]), + ) + + result = logic.query_files_by_selector( + **SELECTOR, + query="SELECT * FROM output", + index=1, + ) + + assert result.get("ok") is True + assert len(captured["urls"]) == 1 + # CRITICAL: index=1 resolves to the SECOND PARQUET file (c.parquet), + # not the second file in the raw sorted listing (b.nc). + assert captured["urls"][0].endswith("c.parquet") + + +def test_index_zero_against_all_netcdf_returns_no_supported_files(monkeypatch): + """index=0 against an all-NetCDF directory: no parquet exists at any + index, so no_supported_files: (not unsupported_format:) — the request + shape was valid; the selector yielded no queryable files.""" + listing = [_full(n) for n in ("a.nc", "b.nc")] + _install_fs(monkeypatch, listing) + + result = logic.query_files_by_selector( + **SELECTOR, + query="SELECT * FROM output", + index=0, + ) + + assert result.get("ok") is False + assert result["error"]["code"] == "no_supported_files" + + +# --------------------------------------------------------------------------- +# Missing required selector args +# --------------------------------------------------------------------------- + + +def test_missing_model_returns_invalid_args(monkeypatch): + selector_no_model = dict(SELECTOR) + selector_no_model["model"] = None + + result = logic.query_files_by_selector( + **selector_no_model, + query="SELECT * FROM output", + ) + + assert result.get("ok") is False + # _require pattern: existing error class + assert "invalid_args" in result["error"]["code"] or "validation" in result["error"]["code"] + + +# --------------------------------------------------------------------------- +# Null-literal coercion for ensemble arg +# +# Small models sometimes emit string literals like "" / "None" / "null" +# for an Optional[str] arg. The BeforeValidator on the @mcp.tool wrapper +# coerces these to None before validation. Tested at the helper level +# (cheap unit test) AND via the in-process Client (proves the wrapper is +# wired correctly). +# --------------------------------------------------------------------------- + + +import pytest +from nextgen_mcp.utils import _coerce_none_string + + +@pytest.mark.parametrize( + "raw,expected", + [ + ("", None), + ("None", None), + ("null", None), + ("nil", None), + ("undefined", None), + ("NULL", None), + (" None ", None), # whitespace stripped + ("", None), + ("1", "1"), # real ensemble value passes through + ("16", "16"), # historical ensemble value passes through + (None, None), # actual None passes through + (1, 1), # non-string passes through unchanged + ], +) +def test_coerce_none_string(raw, expected): + assert _coerce_none_string(raw) == expected + + +def test_query_files_by_selector_tool_accepts_nil_string_for_ensemble(monkeypatch): + """Integration: ensemble='' (string literal from small model) is + coerced to None at the BeforeValidator stage, BEFORE the pattern check + fires. Tool call succeeds without a validation-envelope round-trip. + """ + import asyncio + from fastmcp import Client + from nextgen_mcp.mcp_server import mcp + + listing = [_full("a.parquet")] + _install_fs(monkeypatch, listing) + _install_fake_parquets_query( + monkeypatch, + lambda urls, q: pd.DataFrame([{"filename": "a.parquet", "feature_id": 1}]), + ) + + async def go(): + async with Client(mcp) as c: + return await c.call_tool( + "query_files_by_selector", + { + **SELECTOR, + "ensemble": "", # the offending small-model literal + "query": "SELECT * FROM output", + }, + ) + + result = asyncio.get_event_loop_policy().new_event_loop().run_until_complete(go()) + payload = result.structured_content + # Tool call should succeed — ensemble='' coerced to None, then + # ignored for short_range, no pattern-mismatch envelope fires. + assert payload.get("ok") is True, f"expected success; got {payload!r}" diff --git a/test_mcp/test_query_output_files.py b/test_mcp/test_query_output_files.py index aeeaf6e..1582914 100644 --- a/test_mcp/test_query_output_files.py +++ b/test_mcp/test_query_output_files.py @@ -1,18 +1,18 @@ -"""Tests for ``query_output_files_from_output_selector`` — the multi-file -parquet-query MCP tool. +"""Tests for the ``_duckdb_query_parquets`` helper and shared payload +summarizer. -The function resolves an S3 directory from (model, date, forecast, cycle, -vpu, ensemble), lists the files there, filters to parquet, and runs a -single DuckDB query across all of them with a ``filename`` provenance -column. These tests pin the four code paths that matter: +The four behavior tests that previously lived here (against the deleted +``query_output_files_from_output_selector``) were superseded in v0.5.0 by: - - happy path (3 parquet files, query returns rows) - - directory contains only ``.nc`` files (not_found with count) - - directory is empty (not_found with count=0) - - DuckDB raises ``BinderException`` (LLM-recoverable invalid_query envelope) + - ``test_query_files_by_selector.py``: end-to-end coverage of the new + unified tool including unsupported_format / no_supported_files / + excluded_netcdf_count envelopes + - ``test_exception_handling.py``: migrated BinderException recovery test + targeting ``query_files_by_selector`` + ``_duckdb_query_parquets`` -Following the same monkeypatch-the-helpers pattern as -``test_exception_handling.py`` — we never hit live S3 or DuckDB. +What remains here are the helper-level live-DuckDB tests against +``_duckdb_query_parquets`` (filename basename, substr-on-filename) and the +``_summarize_tool_result`` helper. """ from __future__ import annotations @@ -50,136 +50,6 @@ def _install_fs(monkeypatch: pytest.MonkeyPatch, listing: List[str]) -> None: monkeypatch.setattr(logic, "s3_filesystem", lambda: _MockFs(listing)) -def test_query_output_files_happy_path(monkeypatch: pytest.MonkeyPatch) -> None: - """Three parquet files resolved; DuckDB returns rows with filename column.""" - listing = [ - "ciroh-community-ngen-datastream/outputs/cfe_nom/v2.2_hydrofabric/20260501/short_range/00/06/ngen-run/outputs/troute/file_a.parquet", - "ciroh-community-ngen-datastream/outputs/cfe_nom/v2.2_hydrofabric/20260501/short_range/00/06/ngen-run/outputs/troute/file_b.parquet", - "ciroh-community-ngen-datastream/outputs/cfe_nom/v2.2_hydrofabric/20260501/short_range/00/06/ngen-run/outputs/troute/file_c.parquet", - ] - _install_fs(monkeypatch, listing) - - captured_urls: dict = {} - - def _fake_query(file_urls: List[str], query: str) -> pd.DataFrame: - captured_urls["urls"] = file_urls - captured_urls["query"] = query - # Simulate one row per file with filename provenance + a real column. - return pd.DataFrame( - [ - {"filename": "s3://bucket/file_a.parquet", "feature_id": 1}, - {"filename": "s3://bucket/file_b.parquet", "feature_id": 2}, - {"filename": "s3://bucket/file_c.parquet", "feature_id": 3}, - ] - ) - - # Patch at the import site inside logic.query_output_files_from_output_selector - # (the function does ``from .utils_rest import _duckdb_query_parquets`` - # at call time, so patching the source module is the reliable seam). - from nextgen_mcp import utils_rest - monkeypatch.setattr(utils_rest, "_duckdb_query_parquets", _fake_query) - - result = logic.query_output_files_from_output_selector( - **SELECTOR, - query="SELECT filename, feature_id FROM output", - ) - - assert result.get("ok") is True, result - assert result["file_count"] == 3 - assert len(result["files"]) == 3 - # File names are extracted from the S3 path basename. - assert {f["name"] for f in result["files"]} == { - "file_a.parquet", - "file_b.parquet", - "file_c.parquet", - } - assert result["file_type"] == "parquet" - assert "filename" in result["columns"] - assert result["rows"] == 3 - assert len(result["data"]) == 3 - # Verify the helper actually received the 3 parquet URLs. - assert len(captured_urls["urls"]) == 3 - assert all(u.endswith(".parquet") for u in captured_urls["urls"]) - - -def test_query_output_files_not_found_only_netcdf(monkeypatch: pytest.MonkeyPatch) -> None: - """Selector resolves to a dir containing only .nc files → not_found. - - The envelope's ``count`` reflects the unfiltered total so the LLM can - tell the directory had files, just not parquet — which is the signal - to redirect to the singular ``query_output_file_from_output_selector``. - """ - listing = [ - "ciroh-community-ngen-datastream/outputs/cfe_nom/v2.2_hydrofabric/20260501/short_range/00/06/ngen-run/outputs/troute/out_a.nc", - "ciroh-community-ngen-datastream/outputs/cfe_nom/v2.2_hydrofabric/20260501/short_range/00/06/ngen-run/outputs/troute/out_b.nc", - ] - _install_fs(monkeypatch, listing) - - result = logic.query_output_files_from_output_selector( - **SELECTOR, - query="SELECT * FROM output", - ) - - assert result.get("ok") is False - assert result["error"]["code"] == "not_found" - assert result["file_count"] == 0 - assert result["files"] == [] - # count reflects unfiltered total — directory had 2 .nc files. - assert result["count"] == 2 - - -def test_query_output_files_not_found_empty_dir(monkeypatch: pytest.MonkeyPatch) -> None: - """Empty selector directory → not_found with count=0.""" - _install_fs(monkeypatch, []) - - result = logic.query_output_files_from_output_selector( - **SELECTOR, - query="SELECT * FROM output", - ) - - assert result.get("ok") is False - assert result["error"]["code"] == "not_found" - assert result["file_count"] == 0 - assert result["count"] == 0 - - -def test_query_output_files_binder_exception_returns_envelope(monkeypatch: pytest.MonkeyPatch) -> None: - """DuckDB BinderException → recoverable invalid_query envelope with fix_hint. - - Same recovery contract as the singular ``query_output_file`` — the LLM - gets ``available_columns`` parsed from DuckDB's "Candidate bindings" - message and ``fix_hint`` text telling it to retry once with a real - column. Pins parity with the singular tool's envelope shape. - """ - listing = [ - "ciroh-community-ngen-datastream/outputs/cfe_nom/v2.2_hydrofabric/20260501/short_range/00/06/ngen-run/outputs/troute/file_a.parquet", - "ciroh-community-ngen-datastream/outputs/cfe_nom/v2.2_hydrofabric/20260501/short_range/00/06/ngen-run/outputs/troute/file_b.parquet", - ] - _install_fs(monkeypatch, listing) - - def _raise_binder(file_urls: List[str], query: str) -> pd.DataFrame: - raise duckdb.BinderException( - 'Referenced column "variable" not found in FROM clause!\n' - 'Candidate bindings: "output.feature_id", "output.velocity"' - ) - - from nextgen_mcp import utils_rest - monkeypatch.setattr(utils_rest, "_duckdb_query_parquets", _raise_binder) - - result = logic.query_output_files_from_output_selector( - **SELECTOR, - query="SELECT * FROM output WHERE variable = 'velocity'", - ) - - assert result.get("ok") is False - assert result["error"]["code"] == "invalid_query" - assert result["available_columns"] == ["feature_id", "velocity"] - assert "feature_id" in result["fix_hint"] - # File listing still surfaced in the envelope so the LLM has context. - assert result["file_count"] == 2 - assert len(result["files"]) == 2 - - def test_duckdb_query_parquets_filename_is_basename(tmp_path) -> None: """Integration: filename column exposed to SQL is basename-only. @@ -226,16 +96,13 @@ def test_duckdb_query_parquets_filename_is_basename(tmp_path) -> None: def test_summarize_tool_result_includes_aggregates_and_sample_row() -> None: - """Bug 1 regression: log summary must show enough to spot-check the - result without dumping the full `data` array. - - Observed 2026-05-18 against the deployed server: the multi-file tool's - completion log line was ``LOGGER.info("...result: %s", result)`` which - repr'd the full envelope including a 240-row ``data`` array — a - ~19 KB single-line log entry per call. The first fix replaced that - with a too-sparse stats line; operators couldn't tell whether the - data looked right. The current contract is: aggregates (file_count, - rows, columns *by name*) + ONE sample row, all under ~500 chars. + """Log summary must show enough to spot-check the result without dumping + the full ``data`` array. + + Contract: aggregates (file_count, rows, columns *by name*) + ONE sample + row, all under ~500 chars. Reprs of full data arrays produce ~19 KB + single-line log entries per call; too-sparse stats lines hide whether + the data looked right. This shape is the middle ground. """ from nextgen_mcp.utils import _summarize_tool_result diff --git a/test_mcp/test_tool_descriptions.py b/test_mcp/test_tool_descriptions.py new file mode 100644 index 0000000..f6e0aaf --- /dev/null +++ b/test_mcp/test_tool_descriptions.py @@ -0,0 +1,66 @@ +"""Lockstep contract tests for tool descriptions. + +Positive assertions: the description must contain load-bearing constraints +(parquet-only, error class names, provenance column names). + +Negative assertions: the description must NOT contain concrete example values. +LLMs copy concrete examples verbatim — any ``s3://`` URL, example filename, +or inline SQL in the description leaks into tool calls. +""" + +from __future__ import annotations + +import re + +from nextgen_mcp._tool_descriptions import QUERY_FILES_BY_SELECTOR_DESCRIPTION + + +def test_query_files_by_selector_description_positive_invariants() -> None: + desc = QUERY_FILES_BY_SELECTOR_DESCRIPTION + lower = desc.lower() + + # Names the parquet-only constraint + assert "parquet only" in lower + + # Names the error class the LLM may receive on NetCDF-target calls + assert "unsupported_format" in desc + + # Names the provenance columns the LLM can reference in SQL + assert "filename" in desc + assert "source_path" in desc + + # Names the file_name / index filter args so the LLM knows the filter shape + assert "file_name" in desc + assert "index" in desc + + # Names _excluded_netcdf_count so the LLM knows to look for it + assert "_excluded_netcdf_count" in desc + + +def test_query_files_by_selector_description_negative_invariants() -> None: + """Description must not embed concrete example values that LLMs would copy.""" + desc = QUERY_FILES_BY_SELECTOR_DESCRIPTION + + # No concrete URLs - LLMs would copy these into tool calls + assert "s3://" not in desc + assert "https://" not in desc + + # No example filenames - "*.parquet" or "troute_output_..." style + assert "troute_output" not in desc + assert "*.parquet" not in desc + + # The phrase "parquet only" is the constraint statement and uses lowercase + # "parquet" deliberately. Bare ".parquet" filename suffixes (with the dot) + # would be example values; assert none appear. + assert ".parquet" not in desc + + # No inline SQL examples - just structural mentions ("SELECT or WITH...SELECT") + # are acceptable, but "SELECT * FROM ..." or specific column names are not + assert "SELECT * FROM" not in desc + assert "WHERE feature_id" not in desc + + # No specific selector example values (model names, vpu IDs, etc.) + # The description should describe the shape, not name specific instances. + assert "cfe_nom" not in desc + assert "short_range" not in desc + assert not re.search(r"\bVPU_\d+", desc)