kumoai 2.9.0.dev202509081831__cp312-cp312-win_amd64.whl → 2.13.0.dev202511201731__cp312-cp312-win_amd64.whl

kumoai/__init__.py

@@ -184,15 +184,12 @@ def init(
         snowflake_credentials
     ) if not api_key and snowflake_credentials else None
     client = KumoClient(url=url, api_key=api_key, spcs_token=spcs_token)
-    if client.authenticate():
-        global_state._url = client._url
-        global_state._api_key = client._api_key
-        global_state._snowflake_credentials = snowflake_credentials
-        global_state._spcs_token = client._spcs_token
-        global_state._snowpark_session = snowpark_session
-    else:
-        raise ValueError("Client authentication failed. Please check if you "
-                         "have a valid API key.")
+    client.authenticate()
+    global_state._url = client._url
+    global_state._api_key = client._api_key
+    global_state._snowflake_credentials = snowflake_credentials
+    global_state._spcs_token = client._spcs_token
+    global_state._snowpark_session = snowpark_session
 
     if not api_key and snowflake_credentials:
         # Refresh token every 10 minutes (expires in 1 hour):
@@ -200,9 +197,11 @@ def init(
 
     logger = logging.getLogger('kumoai')
     log_level = logging.getLevelName(logger.getEffectiveLevel())
+
     logger.info(
-        "Successfully initialized the Kumo SDK against deployment %s, with "
-        "log level %s.", url, log_level)
+        f"Successfully initialized the Kumo SDK (version {__version__}) "
+        f"against deployment {url}, with "
+        f"log level {log_level}.")
 
 
 def set_log_level(level: str) -> None:

kumoai/_version.py

@@ -1 +1 @@
-__version__ = '2.9.0.dev202509081831'
+__version__ = '2.13.0.dev202511201731'

kumoai/client/client.py

@@ -20,7 +20,6 @@ if TYPE_CHECKING:
     )
     from kumoai.client.online import OnlineServingEndpointAPI
     from kumoai.client.pquery import PQueryAPI
-    from kumoai.client.rfm import RFMAPI
     from kumoai.client.source_table import SourceTableAPI
     from kumoai.client.table import TableAPI
 
@@ -33,6 +32,7 @@ class KumoClient:
         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.
@@ -42,11 +42,14 @@ class KumoClient:
             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
@@ -69,11 +72,15 @@ class KumoClient:
             self._session.headers.update(
                 {'Authorization': f'Snowflake Token={self._spcs_token}'})
 
-    def authenticate(self) -> bool:
-        r"""Raises an exception if authentication fails. Succeeds if the
-        client is properly formed.
-        """
-        return self._session.get(f"{self._url}/v1/connectors").ok
+    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
@@ -158,12 +165,6 @@ class KumoClient:
         from kumoai.client.online import OnlineServingEndpointAPI
         return OnlineServingEndpointAPI(self)
 
-    @property
-    def rfm_api(self) -> 'RFMAPI':
-        r"""Returns the typed RFM API."""
-        from kumoai.client.rfm import RFMAPI
-        return RFMAPI(self)
-
     def _request(self, endpoint: Endpoint, **kwargs: Any) -> requests.Response:
         r"""Send a HTTP request to the specified endpoint."""
         endpoint_str = endpoint.get_path()
@@ -184,7 +185,7 @@ class KumoClient:
         :meth:`requests.Session.get`.
         """
         url = self._format_endpoint_url(endpoint)
-        return self._session.get(url=url, **kwargs)
+        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
@@ -192,7 +193,7 @@ class KumoClient:
         :meth:`requests.Session.post`.
         """
         url = self._format_endpoint_url(endpoint)
-        return self._session.post(url=url, **kwargs)
+        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
@@ -200,7 +201,7 @@ class KumoClient:
         :meth:`requests.Session.patch`.
         """
         url = self._format_endpoint_url(endpoint)
-        return self._session.patch(url=url, **kwargs)
+        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
@@ -208,7 +209,7 @@ class KumoClient:
         :meth:`requests.Session.delete`.
         """
         url = self._format_endpoint_url(endpoint)
-        return self._session.delete(url=url, **kwargs)
+        return self._session.delete(url=url, verify=self._verify_ssl, **kwargs)
 
     def _format_endpoint_url(self, endpoint: str) -> str:
         if endpoint[0] == "/":

kumoai/client/endpoints.py

@@ -147,3 +147,4 @@ class RFMEndpoints:
     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/rfm.py

@@ -1,6 +1,11 @@
+from typing import Any
+
 from kumoapi.json_serde import to_json_dict
 from kumoapi.rfm import (
     RFMEvaluateResponse,
+    RFMExplanationResponse,
+    RFMParseQueryRequest,
+    RFMParseQueryResponse,
     RFMPredictResponse,
     RFMValidateQueryRequest,
     RFMValidateQueryResponse,
@@ -25,28 +30,35 @@ class RFMAPI:
         Returns:
             RFMPredictResponse containing the predictions
         """
-        # Send binary data to the predict endpoint
         response = self._client._request(
-            RFMEndpoints.predict, data=request,
-            headers={'Content-Type': 'application/x-protobuf'})
+            RFMEndpoints.predict,
+            data=request,
+            headers={'Content-Type': 'application/x-protobuf'},
+        )
         raise_on_error(response)
         return parse_response(RFMPredictResponse, response)
 
-    def explain(self, request: bytes) -> RFMPredictResponse:
+    def explain(
+        self,
+        request: bytes,
+        skip_summary: bool = False,
+    ) -> RFMExplanationResponse:
         """Explain the RFM model on the given context.
 
         Args:
             request: The predict request as serialized protobuf.
+            skip_summary: Whether to skip generating a human-readable summary
+                of the explanation.
 
         Returns:
             RFMPredictResponse containing the explanations
         """
-        # Send binary data to the explain endpoint
+        params: dict[str, Any] = {'generate_summary': not skip_summary}
         response = self._client._request(
-            RFMEndpoints.explain, data=request,
+            RFMEndpoints.explain, data=request, params=params,
             headers={'Content-Type': 'application/x-protobuf'})
         raise_on_error(response)
-        return parse_response(RFMPredictResponse, response)
+        return parse_response(RFMExplanationResponse, response)
 
     def evaluate(self, request: bytes) -> RFMEvaluateResponse:
         """Evaluate the RFM model on the given context.
@@ -57,7 +69,6 @@ class RFMAPI:
         Returns:
             RFMEvaluateResponse containing the computed metrics
         """
-        # Send binary data to the evaluate endpoint
         response = self._client._request(
             RFMEndpoints.evaluate, data=request,
             headers={'Content-Type': 'application/x-protobuf'})
@@ -81,3 +92,21 @@ class RFMAPI:
                                          json=to_json_dict(request))
         raise_on_error(response)
         return parse_response(RFMValidateQueryResponse, response)
+
+    def parse_query(
+        self,
+        request: RFMParseQueryRequest,
+    ) -> RFMParseQueryResponse:
+        """Validate a predictive query against a graph.
+
+        Args:
+            request: The request object containing
+                the query and graph definition
+
+        Returns:
+            RFMParseQueryResponse containing the QueryDefinition
+        """
+        response = self._client._request(RFMEndpoints.parse_query,
+                                         json=to_json_dict(request))
+        raise_on_error(response)
+        return parse_response(RFMParseQueryResponse, response)

kumoai/connector/file_upload_connector.py

@@ -1,7 +1,5 @@
-import os
 from typing import List
 
-from kumoapi.data_source import DeleteUploadedFileRequest
 from kumoapi.source_table import (
     DataSourceType,
     FileType,
@@ -13,15 +11,7 @@ from typing_extensions import override
 
 from kumoai import global_state
 from kumoai.connector.base import Connector
-from kumoai.connector.utils import (
-    CONNECTOR_ID_MAP,
-    MAX_PARTITION_SIZE,
-    MIN_PARTITION_SIZE,
-    _upload_partitioned_csv,
-    _upload_partitioned_parquet,
-    _upload_single_file,
-    logger,
-)
+from kumoai.connector.utils import delete_uploaded_table, upload_table
 
 
 class FileUploadConnector(Connector):
@@ -61,7 +51,6 @@ class FileUploadConnector(Connector):
     def name(self) -> str:
         return f'{self._file_type}_upload_connector'
 
-    @override
     @property
     def source_type(self) -> DataSourceType:
         return DataSourceType.S3
@@ -91,100 +80,92 @@ class FileUploadConnector(Connector):
         auto_partition: bool = True,
         partition_size_mb: int = 250,
     ) -> None:
-        r"""Synchronously uploads a table located on your
-        local machine to the Kumo data plane.
-
-        Tables uploaded in this way can be accessed with
-        this ``FileUploadConnector`` using the provided name,
-        for example: ``connector_obj["my_table"]``
-
-        For files larger than 1GB, the table will be automatically partitioned
-        into smaller chunks and uploaded with common prefix that allows
-        FileUploadConnector to union them when reading.
+        r"""Upload a table to Kumo from a local or remote path.
+
+        Supports ``s3://``, ``gs://``, ``abfs://``, ``abfss://``, and ``az://``
+
+        Tables uploaded this way can be accessed from this
+        ``FileUploadConnector`` using the provided name, e.g.,
+        ``connector_obj["my_table"]``.
+
+        Local files
+        -----------
+        - Accepts one ``.parquet`` or ``.csv`` file (must match this
+          connector’s ``file_type``).
+        - If the file is > 1 GiB and ``auto_partition=True``, it is split
+          into ~``partition_size_mb`` MiB parts and uploaded under a common
+          prefix so the connector can read them as one table.
+
+        Remote paths
+        ------------
+        - **Single file** (``.parquet``/``.csv``): validated and uploaded via
+          multipart PUT. Files > 1 GiB are rejected — re-shard to ~200 MiB
+          and upload the directory instead.
+        - **Directory**: must contain only one format (all Parquet or all CSV)
+          matching this connector’s ``file_type``. Files are validated
+          (consistent schema; CSV headers sanitized) and uploaded in parallel
+          with memory-safe budgeting.
 
         .. warning::
-            Uploaded tables must be single files, either in parquet or CSV
-            format(must match connector type).
-            Partitioned tables are not currently supported.
+           For local uploads, input must be a single CSV or Parquet file
+           (matching the connector type). For remote uploads, mixed
+           CSV/Parquet directories are not supported. Remote single files
+           larger than 1 GiB are not supported.
 
+        Examples:
+        ---------
         .. code-block:: python
 
             import kumoai
-            connector = kumoai.FileUploadConnector(file_type="parquet")
+            conn = kumoai.FileUploadConnector(file_type="parquet")
+
+            # Local: small file
+            conn.upload(name="users", path="/data/users.parquet")
 
-            # Upload a small table
-            connector.upload(name="users", path="/data/users.parquet")
+            # Local: large file (auto-partitions)
+            conn.upload(
+                name="txns",
+                path="/data/large_txns.parquet",
+            )
 
-            # Upload a large parquet table (will be automatically partitioned)
-            connector.upload(name="transactions",
-                        path="/data/large_transactions.parquet")
+            # Local: disable auto-partitioning (raises if > 1 GiB)
+            conn.upload(
+                name="users",
+                path="/data/users.parquet",
+                auto_partition=False,
+            )
 
-            # Disable auto-partitioning (will raise error for large files)
-            upload(name="users", path="/data/users.parquet",
-                        auto_partition=False)
+            # CSV connector
+            csv_conn = kumoai.FileUploadConnector(file_type="csv")
+            csv_conn.upload(name="sales", path="/data/sales.csv")
 
-            # Create a file upload connector for CSV files.
-            connectorCSV = kumoai.FileUploadConnector(file_type="csv")
+            # Remote: single file (<= 1 GiB)
+            conn.upload(name="logs", path="s3://bkt/path/logs.parquet")
 
-            # Upload a large CSV table (will be automatically partitioned)
-            connectorCSV.upload(name="sales", path="/data/large_sales.csv")
+            # Remote: directory of shards (uniform format)
+            csv_conn.upload(name="events", path="gs://mybkt/events_csv/")
 
         Args:
-            name: The name of the table to be uploaded. The uploaded table can
-                be accessed from the
-                :class:`~kumoai.connector.FileUploadConnector` with this name.
-            path: The full path of the table to be uploaded, on the local
-                machine. File Type must match the connector type.
-            auto_partition: Whether to automatically
-                partition large files (>1GB).
-                If False and file is >1GB, raises ValueError. Supports both
-                Parquet and CSV files.
-            partition_size_mb: The size of each partition in MB. Only used if
-                auto_partition is True.
+            name:
+                Table name to create in Kumo; access later via this connector.
+            path:
+                Local path or remote URL to a ``.parquet``/``.csv`` file or a
+                directory (uniform format). The format must match this
+                connector’s ``file_type``.
+            auto_partition:
+                Local-only. If ``True`` and the local file is > 1 GiB, split
+                into ~``partition_size_mb`` MiB parts.
+            partition_size_mb:
+                Local-only. Target partition size (100–1000 MiB) when
+                ``auto_partition`` is ``True``.
         """
-        # Validate file type matches connector type
-        if not path.lower().endswith("." + self._file_type):
-            raise ValueError(f"File {path} must match connector path type: "
-                             f"{self._file_type}.")
-
-        # Validate file type
-        if not (path.endswith(".parquet") or path.endswith(".csv")):
-            raise ValueError(f"Path {path} must be either a CSV or Parquet "
-                             f"file. Partitioned data is not currently "
-                             f"supported.")
-
-        file_size = os.path.getsize(path)
-
-        # Route based on file size
-        if file_size < MAX_PARTITION_SIZE:
-            return _upload_single_file(name, path)
-
-        if not auto_partition:
-            raise ValueError(f"File {path} is {file_size / (1024**3):.2f}GB, "
-                             f"which exceeds the 1GB limit. Enable "
-                             f"auto_partition=True to automatically partition "
-                             f"large files.")
-
-        # Partition and upload large files
-        partition_size = partition_size_mb * 1024**2
-        if (partition_size > MAX_PARTITION_SIZE
-                or partition_size < MIN_PARTITION_SIZE):
-            raise ValueError(f"Partition size {partition_size_mb}MB must be "
-                             f"between {MIN_PARTITION_SIZE / 1024**2}MB and "
-                             f"{MAX_PARTITION_SIZE / 1024**2}MB.")
-
-        logger.info(
-            "File %s is large with size %s, partitioning for upload...", path,
-            file_size)
-        if path.endswith('.parquet'):
-            _upload_partitioned_parquet(name, path, partition_size)
-        else:
-            _upload_partitioned_csv(name, path, partition_size)
+        upload_table(name=name, path=path, auto_partition=auto_partition,
+                     partition_size_mb=partition_size_mb,
+                     file_type=self._file_type)
 
     def delete(
         self,
         name: str,
-        file_type: str,
     ) -> None:
         r"""Synchronously deletes a previously uploaded table from the Kumo
         data plane.
@@ -194,27 +175,15 @@ class FileUploadConnector(Connector):
             # Assume we have uploaded a `.parquet` table named `users`, and a
             # `FileUploadConnector` has been created called `connector`, and
             # we want to delete this table from Kumo:
-            connector.delete(name="users", file_type="parquet")
+            connector.delete(name="users")
 
         Args:
             name: The name of the table to be deleted. This table must have
                 previously been uploaded with a call to
                 :meth:`~kumoai.connector.FileUploadConnector.upload`.
-            file_type: The file type of the table to be deleted; this can
-                either be :obj:`"parquet"` or :obj:`"csv"`, and must match the
-                connector file_type.
         """
-        if file_type.lower() != self._file_type:
-            raise ValueError(f"File type {file_type} does not match "
-                             f"connector file type {self._file_type}.")
-
         if not self.has_table(name):
             raise ValueError(f"The table '{name}' does not exist in {self}. "
                              f"Please check the existence of the source data.")
 
-        req = DeleteUploadedFileRequest(
-            source_table_name=name,
-            connector_id=CONNECTOR_ID_MAP[file_type],
-        )
-        global_state.client.connector_api.delete_file_upload(req)
-        logger.info("Successfully deleted table %s from Kumo.", name)
+        delete_uploaded_table(name, self._file_type)