From 86e3d797f3c24cc2d6f424f767e8480510e431c6 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Fri, 12 Jun 2026 10:27:40 +0100 Subject: [PATCH] Add CDS api tools --- pyproject.toml | 2 + toolsets/cds/pyproject.toml | 19 + toolsets/cds/src/cds/__init__.py | 1 + toolsets/cds/src/cds/client.py | 14 + toolsets/cds/src/cds/settings.py | 21 ++ toolsets/cds/src/cds/tools/__init__.py | 38 ++ toolsets/cds/src/cds/tools/_client.py | 12 + toolsets/cds/src/cds/tools/_errors.py | 79 +++++ toolsets/cds/src/cds/tools/_retry.py | 33 ++ .../cds/src/cds/tools/apply_constraints.py | 50 +++ .../cds/src/cds/tools/check_credentials.py | 26 ++ .../cds/src/cds/tools/get_dataset_schema.py | 116 +++++++ toolsets/cds/src/cds/tools/get_job_status.py | 42 +++ toolsets/cds/src/cds/tools/get_results.py | 56 +++ toolsets/cds/src/cds/tools/list_jobs.py | 64 ++++ toolsets/cds/src/cds/tools/search_datasets.py | 49 +++ toolsets/cds/src/cds/tools/submit_request.py | 57 +++ toolsets/cds/tests/test_cds.py | 326 ++++++++++++++++++ toolsets/cds/toolset.yaml | 6 + uv.lock | 24 ++ 20 files changed, 1035 insertions(+) create mode 100644 toolsets/cds/pyproject.toml create mode 100644 toolsets/cds/src/cds/__init__.py create mode 100644 toolsets/cds/src/cds/client.py create mode 100644 toolsets/cds/src/cds/settings.py create mode 100644 toolsets/cds/src/cds/tools/__init__.py create mode 100644 toolsets/cds/src/cds/tools/_client.py create mode 100644 toolsets/cds/src/cds/tools/_errors.py create mode 100644 toolsets/cds/src/cds/tools/_retry.py create mode 100644 toolsets/cds/src/cds/tools/apply_constraints.py create mode 100644 toolsets/cds/src/cds/tools/check_credentials.py create mode 100644 toolsets/cds/src/cds/tools/get_dataset_schema.py create mode 100644 toolsets/cds/src/cds/tools/get_job_status.py create mode 100644 toolsets/cds/src/cds/tools/get_results.py create mode 100644 toolsets/cds/src/cds/tools/list_jobs.py create mode 100644 toolsets/cds/src/cds/tools/search_datasets.py create mode 100644 toolsets/cds/src/cds/tools/submit_request.py create mode 100644 toolsets/cds/tests/test_cds.py create mode 100644 toolsets/cds/toolset.yaml diff --git a/pyproject.toml b/pyproject.toml index a00c0cc..341fde0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,6 +10,7 @@ dependencies = [ "mcp-agent", "dataset-search", "aoi-generator", + "cds", ] [dependency-groups] @@ -32,6 +33,7 @@ mcp-cli = { workspace = true } mcp-agent = { workspace = true } dataset-search = { workspace = true } aoi-generator = { workspace = true } +cds = { workspace = true } [tool.ruff] line-length = 88 diff --git a/toolsets/cds/pyproject.toml b/toolsets/cds/pyproject.toml new file mode 100644 index 0000000..6ea8f0e --- /dev/null +++ b/toolsets/cds/pyproject.toml @@ -0,0 +1,19 @@ +[project] +name = "cds" +version = "0.1.0" +description = "Copernicus Climate Data Store (CDS) tools: search datasets, inspect schemas, submit and track retrieval jobs." +requires-python = ">=3.12,<3.14" +dependencies = [ + "langchain-core<2.0.0,>=1.4.6", + "mcp-runtime", + "httpx>=0.27", + "pydantic-settings>=2.5", + "tenacity>=9.0", +] + +[tool.uv.sources] +mcp-runtime = { workspace = true } + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" diff --git a/toolsets/cds/src/cds/__init__.py b/toolsets/cds/src/cds/__init__.py new file mode 100644 index 0000000..6d3ba93 --- /dev/null +++ b/toolsets/cds/src/cds/__init__.py @@ -0,0 +1 @@ +"""cds toolset: Copernicus Climate Data Store tools, ported from cds-assistant.""" diff --git a/toolsets/cds/src/cds/client.py b/toolsets/cds/src/cds/client.py new file mode 100644 index 0000000..2c33019 --- /dev/null +++ b/toolsets/cds/src/cds/client.py @@ -0,0 +1,14 @@ +import httpx + +from .settings import settings + + +def make_client() -> httpx.AsyncClient: + return httpx.AsyncClient( + base_url=settings.cds_api_url, + headers={ + "PRIVATE-TOKEN": settings.cds_api_key, + "Content-Type": "application/json", + }, + timeout=httpx.Timeout(connect=10.0, read=120.0, write=10.0, pool=10.0), + ) diff --git a/toolsets/cds/src/cds/settings.py b/toolsets/cds/src/cds/settings.py new file mode 100644 index 0000000..840f117 --- /dev/null +++ b/toolsets/cds/src/cds/settings.py @@ -0,0 +1,21 @@ +from pydantic_settings import BaseSettings, SettingsConfigDict + + +class Settings(BaseSettings): + """CDS API configuration, read from the environment (or a .env in the cwd). + + `cds_api_key` defaults to empty so the toolset imports and serves without + credentials; calls then return structured auth errors until CDS_API_KEY + is provided (in k8s, via the Secret listed in toolset.yaml). + """ + + model_config = SettingsConfigDict( + env_file=".env", env_file_encoding="utf-8", extra="ignore" + ) + + cds_api_key: str = "" + cds_api_url: str = "https://cds.climate.copernicus.eu/api/retrieve/v1" + cds_catalogue_url: str = "https://cds.climate.copernicus.eu/api/catalogue/v1" + + +settings = Settings() diff --git a/toolsets/cds/src/cds/tools/__init__.py b/toolsets/cds/src/cds/tools/__init__.py new file mode 100644 index 0000000..76ef4dd --- /dev/null +++ b/toolsets/cds/src/cds/tools/__init__.py @@ -0,0 +1,38 @@ +"""LangChain tools for the Copernicus Climate Data Store (CDS). + +Covers the full retrieval workflow: search the catalogue, inspect a +dataset's schema and constraints, submit requests, then poll jobs and +fetch download links. +""" + +from .apply_constraints import apply_constraints +from .check_credentials import check_credentials +from .get_dataset_schema import get_dataset_schema +from .get_job_status import get_job_status +from .get_results import get_results +from .list_jobs import list_jobs +from .search_datasets import search_datasets +from .submit_request import submit_request + +__all__ = [ + "TOOLS", + "apply_constraints", + "check_credentials", + "get_dataset_schema", + "get_job_status", + "get_results", + "list_jobs", + "search_datasets", + "submit_request", +] + +TOOLS = [ + search_datasets, + get_dataset_schema, + apply_constraints, + submit_request, + get_job_status, + get_results, + list_jobs, + check_credentials, +] diff --git a/toolsets/cds/src/cds/tools/_client.py b/toolsets/cds/src/cds/tools/_client.py new file mode 100644 index 0000000..3307917 --- /dev/null +++ b/toolsets/cds/src/cds/tools/_client.py @@ -0,0 +1,12 @@ +import httpx + +from ..client import make_client + +_client: httpx.AsyncClient | None = None + + +def get_client() -> httpx.AsyncClient: + global _client + if _client is None: + _client = make_client() + return _client diff --git a/toolsets/cds/src/cds/tools/_errors.py b/toolsets/cds/src/cds/tools/_errors.py new file mode 100644 index 0000000..3ddd7a7 --- /dev/null +++ b/toolsets/cds/src/cds/tools/_errors.py @@ -0,0 +1,79 @@ +from typing import Any + +import httpx + + +def auth_error(detail: str = "Invalid or missing API key.") -> dict[str, Any]: + return {"error": "auth", "detail": detail} + + +def bad_request_error(detail: str) -> dict[str, Any]: + return {"error": "bad_request", "detail": detail} + + +def licence_error(dataset: str = "") -> dict[str, Any]: + if dataset: + url = f"https://cds.climate.copernicus.eu/datasets/{dataset}" + detail = f"Dataset licence not accepted. Accept the terms at {url} then retry." + else: + detail = ( + "Dataset licence not accepted. " + "Visit the dataset page on CDS and accept the terms before retrying." + ) + return {"error": "licence", "detail": detail} + + +def queue_limit_error() -> dict[str, Any]: + return { + "error": "queue_limit", + "detail": ( + "Concurrent queued-job limit reached. Wait for existing jobs to complete." + ), + } + + +def not_found_error(detail: str) -> dict[str, Any]: + return {"error": "not_found", "detail": detail} + + +def not_ready_error(job_id: str, status: str) -> dict[str, Any]: + return { + "error": "not_ready", + "detail": ( + f"Job {job_id!r} is not yet successful (status: {status!r}). " + "Poll again later." + ), + } + + +def transient_error(detail: str) -> dict[str, Any]: + return {"error": "transient", "detail": detail} + + +def classify_submit_failure(message: str, dataset: str = "") -> dict[str, Any]: + """Map a failed-job traceback/message to the appropriate error code.""" + lower = message.lower() + if "licence" in lower or "license" in lower or "terms" in lower: + return licence_error(dataset) + if "queue" in lower: + return queue_limit_error() + return bad_request_error(message) + + +def classify_http_error(resp: httpx.Response, dataset: str = "") -> dict[str, Any]: + if resp.status_code == 401: + return auth_error() + try: + body: Any = resp.json() + detail: Any = body.get("detail", resp.text) + except Exception: + body = {} + detail = resp.text + if resp.status_code == 403: + detail_str = str(detail).lower() + if any(w in detail_str for w in ("licence", "license", "terms")): + return licence_error(dataset) + return auth_error(str(detail)) + if resp.status_code in (400, 422): + return bad_request_error(str(detail)) + return transient_error(f"HTTP {resp.status_code}: {str(detail)[:200]}") diff --git a/toolsets/cds/src/cds/tools/_retry.py b/toolsets/cds/src/cds/tools/_retry.py new file mode 100644 index 0000000..f2a37f8 --- /dev/null +++ b/toolsets/cds/src/cds/tools/_retry.py @@ -0,0 +1,33 @@ +from collections.abc import Callable +from typing import Any, TypeVar + +import httpx +from tenacity import ( + retry, + retry_if_exception_type, + stop_after_attempt, + wait_exponential, +) + +# Exceptions that warrant automatic backoff retry. +# 5xx errors must be raised explicitly via raise_for_status() to trigger this. +TRANSIENT_EXC = ( + httpx.TimeoutException, + httpx.ConnectError, + httpx.RemoteProtocolError, + httpx.HTTPStatusError, +) + +_F = TypeVar("_F", bound=Callable[..., Any]) + +_retry = retry( + retry=retry_if_exception_type(TRANSIENT_EXC), + stop=stop_after_attempt(3), + wait=wait_exponential(multiplier=1, min=2, max=10), + reraise=True, +) + + +def with_retry(fn: _F) -> _F: + """Typed wrapper around tenacity retry so decorated functions keep their signature.""" + return _retry(fn) diff --git a/toolsets/cds/src/cds/tools/apply_constraints.py b/toolsets/cds/src/cds/tools/apply_constraints.py new file mode 100644 index 0000000..d85761d --- /dev/null +++ b/toolsets/cds/src/cds/tools/apply_constraints.py @@ -0,0 +1,50 @@ +from typing import Any + +from langchain_core.tools import tool + +from ._client import get_client +from ._errors import classify_http_error, not_found_error, transient_error +from ._retry import TRANSIENT_EXC, with_retry + + +@with_retry +async def _call(dataset: str, partial_request: dict[str, Any]) -> dict[str, Any]: + resp = await get_client().post( + f"/processes/{dataset}/constraints", + json={"inputs": partial_request}, + ) + if resp.status_code >= 500: + resp.raise_for_status() + if resp.status_code == 404: + return not_found_error(f"Dataset {dataset!r} not found.") + if resp.status_code != 200: + return classify_http_error(resp) + + raw: dict[str, list[str]] = resp.json() + # Omit parameters with empty lists — those are unconstrained (e.g. area) + valid_values = {k: v for k, v in raw.items() if v} + + return {"dataset": dataset, "valid_values": valid_values} + + +@tool +async def apply_constraints( + dataset: str, partial_request: dict[str, Any] +) -> dict[str, Any]: + """Get valid parameter values for a CDS dataset given a partial request. + + Pass whatever parameters you already know (e.g. year and month) and receive + back the valid values for all remaining parameters. Use this to verify that + requested dates exist in the archive — especially for recent or current-year data. + + Args: + dataset: Dataset identifier, e.g. "reanalysis-era5-land". + partial_request: Partial request dict, e.g. {"year": "2026", "month": "05"}. + + Returns a dict with 'dataset' and 'valid_values' mapping each constrained + parameter to its list of currently valid values. + """ + try: + return await _call(dataset, partial_request) + except TRANSIENT_EXC as exc: + return transient_error(str(exc)) diff --git a/toolsets/cds/src/cds/tools/check_credentials.py b/toolsets/cds/src/cds/tools/check_credentials.py new file mode 100644 index 0000000..0636d9c --- /dev/null +++ b/toolsets/cds/src/cds/tools/check_credentials.py @@ -0,0 +1,26 @@ +from typing import Any + +from langchain_core.tools import tool + +from ._client import get_client +from ._errors import classify_http_error, transient_error +from ._retry import TRANSIENT_EXC, with_retry + + +@with_retry +async def _call() -> dict[str, Any]: + resp = await get_client().get("/jobs", params={"limit": 1}) + if resp.status_code >= 500: + resp.raise_for_status() + if resp.status_code == 200: + return {"ok": True} + return classify_http_error(resp) + + +@tool +async def check_credentials() -> dict[str, Any]: + """Validate the CDS API key. Returns {"ok": True} on success or a structured error dict.""" + try: + return await _call() + except TRANSIENT_EXC as exc: + return transient_error(str(exc)) diff --git a/toolsets/cds/src/cds/tools/get_dataset_schema.py b/toolsets/cds/src/cds/tools/get_dataset_schema.py new file mode 100644 index 0000000..0521501 --- /dev/null +++ b/toolsets/cds/src/cds/tools/get_dataset_schema.py @@ -0,0 +1,116 @@ +from typing import Any + +from langchain_core.tools import tool + +from ._client import get_client +from ._errors import classify_http_error, not_found_error, transient_error +from ._retry import TRANSIENT_EXC, with_retry + + +def _extract_enum(schema: dict[str, Any]) -> list[Any] | None: + """Return enum values from a schema dict, handling const and items.enum.""" + if "const" in schema: + return [schema["const"]] + if "enum" in schema: + enum_vals = schema["enum"] + return enum_vals if isinstance(enum_vals, list) else [enum_vals] + if schema.get("type") == "array": + items = schema.get("items", {}) + if isinstance(items, dict) and "enum" in items: + items_enum = items["enum"] + return items_enum if isinstance(items_enum, list) else [items_enum] + return None + + +def _parse_parameter(spec: dict[str, Any]) -> dict[str, Any]: + if "schema" not in spec: + result: dict[str, Any] = { + "type": "unknown", + "required": spec.get("minOccurs", 1) != 0, + } + if spec.get("title"): + result["title"] = spec["title"] + if spec.get("description"): + result["description"] = spec["description"] + return result + + schema = spec.get("schema", {}) + + # Resolve oneOf / anyOf: union all enum values, take type from first branch + field_type: str | None = None + union_values: list[Any] = [] + for composite_key in ("oneOf", "anyOf"): + if composite_key in schema: + for branch in schema[composite_key]: + if field_type is None: + field_type = branch.get("type") + vals = _extract_enum(branch) + if vals: + for v in vals: + if v not in union_values: + union_values.append(v) + break + + if field_type is None: + field_type = schema.get("type", "string") + + default = schema.get("default") + min_occurs = spec.get("minOccurs", 1) + required = (default is None) and (min_occurs != 0) + + result = {"type": field_type, "required": required} + if default is not None: + result["default"] = default + if spec.get("title"): + result["title"] = spec["title"] + if spec.get("description"): + result["description"] = spec["description"] + + # Enum / values + if union_values: + result["values"] = union_values + else: + vals = _extract_enum(schema) + if vals is not None: + result["values"] = vals + elif field_type in ("integer", "number"): + if "minimum" in schema: + result["minimum"] = schema["minimum"] + if "maximum" in schema: + result["maximum"] = schema["maximum"] + + return result + + +@with_retry +async def _call(dataset: str) -> dict[str, Any]: + resp = await get_client().get(f"/processes/{dataset}") + if resp.status_code >= 500: + resp.raise_for_status() + if resp.status_code == 404: + return not_found_error(f"Dataset {dataset!r} not found.") + if resp.status_code != 200: + return classify_http_error(resp) + + data = resp.json() + raw_inputs: dict[str, Any] = data.get("inputs", {}) + parameters = {name: _parse_parameter(spec) for name, spec in raw_inputs.items()} + + return {"dataset": dataset, "parameters": parameters} + + +@tool +async def get_dataset_schema(dataset: str) -> dict[str, Any]: + """Get the input schema for a CDS dataset: required parameters, valid values, and types. + + Args: + dataset: Dataset identifier, e.g. "reanalysis-era5-land". + + Returns a dict with 'dataset' and 'parameters' mapping each field to its type, + required flag, valid values list, and default (if optional). + Always call this before submit_request to know what parameters are needed. + """ + try: + return await _call(dataset) + except TRANSIENT_EXC as exc: + return transient_error(str(exc)) diff --git a/toolsets/cds/src/cds/tools/get_job_status.py b/toolsets/cds/src/cds/tools/get_job_status.py new file mode 100644 index 0000000..fd9f510 --- /dev/null +++ b/toolsets/cds/src/cds/tools/get_job_status.py @@ -0,0 +1,42 @@ +from typing import Any + +from langchain_core.tools import tool + +from ._client import get_client +from ._errors import classify_http_error, transient_error +from ._retry import TRANSIENT_EXC, with_retry + + +@with_retry +async def _call(job_id: str) -> dict[str, Any]: + resp = await get_client().get(f"/jobs/{job_id}") + if resp.status_code >= 500: + resp.raise_for_status() + if resp.status_code != 200: + return classify_http_error(resp) + + data = resp.json() + status: str = data.get("status", "unknown") + return { + "job_id": job_id, + "status": status, + "results_ready": status == "successful", + "created": data.get("created"), + "started": data.get("started"), + "finished": data.get("finished"), + } + + +@tool +async def get_job_status(job_id: str) -> dict[str, Any]: + """Get the current status of a CDS job. + + Args: + job_id: The job ID returned by submit_request. + + Returns dict with job_id, status, results_ready flag, and timestamps. + """ + try: + return await _call(job_id) + except TRANSIENT_EXC as exc: + return transient_error(str(exc)) diff --git a/toolsets/cds/src/cds/tools/get_results.py b/toolsets/cds/src/cds/tools/get_results.py new file mode 100644 index 0000000..7ce5c8c --- /dev/null +++ b/toolsets/cds/src/cds/tools/get_results.py @@ -0,0 +1,56 @@ +from typing import Any + +from langchain_core.tools import tool + +from ._client import get_client +from ._errors import classify_http_error, not_ready_error, transient_error +from ._retry import TRANSIENT_EXC, with_retry + + +@with_retry +async def _call(job_id: str) -> dict[str, Any]: + client = get_client() + + # Confirm the job is successful before fetching assets. + status_resp = await client.get(f"/jobs/{job_id}") + if status_resp.status_code >= 500: + status_resp.raise_for_status() + if status_resp.status_code != 200: + return classify_http_error(status_resp) + + status: str = status_resp.json().get("status", "unknown") + if status != "successful": + return not_ready_error(job_id, status) + + results_resp = await client.get(f"/jobs/{job_id}/results") + if results_resp.status_code >= 500: + results_resp.raise_for_status() + if results_resp.status_code != 200: + return classify_http_error(results_resp) + + try: + val = results_resp.json()["asset"]["value"] + return { + "href": val["href"], + "content_type": val.get("type"), + "size": val.get("file:size"), + "checksum": val.get("file:checksum"), + } + except (KeyError, TypeError): + return transient_error("Unexpected results shape from CDS API.") + + +@tool +async def get_results(job_id: str) -> dict[str, Any]: + """Get the download link for a completed CDS job. + + Args: + job_id: The job ID of a job with status "successful". + + Returns dict with href, content_type, size, checksum — or a structured error. + The href is a public URL; no token needed to download the file. + """ + try: + return await _call(job_id) + except TRANSIENT_EXC as exc: + return transient_error(str(exc)) diff --git a/toolsets/cds/src/cds/tools/list_jobs.py b/toolsets/cds/src/cds/tools/list_jobs.py new file mode 100644 index 0000000..25c0c1b --- /dev/null +++ b/toolsets/cds/src/cds/tools/list_jobs.py @@ -0,0 +1,64 @@ +from typing import Any + +from langchain_core.tools import tool + +from ._client import get_client +from ._errors import classify_http_error, transient_error +from ._retry import TRANSIENT_EXC, with_retry + + +@with_retry +async def _call(status: list[str] | None, limit: int) -> dict[str, Any]: + params: dict[str, Any] = {"limit": limit} + if status: + params["status"] = status # httpx repeats the key for each value + + resp = await get_client().get("/jobs", params=params) + if resp.status_code >= 500: + resp.raise_for_status() + if resp.status_code != 200: + return classify_http_error(resp) + + data = resp.json() + jobs = [ + { + "job_id": j.get("jobID") or j.get("job_id"), + "dataset": j.get("processID"), + "status": j.get("status"), + "results_ready": j.get("status") == "successful", + "created": j.get("created"), + "finished": j.get("finished"), + } + for j in data.get("jobs", []) + ] + + result: dict[str, Any] = {"jobs": jobs, "count": len(jobs)} + next_link = next( + (lnk["href"] for lnk in data.get("links", []) if lnk.get("rel") == "next"), + None, + ) + if next_link: + result["next"] = next_link + + return result + + +@tool +async def list_jobs( + status: list[str] | None = None, + limit: int = 100, +) -> dict[str, Any]: + """List CDS jobs, optionally filtered by status. + + Args: + status: One or more statuses to filter by. Valid values: "accepted", + "running", "successful", "failed". Omit to return all statuses. + limit: Maximum number of jobs to return (default 100). + + Returns dict with a "jobs" list and "count". Includes "next" cursor URL when + more results are available. + """ + try: + return await _call(status, limit) + except TRANSIENT_EXC as exc: + return transient_error(str(exc)) diff --git a/toolsets/cds/src/cds/tools/search_datasets.py b/toolsets/cds/src/cds/tools/search_datasets.py new file mode 100644 index 0000000..bcc09c5 --- /dev/null +++ b/toolsets/cds/src/cds/tools/search_datasets.py @@ -0,0 +1,49 @@ +from typing import Any + +import httpx +from langchain_core.tools import tool + +from ..settings import settings +from ._errors import transient_error +from ._retry import TRANSIENT_EXC, with_retry + + +@with_retry +async def _call(query: str) -> list[dict[str, Any]]: + async with httpx.AsyncClient( + timeout=httpx.Timeout(connect=10.0, read=30.0, write=10.0, pool=10.0) + ) as client: + resp = await client.get( + f"{settings.cds_catalogue_url}/datasets", + params={"q": query, "limit": 10}, + ) + if resp.status_code >= 500: + resp.raise_for_status() + if resp.status_code != 200: + return [{"error": "catalogue_error", "detail": resp.text[:200]}] + + data = resp.json() + results = ( + data.get("results") or data.get("items") or data.get("collections") or [] + ) + return [ + {"id": item.get("id", ""), "title": item.get("title", "")} + for item in results + if item.get("id") + ] + + +@tool +async def search_datasets(query: str) -> list[dict[str, Any]]: + """Search for CDS datasets by keyword. + + Args: + query: Search term, e.g. "ERA5 temperature" or "sea surface salinity". + + Returns a list of matching datasets, each with 'id' and 'title'. + Use the 'id' as the dataset argument in get_dataset_schema and submit_request. + """ + try: + return await _call(query) + except TRANSIENT_EXC as exc: + return [transient_error(str(exc))] diff --git a/toolsets/cds/src/cds/tools/submit_request.py b/toolsets/cds/src/cds/tools/submit_request.py new file mode 100644 index 0000000..045c576 --- /dev/null +++ b/toolsets/cds/src/cds/tools/submit_request.py @@ -0,0 +1,57 @@ +from typing import Any + +from langchain_core.tools import tool + +from ._client import get_client +from ._errors import classify_http_error, classify_submit_failure, transient_error +from ._retry import TRANSIENT_EXC, with_retry + + +def _failure_message(data: dict[str, Any]) -> str: + meta = data.get("metadata") or {} + results = meta.get("results") or {} + return ( + results.get("traceback") + or results.get("message") + or data.get("message") + or "Unknown failure reason." + ) + + +@with_retry +async def _call(dataset: str, request: dict[str, Any]) -> dict[str, Any]: + resp = await get_client().post( + f"/processes/{dataset}/execution", + json={"inputs": request}, + ) + if resp.status_code >= 500: + resp.raise_for_status() + if resp.status_code not in (200, 201): + return classify_http_error(resp, dataset) + + data = resp.json() + status: str = data.get("status", "accepted") + + if status == "failed": + return classify_submit_failure(_failure_message(data), dataset) + + return { + "job_id": data.get("jobID") or data.get("job_id") or "", + "status": status, + } + + +@tool +async def submit_request(dataset: str, request: dict[str, Any]) -> dict[str, Any]: + """Submit a data request to the CDS. + + Args: + dataset: Dataset identifier, e.g. "reanalysis-era5-land". + request: Bare request parameters dict — do NOT wrap in "inputs", that is added internally. + + Returns dict with job_id and status, or a structured error dict. + """ + try: + return await _call(dataset, request) + except TRANSIENT_EXC as exc: + return transient_error(str(exc)) diff --git a/toolsets/cds/tests/test_cds.py b/toolsets/cds/tests/test_cds.py new file mode 100644 index 0000000..e4caa9e --- /dev/null +++ b/toolsets/cds/tests/test_cds.py @@ -0,0 +1,326 @@ +"""Unit tests for the CDS toolset (ported from cds-assistant).""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx + +from cds.tools import TOOLS +from cds.tools.apply_constraints import _call as _apply_call +from cds.tools.get_dataset_schema import _call as _schema_call +from cds.tools.get_dataset_schema import _parse_parameter +from cds.tools.search_datasets import _call as _search_call + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +ERA5_PROCESS_RESPONSE = { + "id": "reanalysis-era5-land", + "inputs": { + "variable": { + "schema": { + "type": "array", + "items": {"enum": ["2m_temperature", "total_precipitation"]}, + } + }, + "year": {"schema": {"type": "string", "enum": ["2020", "2024"]}}, + "month": {"schema": {"type": "string", "enum": ["01", "12"]}}, + "data_format": {"schema": {"type": "string", "enum": ["grib", "netcdf"]}}, + "area": {"schema": {"type": "array", "default": [90, -180, -90, 180]}}, + "download_format": { + "schema": { + "type": "string", + "enum": ["zip", "unarchived"], + "default": "unarchived", + } + }, + }, +} + +CONSTRAINTS_RESPONSE = { + "variable": ["2m_temperature", "total_precipitation"], + "day": ["01", "02", "28"], + "time": ["00:00", "12:00"], + "data_format": ["grib", "netcdf"], + "area": [], # unconstrained — should be filtered + "download_format": ["unarchived", "zip"], +} + +CATALOGUE_RESPONSE = { + "results": [ + {"id": "reanalysis-era5-land", "title": "ERA5-Land hourly data"}, + { + "id": "reanalysis-era5-pressure-levels", + "title": "ERA5 hourly on pressure levels", + }, + ] +} + + +def _mock_client(method: str, response: httpx.Response) -> MagicMock: + client = MagicMock() + setattr(client, method, AsyncMock(return_value=response)) + return client + + +# --------------------------------------------------------------------------- +# TOOLS export +# --------------------------------------------------------------------------- + + +def test_tools_exported() -> None: + names = {tool.name for tool in TOOLS} + assert names == { + "search_datasets", + "get_dataset_schema", + "apply_constraints", + "submit_request", + "get_job_status", + "get_results", + "list_jobs", + "check_credentials", + } + + +# --------------------------------------------------------------------------- +# _parse_parameter (pure function — no mocking needed) +# --------------------------------------------------------------------------- + + +def test_parse_string_field_with_enum() -> None: + spec = {"schema": {"type": "string", "enum": ["grib", "netcdf"]}} + result = _parse_parameter(spec) + assert result == {"type": "string", "required": True, "values": ["grib", "netcdf"]} + + +def test_parse_array_field_with_items_enum() -> None: + spec = {"schema": {"type": "array", "items": {"enum": ["00:00", "06:00", "12:00"]}}} + result = _parse_parameter(spec) + assert result == { + "type": "array", + "required": True, + "values": ["00:00", "06:00", "12:00"], + } + + +def test_parse_optional_field_with_default_only() -> None: + spec = {"schema": {"type": "array", "default": [90, -180, -90, 180]}} + result = _parse_parameter(spec) + assert result["required"] is False + assert result["default"] == [90, -180, -90, 180] + assert "values" not in result + + +def test_parse_optional_field_with_default_and_enum() -> None: + spec = { + "schema": { + "type": "string", + "enum": ["zip", "unarchived"], + "default": "unarchived", + } + } + result = _parse_parameter(spec) + assert result["required"] is False + assert result["values"] == ["zip", "unarchived"] + assert result["default"] == "unarchived" + + +# --------------------------------------------------------------------------- +# get_dataset_schema +# --------------------------------------------------------------------------- + + +async def test_get_dataset_schema_parses_all_fields() -> None: + client = _mock_client("get", httpx.Response(200, json=ERA5_PROCESS_RESPONSE)) + + with patch("cds.tools.get_dataset_schema.get_client", return_value=client): + result = await _schema_call("reanalysis-era5-land") + + assert result["dataset"] == "reanalysis-era5-land" + params = result["parameters"] + + assert params["variable"] == { + "type": "array", + "required": True, + "values": ["2m_temperature", "total_precipitation"], + } + assert params["year"] == { + "type": "string", + "required": True, + "values": ["2020", "2024"], + } + assert params["area"]["required"] is False + assert params["area"]["default"] == [90, -180, -90, 180] + assert params["download_format"]["required"] is False + assert params["download_format"]["values"] == ["zip", "unarchived"] + + +async def test_get_dataset_schema_not_found() -> None: + client = _mock_client("get", httpx.Response(404, json={"detail": "not found"})) + + with patch("cds.tools.get_dataset_schema.get_client", return_value=client): + result = await _schema_call("bad-dataset") + + assert result["error"] == "not_found" + assert "bad-dataset" in result["detail"] + + +# --------------------------------------------------------------------------- +# apply_constraints +# --------------------------------------------------------------------------- + + +async def test_apply_constraints_filters_empty_lists() -> None: + client = _mock_client("post", httpx.Response(200, json=CONSTRAINTS_RESPONSE)) + + with patch("cds.tools.apply_constraints.get_client", return_value=client): + result = await _apply_call( + "reanalysis-era5-land", {"year": "2024", "month": "01"} + ) + + assert result["dataset"] == "reanalysis-era5-land" + assert "area" not in result["valid_values"] # empty list removed + assert result["valid_values"]["day"] == ["01", "02", "28"] + + +async def test_apply_constraints_sends_inputs_wrapper() -> None: + client = _mock_client("post", httpx.Response(200, json=CONSTRAINTS_RESPONSE)) + + with patch("cds.tools.apply_constraints.get_client", return_value=client): + await _apply_call("reanalysis-era5-land", {"year": "2024", "month": "01"}) + + _, call_kwargs = client.post.call_args + assert call_kwargs["json"] == {"inputs": {"year": "2024", "month": "01"}} + + +async def test_apply_constraints_not_found() -> None: + client = _mock_client("post", httpx.Response(404, json={"detail": "not found"})) + + with patch("cds.tools.apply_constraints.get_client", return_value=client): + result = await _apply_call("bad-dataset", {}) + + assert result["error"] == "not_found" + + +# --------------------------------------------------------------------------- +# search_datasets +# --------------------------------------------------------------------------- + + +async def test_search_datasets_returns_id_and_title() -> None: + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = CATALOGUE_RESPONSE + + mock_client = AsyncMock() + mock_client.get.return_value = mock_response + mock_client.__aenter__.return_value = mock_client + mock_client.__aexit__.return_value = None + + with patch("cds.tools.search_datasets.httpx.AsyncClient", return_value=mock_client): + results = await _search_call("ERA5") + + assert len(results) == 2 + assert results[0] == { + "id": "reanalysis-era5-land", + "title": "ERA5-Land hourly data", + } + assert results[1]["id"] == "reanalysis-era5-pressure-levels" + + +async def test_search_datasets_passes_q_param() -> None: + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"results": []} + + mock_client = AsyncMock() + mock_client.get.return_value = mock_response + mock_client.__aenter__.return_value = mock_client + mock_client.__aexit__.return_value = None + + with patch("cds.tools.search_datasets.httpx.AsyncClient", return_value=mock_client): + await _search_call("sea surface temperature") + + _, call_kwargs = mock_client.get.call_args + assert call_kwargs["params"]["q"] == "sea surface temperature" + assert call_kwargs["params"]["limit"] == 10 + + +# --------------------------------------------------------------------------- +# _parse_parameter — new schema types +# --------------------------------------------------------------------------- + + +def test_parse_integer_field_with_range() -> None: + spec = {"schema": {"type": "integer", "minimum": 1, "maximum": 1000}} + result = _parse_parameter(spec) + assert result["type"] == "integer" + assert result["required"] is True + assert result["minimum"] == 1 + assert result["maximum"] == 1000 + assert "values" not in result + + +def test_parse_number_field_no_range() -> None: + spec = {"schema": {"type": "number"}} + result = _parse_parameter(spec) + assert result["type"] == "number" + assert result["required"] is True + assert "minimum" not in result + assert "maximum" not in result + + +def test_parse_boolean_field() -> None: + spec = {"schema": {"type": "boolean"}} + result = _parse_parameter(spec) + assert result["type"] == "boolean" + assert result["required"] is True + assert "values" not in result + + +def test_parse_one_of_schema() -> None: + spec = { + "schema": { + "oneOf": [ + {"type": "string", "enum": ["grib"]}, + {"type": "string", "enum": ["netcdf", "netcdf4"]}, + ] + } + } + result = _parse_parameter(spec) + assert result["type"] == "string" + assert set(result["values"]) == {"grib", "netcdf", "netcdf4"} + + +def test_parse_const_value() -> None: + spec = {"schema": {"type": "string", "const": "netcdf"}} + result = _parse_parameter(spec) + assert result["values"] == ["netcdf"] + assert result["required"] is True + + +def test_parse_spec_with_title_and_description() -> None: + spec = { + "title": "Output format", + "description": "Format of the downloaded file.", + "schema": {"type": "string", "enum": ["grib", "netcdf"]}, + } + result = _parse_parameter(spec) + assert result["title"] == "Output format" + assert result["description"] == "Format of the downloaded file." + assert result["values"] == ["grib", "netcdf"] + + +def test_parse_min_occurs_zero_marks_optional() -> None: + spec = {"minOccurs": 0, "schema": {"type": "string", "enum": ["a", "b"]}} + result = _parse_parameter(spec) + assert result["required"] is False + assert result["values"] == ["a", "b"] + + +def test_parse_no_schema_key() -> None: + spec = {"title": "Mystery field"} + result = _parse_parameter(spec) + assert result["type"] == "unknown" + assert result["required"] is True + assert result["title"] == "Mystery field" diff --git a/toolsets/cds/toolset.yaml b/toolsets/cds/toolset.yaml new file mode 100644 index 0000000..7edd03d --- /dev/null +++ b/toolsets/cds/toolset.yaml @@ -0,0 +1,6 @@ +# Optional Helm values overrides for this toolset (see charts/mcp-toolset/values.yaml). +# The cds-credentials Secret must exist in the namespace and provide CDS_API_KEY +# (create it out-of-band: kubectl -n mcp-toolsets create secret generic +# cds-credentials --from-literal=CDS_API_KEY=...). +secrets: + - cds-credentials diff --git a/uv.lock b/uv.lock index 47608b4..b1f9b40 100644 --- a/uv.lock +++ b/uv.lock @@ -9,6 +9,7 @@ resolution-markers = [ [manifest] members = [ "aoi-generator", + "cds", "dataset-search", "mcp-agent", "mcp-cli", @@ -224,6 +225,27 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/99/37/e8730c3587a65eb5645d4aba2d27aae48e8003614d6aaf15dda67f702f1f/bidict-0.23.1-py3-none-any.whl", hash = "sha256:5dae8d4d79b552a71cbabc7deb25dfe8ce710b17ff41711e13010ead2abfc3e5", size = 32764, upload-time = "2024-02-18T19:09:04.156Z" }, ] +[[package]] +name = "cds" +version = "0.1.0" +source = { editable = "toolsets/cds" } +dependencies = [ + { name = "httpx" }, + { name = "langchain-core" }, + { name = "mcp-runtime" }, + { name = "pydantic-settings" }, + { name = "tenacity" }, +] + +[package.metadata] +requires-dist = [ + { name = "httpx", specifier = ">=0.27" }, + { name = "langchain-core", specifier = ">=1.4.6,<2.0.0" }, + { name = "mcp-runtime", editable = "packages/mcp-runtime" }, + { name = "pydantic-settings", specifier = ">=2.5" }, + { name = "tenacity", specifier = ">=9.0" }, +] + [[package]] name = "certifi" version = "2026.5.20" @@ -1152,6 +1174,7 @@ version = "0.1.0" source = { virtual = "." } dependencies = [ { name = "aoi-generator" }, + { name = "cds" }, { name = "dataset-search" }, { name = "mcp-agent" }, { name = "mcp-cli" }, @@ -1169,6 +1192,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "aoi-generator", editable = "toolsets/aoi-generator" }, + { name = "cds", editable = "toolsets/cds" }, { name = "dataset-search", editable = "toolsets/dataset-search" }, { name = "mcp-agent", editable = "packages/mcp-agent" }, { name = "mcp-cli", editable = "packages/mcp-cli" },