mlspace-plugin 0.2.1__py3-none-any.whl

mlspace_mcp/tls.py

@@ -0,0 +1,46 @@
+"""One verified TLS policy for setup and runtime; no global SSL monkey patch."""
+from __future__ import annotations
+
+import os
+import ssl
+
+import httpx
+
+
+def tls_context(ca_file: str = '') -> ssl.SSLContext:
+    """Use OS trust and optionally an administrator-supplied PEM bundle."""
+    extra = ca_file or os.environ.get('SSL_CERT_FILE', '')
+    context: ssl.SSLContext
+    try:
+        import truststore
+        context = truststore.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+    except ImportError:
+        context = ssl.create_default_context()
+    if extra:
+        context.load_verify_locations(cafile=extra)
+    capath = os.environ.get('SSL_CERT_DIR')
+    if capath:
+        context.load_verify_locations(capath=capath)
+    return context
+
+
+def safe_error(exc: Exception) -> str:
+    """Describe the cause without printing exception bodies or submitted values."""
+    chain: list[BaseException] = []
+    current: BaseException | None = exc
+    while current is not None and all(current is not item for item in chain):
+        chain.append(current)
+        current = current.__cause__ or current.__context__
+    if any(isinstance(item, ssl.SSLCertVerificationError) or
+           (isinstance(item, (ssl.SSLError, httpx.ConnectError)) and
+            'CERTIFICATE_VERIFY_FAILED' in str(item)) for item in chain):
+        return ('TLS certificate verification failed. Use the corporate CA from your '
+                'administrator: --ca-file /path/to/ca.pem. TLS verification remains enabled.')
+    status = getattr(exc, 'status', None)
+    if isinstance(status, int):
+        return f'API request failed (HTTP {status}); check access keys and workspace permissions.'
+    if any(isinstance(item, httpx.TimeoutException) for item in chain):
+        return 'API request timed out; check network/VPN and retry.'
+    if any(isinstance(item, httpx.ConnectError) for item in chain):
+        return 'Could not connect; check network/VPN and the API endpoint.'
+    return type(exc).__name__

mlspace_mcp/tools/__init__.py

@@ -0,0 +1 @@
+"""Domain umbrella-tool definitions. Each module exposes ``DOMAIN: DomainTool``."""

mlspace_mcp/tools/allocations.py

@@ -0,0 +1,92 @@
+"""``mlspace_allocations`` — inspect MLSpace compute allocations and their defaults.
+
+Pure declarative data: a single ``DOMAIN: DomainTool``. All request mechanics
+(headers, path templating, query encoding, bodies, readonly, confirm,
+formatting) live in the core. See ``inference.py`` for the canonical shape.
+"""
+
+from __future__ import annotations
+
+from ..registry import DomainTool, Op, Param
+
+# action -> Op. The (method, path) pairs are the single source of truth for the
+# drift test; paths are copied verbatim from the spec, placeholders included.
+ACTIONS: dict[str, Op] = {
+    "list": Op(
+        "GET", "/public/v2/allocations/",
+        kind="list",
+        help="List allocations in the workspace (personal account only).",
+    ),
+    "list_defaults": Op(
+        "GET", "/public/v2/allocations/defaults",
+        query_params={"region": "region"},
+        required=("region",),
+        kind="list",
+        help="List default allocations in a region.",
+    ),
+    "get": Op(
+        "GET", "/public/v2/allocations/{allocation_id}",
+        path_params={"allocation_id": "allocation_id"},
+        required=("allocation_id",),
+        help="Get one allocation with its resource status.",
+    ),
+    "list_assignments": Op(
+        "GET", "/public/v2/allocations/{allocation_id}/assignments",
+        path_params={"allocation_id": "allocation_id"},
+        required=("allocation_id",),
+        kind="list",
+        help="List an allocation's assignments by organization structure.",
+    ),
+    "list_instance_types": Op(
+        "GET", "/public/v2/allocations/{allocation_id}/instance-types",
+        path_params={"allocation_id": "allocation_id"},
+        required=("allocation_id",),
+        kind="list",
+        help="List instance types available in an allocation.",
+    ),
+    "get_nodes": Op(
+        "GET", "/public/v2/allocations/{allocation_id}/nodes",
+        path_params={"allocation_id": "allocation_id"},
+        required=("allocation_id",),
+        help="Get the GPU nodes in an allocation.",
+    ),
+    "set_default": Op(
+        "POST", "/public/v2/allocations/{allocation_id}/defaults",
+        path_params={"allocation_id": "allocation_id"},
+        body_field="body",
+        required=("allocation_id", "body"),
+        write=True,
+        help="Set an allocation as default for job/notebook; `body` = DefaultRequest "
+        "{org_structure, default_for}.",
+    ),
+    "unset_default": Op(
+        "DELETE", "/public/v2/allocations/{allocation_id}/defaults",
+        path_params={"allocation_id": "allocation_id"},
+        body_field="body",
+        required=("allocation_id", "body"),
+        write=True,
+        help="Unset an allocation as default for job/notebook; `body` = DefaultRequest "
+        "{org_structure, default_for}.",
+    ),
+}
+
+PARAMS: list[Param] = [
+    Param("allocation_id", str, "Allocation ID (UUID)."),
+    Param("region", str, "Region key (required for list_defaults)."),
+    Param(
+        "body",
+        dict,
+        "JSON body = DefaultRequest for set_default/unset_default: "
+        "{org_structure: {id, type}, default_for: 'job'|'notebook'}.",
+    ),
+]
+
+DOMAIN = DomainTool(
+    domain="allocations",
+    name="mlspace_allocations",
+    title="MLSpace Allocations",
+    summary="Inspect MLSpace compute allocations, their nodes, instance types, "
+    "assignments, and manage default allocations.",
+    actions=ACTIONS,
+    params=PARAMS,
+)

mlspace_mcp/tools/async_inference.py

@@ -0,0 +1,77 @@
+"""``mlspace_async_inference`` — submit and track asynchronous inference requests.
+
+Pure declarative data following the ``inference`` exemplar; all request
+mechanics live in the core (registry.py + client.py).
+"""
+
+from __future__ import annotations
+
+from ..registry import DomainTool, Op, Param
+
+ACTIONS: dict[str, Op] = {
+    "list": Op(
+        "GET", "/public/v2/async_inferences/v1/",
+        query_params={"inference_name": "inference_name"},
+        required=("inference_name",),
+        kind="list",
+        help="List asynchronous requests for an inference service.",
+    ),
+    "predict": Op(
+        "POST",
+        "/public/v2/async_inferences/v1/{region}/{inference_name}/{predict_path}",
+        path_params={
+            "region": "region",
+            "inference_name": "inference_name",
+            "predict_path": "predict_path",
+        },
+        # predict_path is a deliberate multi-segment tail (e.g. v1/models/svc:predict):
+        # it keeps its "/" separators; region and inference_name stay opaque segments.
+        path_tail=frozenset({"predict_path"}),
+        body_field="body",
+        required=("region", "inference_name", "predict_path"),
+        side_effect=True,  # submits a (paid) inference job — not a read
+        help="Send an async prediction request; put the model's input in `body` (free-form).",
+    ),
+    "get_result": Op(
+        "GET", "/public/v2/async_inferences/v1/{request_id}/result",
+        path_params={"request_id": "request_id"},
+        query_params={"inference_name": "inference_name"},
+        required=("request_id", "inference_name"),
+        help="Get the result of an asynchronous request.",
+    ),
+    "get_status": Op(
+        "GET", "/public/v2/async_inferences/v1/{request_id}/status",
+        path_params={"request_id": "request_id"},
+        query_params={"inference_name": "inference_name"},
+        required=("request_id", "inference_name"),
+        help="Get the status of an asynchronous request.",
+    ),
+}
+
+PARAMS: list[Param] = [
+    Param("inference_name", str, "Name of the inference service."),
+    Param(
+        "region",
+        str,
+        "Inference region, e.g. DGX2-INF-001, CCE-INF, INF-002 "
+        "(authoritative list via mlspace_resources configs cluster_type=INF). "
+        "Required for predict.",
+    ),
+    Param(
+        "predict_path",
+        str,
+        "Predict path, usually `v1/models/{inference_name}:predict` (predict only).",
+    ),
+    Param("request_id", str, "ID of the async request (required for get_result/get_status)."),
+    Param("body", dict, "Free-form model input payload for predict."),
+]
+
+DOMAIN = DomainTool(
+    domain="async_inference",
+    name="mlspace_async_inference",
+    title="MLSpace Async Inference",
+    summary="Submit and track asynchronous MLSpace inference requests (predict, status, result).",
+    actions=ACTIONS,
+    params=PARAMS,
+    redact_result_keys=("inference_key",),
+)

mlspace_mcp/tools/build_image.py

@@ -0,0 +1,76 @@
+"""``mlspace_build_image`` — manage MLSpace build-image jobs.
+
+Pure declarative data; the core (registry.py + client.py) owns all request
+mechanics. Mirrors the exemplar ``inference.py`` in shape.
+"""
+
+from __future__ import annotations
+
+from ..registry import DomainTool, Op, Param
+
+# action -> Op. (method, path) pairs are the single source of truth for the
+# drift test; paths are copied verbatim from the spec.
+ACTIONS: dict[str, Op] = {
+    "list": Op(
+        "GET", "/public/v2/build-image/jobs",
+        kind="list",
+        help="List build-image jobs in the workspace.",
+    ),
+    "run": Op(
+        "POST", "/public/v2/build-image/jobs",
+        body_field="body",
+        required=("body",),
+        write=True,
+        help="Run a build-image job. NOTE: this endpoint requires a "
+        "multipart/form-data file upload (a requirements.txt), which this tool "
+        "does NOT support — a JSON body will be rejected by the API. Build images "
+        "via the MLSpace console/CLI instead.",
+    ),
+    "get": Op(
+        "GET", "/public/v2/build-image/jobs/{job_name}",
+        path_params={"job_name": "job_name"},
+        required=("job_name",),
+        help="Get a build-image job's status by name.",
+    ),
+    "get_logs": Op(
+        "GET", "/public/v2/build-image/jobs/{job_name}/logs",
+        path_params={"job_name": "job_name"},
+        required=("job_name",),
+        kind="log",
+        help="Get a build-image job's logs by name.",
+    ),
+    "get_pods": Op(
+        "GET", "/public/v2/build-image/jobs/{job_name}/pods",
+        path_params={"job_name": "job_name"},
+        required=("job_name",),
+        kind="list",
+        help="List a build-image job's pods by name.",
+    ),
+    "delete": Op(
+        "DELETE", "/public/v2/build-image/jobs/{job_name}",
+        path_params={"job_name": "job_name"},
+        required=("job_name",),
+        write=True,
+        confirm=True,
+        help="Delete a build-image job by name.",
+    ),
+}
+
+PARAMS: list[Param] = [
+    Param("job_name", str, "Name of the build-image job."),
+    Param(
+        "body",
+        dict,
+        "Unused: the run endpoint needs a multipart file upload, not a JSON body "
+        "(not supported by this tool — use the console/CLI to build images).",
+    ),
+]
+
+DOMAIN = DomainTool(
+    domain="build_image",
+    name="mlspace_build_image",
+    title="MLSpace Build Image",
+    summary="Manage MLSpace build-image jobs (run, inspect status/logs/pods, delete).",
+    actions=ACTIONS,
+    params=PARAMS,
+)

mlspace_mcp/tools/dalle.py

@@ -0,0 +1,52 @@
+"""``mlspace_dalle`` — call the MLSpace DALL-E async inference API.
+
+Pure declarative data only; the core (registry.py + client.py) owns all request
+mechanics. Two operations: submit an async predict request, then poll for its
+status/result (which may be a generated image).
+"""
+
+from __future__ import annotations
+
+from ..registry import DomainTool, Op, Param
+
+ACTIONS: dict[str, Op] = {
+    "predict": Op(
+        "POST", "/public/v2/dalle/v1/predict/{model_name}",
+        path_params={"model_name": "model_name"},
+        body_field="body",
+        required=("model_name",),
+        side_effect=True,  # submits a (paid) inference job — not a read
+        help="Send an async request to DALL-E inference; put the model's input in `body` (free-form). Returns a request_id.",
+    ),
+    "result": Op(
+        "GET", "/public/v2/dalle/v1/result/{request_id}",
+        path_params={"request_id": "request_id"},
+        required=("request_id",),
+        kind="binary",
+        help="Get status or result (may be a generated image) for a prior predict request_id.",
+    ),
+}
+
+PARAMS: list[Param] = [
+    Param(
+        "model_name",
+        str,
+        "DALL-E model name to run the prediction on (predict only). NOTE: no MLSpace "
+        "endpoint lists DALL-E models — the ai_services catalog returned 403 on the "
+        "tested workspace (RBAC role model access checking error) and the /inference "
+        "list is empty — so this name is NOT discoverable via the API; take it from the "
+        "platform docs/console. A wrong name yields an opaque 404 that is byte-identical "
+        "for 'unknown model' and 'service not deployed'.",
+    ),
+    Param("request_id", str, "Request id returned by predict, used to poll for the result."),
+    Param("body", dict, "Free-form JSON input payload forwarded to the DALL-E model (predict only)."),
+]
+
+DOMAIN = DomainTool(
+    domain="dalle",
+    name="mlspace_dalle",
+    title="MLSpace DALL-E",
+    summary="Call the MLSpace DALL-E async inference API: submit a predict request, then poll for the status/result.",
+    actions=ACTIONS,
+    params=PARAMS,
+)

mlspace_mcp/tools/data_transfer.py

@@ -0,0 +1,260 @@
+"""``mlspace_data_transfer`` — manage MLSpace data-transfer connectors, transfers, and history.
+
+Pure declarative data (see ``inference.py`` exemplar). The core owns all request
+mechanics. Spans API versions v2/v3/v4/v5; paths are copied verbatim from the
+slice, placeholders included.
+"""
+
+from __future__ import annotations
+
+from ..registry import DomainTool, Op, Param
+
+ACTIONS: dict[str, Op] = {
+    # --- connectors (v3 collection + v2 delete) ---------------------------
+    "list_connectors": Op(
+        "GET", "/public/v2/data_transfer/v3/connectors",
+        # page/page_size are the ONLY working shrink lever for this ~27KB list: the
+        # API ignores limit/offset here (limit=5&offset=0 still returned all 65). No
+        # auto-default page_size — the response has no total_count, so a silent clamp
+        # would hide data (breaks principle #1). Pagination stays strictly opt-in.
+        query_params={"page": "page", "page_size": "page_size"},
+        kind="list",
+        help="List all connectors (system, personal, public) in the workspace.",
+    ),
+    "get_connector": Op(
+        "GET", "/public/v2/data_transfer/v3/connectors/{connector_type}/{id_}",
+        path_params={"connector_type": "connector_type", "id_": "id_"},
+        required=("connector_type", "id_"),
+        help="Get one connector by type and id.",
+    ),
+    "create_connector": Op(
+        "POST", "/public/v2/data_transfer/v3/connectors",
+        body_field="body",
+        required=("body",),
+        write=True,
+        help="Create a connector; `body` = ConnectorInputViewModelV3.",
+    ),
+    "update_connector": Op(
+        "POST", "/public/v2/data_transfer/v3/connectors/{connector_type}/{id_}",
+        path_params={"connector_type": "connector_type", "id_": "id_"},
+        body_field="body",
+        required=("connector_type", "id_", "body"),
+        write=True,
+        help="Update a connector; `body` = ConnectorUpdateViewModelV3.",
+    ),
+    "fav_connector": Op(
+        "POST", "/public/v2/data_transfer/v3/connectors/{connector_type}/{id_}/fav",
+        path_params={"connector_type": "connector_type", "id_": "id_"},
+        body_field="body",
+        required=("connector_type", "id_", "body"),
+        write=True,
+        help="Mark/unmark a connector favorite; `body` = FavoriteInputViewModel.",
+    ),
+    "halt_connector": Op(
+        "POST", "/public/v2/data_transfer/v3/connectors/{connector_type}/{id_}/halt",
+        path_params={"connector_type": "connector_type", "id_": "id_"},
+        required=("connector_type", "id_"),
+        write=True,
+        help="Mark a connector inactive (also halts its transfers).",
+    ),
+    "try_connector": Op(
+        "POST", "/public/v2/data_transfer/v3/connectors/{connector_type}/{id_}/try",
+        path_params={"connector_type": "connector_type", "id_": "id_"},
+        required=("connector_type", "id_"),
+        write=True,
+        help="Trial-test a connector (connection check; activates the connector on success).",
+    ),
+    "get_connector_logs": Op(
+        "GET", "/public/v2/data_transfer/v3/connectors/{connector_type}/{id_}/try/logs",
+        path_params={"connector_type": "connector_type", "id_": "id_"},
+        required=("connector_type", "id_"),
+        kind="log",
+        help="Get connector activation (try) logs.",
+    ),
+    "delete_connectors": Op(
+        "DELETE", "/public/v2/data_transfer/v2/connectors",
+        query_params={"ids": "ids"},
+        required=("ids",),
+        write=True,
+        confirm=True,
+        help="Delete connectors by id (ids repeat).",
+    ),
+    "list_sources": Op(
+        "GET", "/public/v2/data_transfer/v2/connectors/sources",
+        kind="list",
+        help="List schemas of all connectors usable as a source.",
+    ),
+    # --- transfers (v4 reads, v5 writes, v2 delete) -----------------------
+    "list_transfers": Op(
+        "GET", "/public/v2/data_transfer/v4/transfer",
+        kind="list",
+        help="List all transfers in the workspace.",
+    ),
+    "get_transfer": Op(
+        "GET", "/public/v2/data_transfer/v4/transfer/{transfer_id}",
+        path_params={"transfer_id": "transfer_id"},
+        required=("transfer_id",),
+        help="Get one transfer by id.",
+    ),
+    "create_transfer": Op(
+        "POST", "/public/v2/data_transfer/v5/transfer",
+        body_field="body",
+        required=("body",),
+        write=True,
+        help="Create a transfer; `body` = TransferInputViewModelV5 (runs now unless crontab is set).",
+    ),
+    "update_transfer": Op(
+        "POST", "/public/v2/data_transfer/v5/transfer/{id_}",
+        path_params={"id_": "id_"},
+        body_field="body",
+        required=("id_", "body"),
+        write=True,
+        help="Update a transfer by id; `body` = TransferInputViewModelV5.",
+    ),
+    "fav_transfer": Op(
+        "POST", "/public/v2/data_transfer/v4/transfer/{id_}/fav",
+        path_params={"id_": "id_"},
+        body_field="body",
+        required=("id_", "body"),
+        write=True,
+        help="Mark/unmark a transfer favorite; `body` = FavoriteInputViewModel.",
+    ),
+    "switch_transfer": Op(
+        "POST", "/public/v2/data_transfer/v4/transfer/{id_}/switch",
+        path_params={"id_": "id_"},
+        body_field="body",
+        required=("id_", "body"),
+        write=True,
+        help="Toggle a transfer active/inactive; `body` = StateInputViewModel.",
+    ),
+    "delete_transfers": Op(
+        "DELETE", "/public/v2/data_transfer/v2/transfer",
+        query_params={"ids": "ids"},
+        required=("ids",),
+        write=True,
+        confirm=True,
+        help="Delete transfers by id (ids repeat).",
+    ),
+    # --- history (v3 reads, v2/v3 ops, v2 delete) -------------------------
+    "list_history": Op(
+        "GET", "/public/v2/data_transfer/v3/history",
+        query_params={
+            "transfer_id": "transfer_id",
+            "source_name": "source_name",
+            "page": "page",
+            "page_size": "page_size",
+        },
+        kind="list",
+        help="List transfer history (paginated; filter by transfer_id/source_name).",
+    ),
+    "get_history_status": Op(
+        "GET", "/public/v2/data_transfer/v3/history/status/stream",
+        query_params={"ids": "ids"},
+        required=("ids",),
+        kind="list",
+        help="Get history entries by id (ids repeat).",
+    ),
+    "get_event_logs": Op(
+        "GET", "/public/v2/data_transfer/v2/events/list",
+        query_params={
+            "offset": "offset",
+            "limit": "limit",
+            "history_id": "history_id",
+            "transfer_id": "transfer_id",
+            "event_type": "event_type",
+            "event_category": "event_category",
+            "as_enum_resolved": "as_enum_resolved",
+        },
+        required=("offset", "limit"),
+        kind="list",
+        help="Get transfer event logs (needs offset+limit and either history_id or transfer_id).",
+    ),
+    "cancel_history": Op(
+        "POST", "/public/v2/data_transfer/v2/history/cancel",
+        body_field="body",
+        required=("body",),
+        write=True,
+        confirm=True,
+        help="Cancel a transfer; `body` = CancelInput.",
+    ),
+    "rerun_history": Op(
+        "POST", "/public/v2/data_transfer/v3/history/rerun",
+        body_field="body",
+        required=("body",),
+        write=True,
+        help="Restart a transfer; `body` = RerunInput.",
+    ),
+    "fav_history": Op(
+        "POST", "/public/v2/data_transfer/v2/history/{id_}/fav",
+        path_params={"id_": "id_"},
+        body_field="favorite",
+        required=("id_", "favorite"),
+        write=True,
+        help="Mark/unmark a history entry favorite; `favorite` is a bare boolean body.",
+    ),
+    "delete_history": Op(
+        "DELETE", "/public/v2/data_transfer/v2/history",
+        query_params={"ids": "ids"},
+        required=("ids",),
+        write=True,
+        confirm=True,
+        help="Delete history entries by id (ids repeat).",
+    ),
+}
+
+PARAMS: list[Param] = [
+    Param("id_", str, "Entity UUID (connector/transfer/history id, depending on action)."),
+    Param("transfer_id", str, "Transfer UUID (get_transfer path; filter for list_history/event logs)."),
+    Param(
+        "connector_type",
+        str,
+        "Connector type. Known values: postgresql, mssql, mysql, clickhouse, "
+        "oracledb, s3amazon, s3custom, s3google, hdfs, s3mlspace, s3datahub, nfs, "
+        "nfsprivate, nfsshared, s3evolution. get_connector accepts the full set; "
+        "try/halt/fav/update only the custom subset (the first 8). "
+        "Discover live values via list_connectors.",
+    ),
+    Param("ids", list[str], "List of UUIDs (repeated query) for delete/status actions."),
+    Param("source_name", str, "Filter history by source_name (list_history)."),
+    Param(
+        "page",
+        int,
+        "Page number (list_history, list_connectors, 1-based). NOTE for list_connectors "
+        "the response is a bare array with no total_count, so a short/empty page does "
+        "NOT prove the end of the list; list_history instead returns {data, total_count}.",
+    ),
+    Param("page_size", int, "Page size (list_history, list_connectors)."),
+    Param("offset", int, "Pagination offset (get_event_logs)."),
+    Param("limit", int, "Pagination limit (get_event_logs)."),
+    Param("history_id", str, "History UUID filter (get_event_logs)."),
+    Param("event_type", int, "Event type filter (get_event_logs)."),
+    Param("event_category", int, "Event category filter (get_event_logs)."),
+    Param("as_enum_resolved", bool, "Resolve enum codes to names (get_event_logs)."),
+    Param(
+        "favorite",
+        bool,
+        "Favorite flag for fav_history (sent as a bare boolean body, not an object).",
+    ),
+    Param(
+        "body",
+        dict,
+        "JSON body. create_connector=ConnectorInputViewModelV3; "
+        "update_connector=ConnectorUpdateViewModelV3; "
+        "fav_connector/fav_transfer=FavoriteInputViewModel {favorite}; "
+        "create_transfer/update_transfer=TransferInputViewModelV5; "
+        "switch_transfer=StateInputViewModel; cancel_history=CancelInput; "
+        "rerun_history=RerunInput {history_id}.",
+    ),
+]
+
+DOMAIN = DomainTool(
+    domain="data_transfer",
+    name="mlspace_data_transfer",
+    title="MLSpace Data Transfer",
+    summary=(
+        "Manage MLSpace data-transfer connectors, transfers, and migration history "
+        "(create/update/try connectors, run/cancel/rerun transfers, inspect logs)."
+    ),
+    actions=ACTIONS,
+    params=PARAMS,
+)