kumoai 2.14.0.dev202601011731__cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

kumoai/client/client.py

@@ -0,0 +1,223 @@
+from typing import TYPE_CHECKING, Any, Optional
+
+import requests
+from requests.adapters import HTTPAdapter
+from urllib3.util import Retry
+
+from kumoai.client.endpoints import Endpoint, HTTPMethod
+
+if TYPE_CHECKING:
+    from kumoai.client.connector import ConnectorAPI
+    from kumoai.client.graph import GraphAPI
+    from kumoai.client.jobs import (
+        ArtifactExportJobAPI,
+        BaselineJobAPI,
+        BatchPredictionJobAPI,
+        DistillationJobAPI,
+        GeneratePredictionTableJobAPI,
+        GenerateTrainTableJobAPI,
+        LLMJobAPI,
+        TrainingJobAPI,
+    )
+    from kumoai.client.online import OnlineServingEndpointAPI
+    from kumoai.client.pquery import PQueryAPI
+    from kumoai.client.source_table import SourceTableAPI
+    from kumoai.client.table import TableAPI
+
+API_VERSION = 'v1'
+
+
+class KumoClient:
+    def __init__(
+        self,
+        url: str,
+        api_key: Optional[str],
+        spcs_token: Optional[str] = None,
+        verify_ssl: bool = True,
+    ) -> None:
+        r"""Creates a client against the Kumo public API, provided a URL of
+        the endpoint and an authentication token.
+
+        Args:
+            url: the public API endpoint URL.
+            api_key: the public API authentication token.
+            spcs_token: the SPCS token used for authentication to access the
+                Kumo API endpoint.
+            verify_ssl: whether to verify SSL certificates. Set to False to
+                skip SSL certificate verification (equivalent to curl -k).
+        """
+        self._url = url
+        self._api_url = f"{url}/{API_VERSION}"
+        self._api_key = api_key
+        self._spcs_token = spcs_token
+        self._verify_ssl = verify_ssl
+
+        retry_strategy = Retry(
+            total=10,  # Maximum number of retries
+            connect=3,  # How many connection-related errors to retry on
+            read=3,  # How many times to retry on read errors
+            status=5,  # How many times to retry on bad status codes (below)
+            # Status codes to retry on.
+            status_forcelist=[408, 429, 500, 502, 503, 504],
+            # Exponential backoff factor: 2, 4, 8, 16 seconds delay)
+            backoff_factor=2.0,
+        )
+        http_adapter = HTTPAdapter(max_retries=retry_strategy)
+        session = requests.Session()
+        session.mount('http://', http_adapter)
+        session.mount('https://', http_adapter)
+        self._session = session
+        if self._api_key:
+            self._session.headers.update({"X-API-Key": self._api_key})
+        elif self._spcs_token:
+            self._session.headers.update(
+                {'Authorization': f'Snowflake Token={self._spcs_token}'})
+
+    def authenticate(self) -> None:
+        """Raises an exception if authentication fails."""
+        try:
+            self._session.get(self._url + '/v1/connectors',
+                              verify=self._verify_ssl).raise_for_status()
+        except Exception:
+            raise ValueError(
+                "Client authentication failed. Please check if you "
+                "have a valid API key/credentials.")
+
+    def set_spcs_token(self, spcs_token: str) -> None:
+        r"""Sets the SPCS token for the client and updates the session
+        headers.
+        """
+        self._spcs_token = spcs_token
+        self._session.headers.update(
+            {'Authorization': f'Snowflake Token={self._spcs_token}'})
+
+    @property
+    def artifact_export_api(self) -> 'ArtifactExportJobAPI':
+        r"""Returns the artifact export API."""
+        from kumoai.client.jobs import ArtifactExportJobAPI
+        return ArtifactExportJobAPI(self)
+
+    @property
+    def connector_api(self) -> 'ConnectorAPI':
+        r"""Returns the typed connector API."""
+        from kumoai.client.connector import ConnectorAPI
+        return ConnectorAPI(self)
+
+    @property
+    def source_table_api(self) -> 'SourceTableAPI':
+        r"""Returns the typed source table API."""
+        from kumoai.client.source_table import SourceTableAPI
+        return SourceTableAPI(self)
+
+    @property
+    def table_api(self) -> 'TableAPI':
+        r"""Returns the typed Kumo Table (snapshot) API."""
+        from kumoai.client.table import TableAPI
+        return TableAPI(self)
+
+    @property
+    def graph_api(self) -> 'GraphAPI':
+        r"""Returns the typed Graph (metadata and snapshot) API."""
+        from kumoai.client.graph import GraphAPI
+        return GraphAPI(self)
+
+    @property
+    def pquery_api(self) -> 'PQueryAPI':
+        r"""Returns the typed Predictive Query API."""
+        from kumoai.client.pquery import PQueryAPI
+        return PQueryAPI(self)
+
+    @property
+    def training_job_api(self) -> 'TrainingJobAPI':
+        r"""Returns the typed Training Job API."""
+        from kumoai.client.jobs import TrainingJobAPI
+        return TrainingJobAPI(self)
+
+    @property
+    def distillation_job_api(self) -> 'DistillationJobAPI':
+        from kumoai.client.jobs import DistillationJobAPI
+        return DistillationJobAPI(self)
+
+    @property
+    def batch_prediction_job_api(self) -> 'BatchPredictionJobAPI':
+        from kumoai.client.jobs import BatchPredictionJobAPI
+        return BatchPredictionJobAPI(self)
+
+    @property
+    def generate_train_table_job_api(self) -> 'GenerateTrainTableJobAPI':
+        r"""Returns the typed Generate-Train-Table Job API."""
+        from kumoai.client.jobs import GenerateTrainTableJobAPI
+        return GenerateTrainTableJobAPI(self)
+
+    @property
+    def generate_prediction_table_job_api(
+            self) -> 'GeneratePredictionTableJobAPI':
+        from kumoai.client.jobs import GeneratePredictionTableJobAPI
+        return GeneratePredictionTableJobAPI(self)
+
+    @property
+    def llm_job_api(self) -> 'LLMJobAPI':
+        from kumoai.client.jobs import LLMJobAPI
+        return LLMJobAPI(self)
+
+    @property
+    def baseline_job_api(self) -> 'BaselineJobAPI':
+        r"""Returns the typed Training Job API."""
+        from kumoai.client.jobs import BaselineJobAPI
+        return BaselineJobAPI(self)
+
+    @property
+    def online_serving_endpoint_api(self) -> 'OnlineServingEndpointAPI':
+        from kumoai.client.online import OnlineServingEndpointAPI
+        return OnlineServingEndpointAPI(self)
+
+    def _request(self, endpoint: Endpoint, **kwargs: Any) -> requests.Response:
+        r"""Send a HTTP request to the specified endpoint."""
+        endpoint_str = endpoint.get_path()
+        if endpoint.method == HTTPMethod.GET:
+            return self._get(endpoint_str, **kwargs)
+        elif endpoint.method == HTTPMethod.POST:
+            return self._post(endpoint_str, **kwargs)
+        elif endpoint.method == HTTPMethod.PATCH:
+            return self._patch(endpoint_str, **kwargs)
+        elif endpoint.method == HTTPMethod.DELETE:
+            return self._delete(endpoint_str, **kwargs)
+        else:
+            raise ValueError(f"Unsupported HTTP method: {endpoint.method}")
+
+    def _get(self, endpoint: str, **kwargs: Any) -> requests.Response:
+        r"""Send a GET request to the specified endpoint, with keyword
+        arguments, returned objects, and exceptions raised corresponding to
+        :meth:`requests.Session.get`.
+        """
+        url = self._format_endpoint_url(endpoint)
+        return self._session.get(url=url, verify=self._verify_ssl, **kwargs)
+
+    def _post(self, endpoint: str, **kwargs: Any) -> requests.Response:
+        r"""Send a POST request to the specified endpoint, with keyword
+        arguments, returned objects, and exceptions raised corresponding to
+        :meth:`requests.Session.post`.
+        """
+        url = self._format_endpoint_url(endpoint)
+        return self._session.post(url=url, verify=self._verify_ssl, **kwargs)
+
+    def _patch(self, endpoint: str, **kwargs: Any) -> requests.Response:
+        r"""Send a PATCH request to the specified endpoint, with keyword
+        arguments, returned objects, and exceptions raised corresponding to
+        :meth:`requests.Session.patch`.
+        """
+        url = self._format_endpoint_url(endpoint)
+        return self._session.patch(url=url, verify=self._verify_ssl, **kwargs)
+
+    def _delete(self, endpoint: str, **kwargs: Any) -> requests.Response:
+        r"""Send a DELETE request to the specified endpoint, with keyword
+        arguments, returned objects, and exceptions raised corresponding to
+        :meth:`requests.Session.delete`.
+        """
+        url = self._format_endpoint_url(endpoint)
+        return self._session.delete(url=url, verify=self._verify_ssl, **kwargs)
+
+    def _format_endpoint_url(self, endpoint: str) -> str:
+        if endpoint[0] == "/":
+            endpoint = endpoint[1:]
+        return f"{self._api_url}/{endpoint}"

kumoai/client/connector.py

@@ -0,0 +1,110 @@
+from http import HTTPStatus
+from typing import List, Optional
+
+from kumoapi.data_source import (
+    CompleteFileUploadRequest,
+    ConnectorResponse,
+    CreateConnectorArgs,
+    DataSourceType,
+    DeleteUploadedFileRequest,
+    StartFileUploadRequest,
+    StartFileUploadResponse,
+)
+from kumoapi.json_serde import to_json_dict
+
+from kumoai.client import KumoClient
+from kumoai.client.endpoints import ConnectorEndpoints
+from kumoai.client.utils import parse_response, raise_on_error
+from kumoai.exceptions import HTTPException
+
+
+class ConnectorAPI:
+    r"""Typed API definition for Kumo connectors."""
+    def __init__(self, client: KumoClient) -> None:
+        self._client = client
+
+    def start_file_upload(
+            self, req: StartFileUploadRequest) -> StartFileUploadResponse:
+        res = self._client._request(
+            ConnectorEndpoints.start_file_upload,
+            json=to_json_dict(req, insecure=True),
+        )
+        raise_on_error(res)
+        return parse_response(StartFileUploadResponse, res)
+
+    def delete_file_upload(self, req: DeleteUploadedFileRequest) -> None:
+        res = self._client._request(
+            ConnectorEndpoints.delete_uploaded_file,
+            json=to_json_dict(req, insecure=True),
+        )
+        raise_on_error(res)
+
+    def complete_file_upload(self, req: CompleteFileUploadRequest) -> None:
+        res = self._client._request(
+            ConnectorEndpoints.complete_file_upload,
+            json=to_json_dict(req, insecure=True),
+        )
+        raise_on_error(res)
+
+    def create(self, create_connector_args: CreateConnectorArgs) -> None:
+        r"""Creates a connector in Kumo."""
+        resp = self._client._request(
+            ConnectorEndpoints.create,
+            json=to_json_dict(create_connector_args, insecure=True),
+        )
+        raise_on_error(resp)
+
+    def create_if_not_exist(
+        self,
+        create_connector_args: CreateConnectorArgs,
+    ) -> bool:
+        r"""Creates a connector in Kumo if the connector does not exist.
+
+        Returns:
+            :obj:`True` if the connector is newly created, :obj:`False`
+            otherwise.
+        """
+        _id = create_connector_args.config.name
+        existing_connector = self.get(_id)
+        if existing_connector:
+            if existing_connector.config == create_connector_args.config:
+                return False
+            raise HTTPException(
+                HTTPStatus.UNPROCESSABLE_ENTITY,
+                f"Connector {_id} already exists, but has a differing "
+                f"configuration. Input: {create_connector_args.config}, "
+                f"existing: {existing_connector.config}",
+            )
+        self.create(create_connector_args)
+        return existing_connector is None
+
+    def get(self, connector_id: str) -> Optional[ConnectorResponse]:
+        r"""Fetches a connector given its ID."""
+        resp = self._client._request(
+            ConnectorEndpoints.get.with_id(connector_id))
+        if resp.status_code == HTTPStatus.NOT_FOUND:
+            return None
+
+        raise_on_error(resp)
+        return parse_response(ConnectorResponse, resp)
+
+    def list(
+        self, data_source_type: Optional[DataSourceType] = None
+    ) -> List[ConnectorResponse]:
+        r"""Lists connectors for a given data source type."""
+        params = {
+            'data_source_type': data_source_type
+        } if data_source_type else {}
+
+        resp = self._client._request(ConnectorEndpoints.list, params=params)
+        raise_on_error(resp)
+        return parse_response(List[ConnectorResponse], resp)
+
+    def delete_if_exists(self, connector_id: str) -> bool:
+        r"""Deletes a connector if it exists."""
+        resp = self._client._request(
+            ConnectorEndpoints.delete.with_id(connector_id))
+        if resp.status_code == HTTPStatus.NOT_FOUND:
+            return False
+        raise_on_error(resp)
+        return True

kumoai/client/endpoints.py

@@ -0,0 +1,150 @@
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Final, Optional
+
+
+class HTTPMethod(Enum):
+    r"""HTTP methods supported by the API."""
+    GET = "GET"
+    POST = "POST"
+    DELETE = "DELETE"
+    PATCH = "PATCH"
+
+
+@dataclass(frozen=True)
+class Endpoint:
+    r"""Represents an API endpoint with its path and HTTP method."""
+    path: Optional[str] = field(default=None)
+    method: HTTPMethod = HTTPMethod.GET
+
+    def validate(self) -> None:
+        pass
+
+    def get_path(self) -> str:
+        if self.path is None:
+            # This should be in validate but is here for type checker
+            raise ValueError("Endpoint requires a path")
+        return self.path
+
+
+@dataclass(frozen=True)
+class IDEndpoint(Endpoint):
+    r"""Represents an API endpoint with an additional id/name in its path."""
+    template_path: Optional[str] = field(default=None)
+
+    def __post_init__(self) -> None:
+        if self.template_path is None:
+            raise ValueError("template_path must be set explicitly")
+        if "{id}" not in self.template_path:
+            raise ValueError("IDEndpoint path must contain '{{id}}': "
+                             f"got '{self.template_path}'")
+
+    def with_id(self, resource_id: str) -> 'IDEndpoint':
+        assert self.template_path is not None
+        return IDEndpoint(template_path=self.template_path,
+                          path=self.template_path.format(id=resource_id),
+                          method=self.method)
+
+    def validate(self) -> None:
+        if self.path is None:
+            raise ValueError(
+                "IDEndpoint requires with_id() to be called with a resource id"
+            )
+
+    @classmethod
+    def from_base(cls, base: str, method: HTTPMethod) -> 'IDEndpoint':
+        r"""Alternate constructor to build a template_path by appending /{id}
+        to base.
+        """
+        template = f"{base}/{{id}}"
+        return cls(template_path=template, method=method)
+
+
+class ConnectorEndpoints:
+    BASE: Final[str] = "/connectors"
+
+    create = Endpoint(BASE, HTTPMethod.POST)
+    list = Endpoint(BASE, HTTPMethod.GET)
+    get = IDEndpoint.from_base(BASE, HTTPMethod.GET)
+    delete = IDEndpoint.from_base(BASE, HTTPMethod.DELETE)
+
+    start_file_upload = Endpoint(f"{BASE}/utils/start_file_upload",
+                                 HTTPMethod.POST)
+    delete_uploaded_file = Endpoint(f"{BASE}/utils/delete_uploaded_file",
+                                    HTTPMethod.POST)
+    complete_file_upload = Endpoint(f"{BASE}/utils/complete_file_upload",
+                                    HTTPMethod.POST)
+
+
+class PQueryEndpoints:
+    BASE: Final[str] = '/predictive_queries'
+
+    create = Endpoint(BASE, HTTPMethod.POST)
+    list = Endpoint(BASE, HTTPMethod.GET)
+    get = IDEndpoint.from_base(BASE, HTTPMethod.GET)
+    delete = IDEndpoint.from_base(BASE, HTTPMethod.DELETE)
+
+    infer_task_type = Endpoint(f"{BASE}/infer_task_type", HTTPMethod.POST)
+    validate = Endpoint(f"{BASE}/validate", HTTPMethod.POST)
+    suggest_training_table_plan = Endpoint(
+        f"{BASE}/train_table_generation_plan", HTTPMethod.POST)
+    suggest_model_plan = Endpoint(f"{BASE}/model_plan", HTTPMethod.POST)
+
+
+class SourceTableEndpoints:
+    BASE: Final[str] = "/source_tables"
+
+    get = IDEndpoint.from_base(BASE, HTTPMethod.GET)
+
+    list_tables = Endpoint(f"{BASE}/list_tables", HTTPMethod.POST)
+    validate_table = Endpoint(f"{BASE}/validate_table", HTTPMethod.POST)
+    get_table_data = Endpoint(f"{BASE}/get_table_data", HTTPMethod.POST)
+    get_table_config = Endpoint(f"{BASE}/get_table_config", HTTPMethod.POST)
+
+
+class GraphEndpoints:
+    BASE: Final[str] = "/graphs"
+    SNAPSHOTS: Final[str] = "/graphsnapshots"
+
+    create = Endpoint(BASE, HTTPMethod.POST)
+    get = IDEndpoint.from_base(BASE, HTTPMethod.GET)
+
+    validate = Endpoint(f"{BASE}/validate", HTTPMethod.POST)
+    infer_links = Endpoint(f"{BASE}/infer_links", HTTPMethod.POST)
+    create_snapshot = Endpoint(SNAPSHOTS, HTTPMethod.POST)
+    get_snapshot = IDEndpoint.from_base(SNAPSHOTS, HTTPMethod.GET)
+    get_edge_stats = IDEndpoint.from_base(f"{SNAPSHOTS}/edge_health",
+                                          method=HTTPMethod.GET)
+
+
+class OnlineServingEndpoints:
+    BASE: Final[str] = "/online_serving_endpoints"
+
+    create = Endpoint(BASE, HTTPMethod.POST)
+    get = IDEndpoint.from_base(BASE, HTTPMethod.GET)
+    list = Endpoint(BASE, HTTPMethod.GET)
+    update = IDEndpoint.from_base(BASE, HTTPMethod.PATCH)
+    delete = IDEndpoint.from_base(BASE, HTTPMethod.DELETE)
+
+
+class TableEndpoints:
+    BASE: Final[str] = "/tables"
+    SNAPSHOTS: Final[str] = "/tablesnapshots"
+
+    create = Endpoint(BASE, HTTPMethod.POST)
+    get = IDEndpoint.from_base(BASE, HTTPMethod.GET)
+
+    create_snapshot = Endpoint(SNAPSHOTS, HTTPMethod.POST)
+    get_snapshot = IDEndpoint.from_base(SNAPSHOTS, HTTPMethod.GET)
+    validate = Endpoint(f"{BASE}/validate", HTTPMethod.POST)
+    infer_metadata = Endpoint(f"{BASE}/infer_metadata", HTTPMethod.POST)
+
+
+class RFMEndpoints:
+    BASE: Final[str] = "/rfm"
+
+    predict = Endpoint(f"{BASE}/predict", HTTPMethod.POST)
+    explain = Endpoint(f"{BASE}/explain", HTTPMethod.POST)
+    evaluate = Endpoint(f"{BASE}/evaluate", HTTPMethod.POST)
+    validate_query = Endpoint(f"{BASE}/validate_query", HTTPMethod.POST)
+    parse_query = Endpoint(f"{BASE}/parse_query", HTTPMethod.POST)

kumoai/client/graph.py

@@ -0,0 +1,120 @@
+from http import HTTPStatus
+from typing import Any, Dict, Optional
+
+from kumoapi.common import ValidationResponse
+from kumoapi.data_snapshot import GraphSnapshotID, GraphSnapshotResource
+from kumoapi.graph import (
+    EdgeHealthResponse,
+    GraphDefinition,
+    GraphResource,
+    GraphValidationRequest,
+)
+from kumoapi.json_serde import to_json_dict
+
+from kumoai.client import KumoClient
+from kumoai.client.endpoints import GraphEndpoints
+from kumoai.client.utils import (
+    parse_id_response,
+    parse_response,
+    raise_on_error,
+)
+
+GraphID = str
+
+
+class GraphAPI:
+    r"""Typed API definition for Kumo graph definition."""
+    def __init__(self, client: KumoClient) -> None:
+        self._client = client
+
+    def create_graph(
+        self,
+        graph_def: GraphDefinition,
+        *,
+        name_alias: Optional[str] = None,
+        force_rename: bool = False,
+    ) -> GraphID:
+        r"""Creates a Graph (metadata definition) resource object in Kumo."""
+        params: Dict[str, Any] = {'force_rename': force_rename}
+        if name_alias:
+            params['name_alias'] = name_alias
+        resp = self._client._request(
+            GraphEndpoints.create,
+            params=params,
+            json=to_json_dict(graph_def),
+        )
+        raise_on_error(resp)
+        return parse_id_response(resp)
+
+    def get_graph_if_exists(
+        self,
+        graph_id_or_name: str,
+    ) -> Optional[GraphResource]:
+        resp = self._client._request(
+            GraphEndpoints.get.with_id(graph_id_or_name))
+        if resp.status_code == HTTPStatus.NOT_FOUND:
+            return None
+
+        raise_on_error(resp)
+        return parse_response(GraphResource, resp)
+
+    def create_snapshot(
+        self,
+        graph_id: GraphID,
+        *,
+        refresh_source: bool = False,
+    ) -> GraphSnapshotID:
+        params: Dict[str, Any] = {
+            'graph_id': graph_id,
+            'refresh_source': refresh_source
+        }
+        resp = self._client._request(GraphEndpoints.create_snapshot,
+                                     params=params)
+        raise_on_error(resp)
+        return GraphSnapshotID(parse_id_response(resp))
+
+    def get_snapshot(
+        self,
+        snapshot_id: GraphSnapshotID,
+    ) -> GraphSnapshotResource:
+        resp = self._client._request(
+            GraphEndpoints.get_snapshot.with_id(snapshot_id))
+        raise_on_error(resp)
+        return parse_response(GraphSnapshotResource, resp)
+
+    def get_snapshot_if_exists(
+        self,
+        snapshot_id: GraphSnapshotID,
+    ) -> Optional[GraphSnapshotResource]:
+        resp = self._client._request(
+            GraphEndpoints.get_snapshot.with_id(snapshot_id))
+        if resp.status_code == HTTPStatus.NOT_FOUND:
+            return None
+
+        raise_on_error(resp)
+        return parse_response(GraphSnapshotResource, resp)
+
+    def get_edge_stats(
+        self,
+        graph_snapshot_id: GraphSnapshotID,
+    ) -> EdgeHealthResponse:
+        r"""Fetches edge statistics given a snapshot id"""
+        resp = self._client._request(
+            GraphEndpoints.get_edge_stats.with_id(graph_snapshot_id))
+        raise_on_error(resp)
+        return parse_response(EdgeHealthResponse, resp)
+
+    def validate_graph(
+        self,
+        request: GraphValidationRequest,
+    ) -> ValidationResponse:
+        response = self._client._request(GraphEndpoints.validate,
+                                         json=to_json_dict(request))
+        raise_on_error(response)
+        return parse_response(ValidationResponse, response)
+
+    def infer_links(self, graph: GraphDefinition) -> GraphDefinition:
+        resp = self._client._request(GraphEndpoints.infer_links,
+                                     json=to_json_dict(graph))
+        raise_on_error(resp)
+        return parse_response(GraphDefinition, resp)