kumoai 2.12.0.dev202511071730__cp310-cp310-win_amd64.whl → 2.13.0.dev202511261731__cp310-cp310-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):

kumoai/_version.py

@@ -1 +1 @@
-__version__ = '2.12.0.dev202511071730'
+__version__ = '2.13.0.dev202511261731'

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
 
@@ -73,12 +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",
-                                 verify=self._verify_ssl).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
@@ -163,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()

kumoai/client/rfm.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 from kumoapi.json_serde import to_json_dict
 from kumoapi.rfm import (
     RFMEvaluateResponse,
@@ -28,25 +30,32 @@ 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) -> RFMExplanationResponse:
+    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(RFMExplanationResponse, response)
@@ -60,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'})

kumoai/experimental/rfm/__init__.py

@@ -31,14 +31,122 @@ Please create a feature request at 'https://github.com/kumo-ai/kumo-rfm'."""
 
     raise RuntimeError(_msg) from e
 
-from typing import Optional, Dict
+from dataclasses import dataclass
+from enum import Enum
+import ipaddress
+import logging
+import re
+import socket
+import threading
+from typing import Optional, Dict, Tuple
 import os
+from urllib.parse import urlparse
 import kumoai
+from kumoai.client.client import KumoClient
+from .sagemaker import (KumoClient_SageMakerAdapter,
+                        KumoClient_SageMakerProxy_Local)
 from .local_table import LocalTable
 from .local_graph import LocalGraph
-from .rfm import KumoRFM
+from .rfm import ExplainConfig, Explanation, KumoRFM
 from .authenticate import authenticate
 
+logger = logging.getLogger('kumoai_rfm')
+
+
+def _is_local_address(host: str | None) -> bool:
+    """Return True if the hostname/IP refers to the local machine."""
+    if not host:
+        return False
+    try:
+        infos = socket.getaddrinfo(host, None)
+        for _, _, _, _, sockaddr in infos:
+            ip = sockaddr[0]
+            ip_obj = ipaddress.ip_address(ip)
+            if ip_obj.is_loopback or ip_obj.is_unspecified:
+                return True
+        return False
+    except Exception:
+        return False
+
+
+class InferenceBackend(str, Enum):
+    REST = "REST"
+    LOCAL_SAGEMAKER = "LOCAL_SAGEMAKER"
+    AWS_SAGEMAKER = "AWS_SAGEMAKER"
+    UNKNOWN = "UNKNOWN"
+
+
+def _detect_backend(
+        url: str) -> Tuple[InferenceBackend, Optional[str], Optional[str]]:
+    parsed = urlparse(url)
+
+    # Remote SageMaker
+    if ("runtime.sagemaker" in parsed.netloc
+            and parsed.path.endswith("/invocations")):
+        # Example: https://runtime.sagemaker.us-west-2.amazonaws.com/
+        # endpoints/Name/invocations
+        match = re.search(r"runtime\.sagemaker\.([a-z0-9-]+)\.amazonaws\.com",
+                          parsed.netloc)
+        region = match.group(1) if match else None
+        m = re.search(r"/endpoints/([^/]+)/invocations", parsed.path)
+        endpoint_name = m.group(1) if m else None
+        return InferenceBackend.AWS_SAGEMAKER, region, endpoint_name
+
+    # Local SageMaker
+    if parsed.port == 8080 and parsed.path.endswith(
+            "/invocations") and _is_local_address(parsed.hostname):
+        return InferenceBackend.LOCAL_SAGEMAKER, None, None
+
+    # Default: regular REST
+    return InferenceBackend.REST, None, None
+
+
+@dataclass
+class RfmGlobalState:
+    _url: str = '__url_not_provided__'
+    _backend: InferenceBackend = InferenceBackend.UNKNOWN
+    _region: Optional[str] = None
+    _endpoint_name: Optional[str] = None
+    _thread_local = threading.local()
+
+    # Thread-safe init-once.
+    _initialized: bool = False
+    _lock: threading.Lock = threading.Lock()
+
+    @property
+    def client(self) -> KumoClient:
+        if self._backend == InferenceBackend.REST:
+            return kumoai.global_state.client
+
+        if hasattr(self._thread_local, '_sagemaker'):
+            # Set the spcs token in the client to ensure it has the latest.
+            return self._thread_local._sagemaker
+
+        sagemaker_client: KumoClient
+        if self._backend == InferenceBackend.LOCAL_SAGEMAKER:
+            sagemaker_client = KumoClient_SageMakerProxy_Local(self._url)
+        else:
+            assert self._backend == InferenceBackend.AWS_SAGEMAKER
+            assert self._region
+            assert self._endpoint_name
+            sagemaker_client = KumoClient_SageMakerAdapter(
+                self._region, self._endpoint_name)
+
+        self._thread_local._sagemaker = sagemaker_client
+        return sagemaker_client
+
+    def reset(self) -> None:  # For testing only.
+        with self._lock:
+            self._initialized = False
+            self._url = '__url_not_provided__'
+            self._backend = InferenceBackend.UNKNOWN
+            self._region = None
+            self._endpoint_name = None
+            self._thread_local = threading.local()
+
+
+global_state = RfmGlobalState()
+
 
 def init(
     url: Optional[str] = None,
@@ -47,19 +155,54 @@ def init(
     snowflake_application: Optional[str] = None,
     log_level: str = "INFO",
 ) -> None:
-    if url is None:
-        url = os.getenv("KUMO_API_URL", "https://kumorfm.ai/api")
+    with global_state._lock:
+        if global_state._initialized:
+            if url != global_state._url:
+                raise ValueError(
+                    "Kumo RFM has already been initialized with a different "
+                    "URL. Re-initialization with a different URL is not "
+                    "supported.")
+            return
+
+        if url is None:
+            url = os.getenv("RFM_API_URL", "https://kumorfm.ai/api")
+
+        backend, region, endpoint_name = _detect_backend(url)
+        if backend == InferenceBackend.REST:
+            # Initialize kumoai.global_state
+            if (kumoai.global_state.initialized
+                    and kumoai.global_state._url != url):
+                raise ValueError(
+                    "Kumo AI SDK has already been initialized with different "
+                    "API URL. Please restart Python interpreter and "
+                    "initialize via kumoai.rfm.init()")
+            kumoai.init(url=url, api_key=api_key,
+                        snowflake_credentials=snowflake_credentials,
+                        snowflake_application=snowflake_application,
+                        log_level=log_level)
+        elif backend == InferenceBackend.AWS_SAGEMAKER:
+            assert region
+            assert endpoint_name
+            KumoClient_SageMakerAdapter(region, endpoint_name).authenticate()
+        else:
+            assert backend == InferenceBackend.LOCAL_SAGEMAKER
+            KumoClient_SageMakerProxy_Local(url).authenticate()
 
-    kumoai.init(url=url, api_key=api_key,
-                snowflake_credentials=snowflake_credentials,
-                snowflake_application=snowflake_application,
-                log_level=log_level)
+        global_state._url = url
+        global_state._backend = backend
+        global_state._region = region
+        global_state._endpoint_name = endpoint_name
+        global_state._initialized = True
+        logger.info("Kumo RFM initialized with backend: %s, url: %s", backend,
+                    url)
 
 
 __all__ = [
     'LocalTable',
     'LocalGraph',
     'KumoRFM',
+    'ExplainConfig',
+    'Explanation',
     'authenticate',
     'init',
 ]

kumoai/experimental/rfm/local_graph_sampler.py

@@ -2,7 +2,6 @@ from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 import pandas as pd
-from kumoapi.model_plan import RunMode
 from kumoapi.rfm.context import EdgeLayout, Link, Subgraph, Table
 from kumoapi.typing import Stype
 
@@ -33,7 +32,6 @@ class LocalGraphSampler:
         entity_table_names: Tuple[str, ...],
         node: np.ndarray,
         time: np.ndarray,
-        run_mode: RunMode,
         num_neighbors: List[int],
         exclude_cols_dict: Dict[str, List[str]],
     ) -> Subgraph:

kumoai/experimental/rfm/local_pquery_driver.py

@@ -264,8 +264,8 @@ class LocalPQueryDriver:
                     reached_end = True
                     break
                 candidate_offset = 0
-                end_offset = self._query.target_timeframe.end_date_offset
-                anchor_time = anchor_time - (end_offset *
+                time_frame = self._query.target_timeframe.timeframe
+                anchor_time = anchor_time - (time_frame *
                                              self._query.num_forecasts)
                 if anchor_time < self._graph_store.min_time:
                     reached_end = True

kumoai/experimental/rfm/rfm.py

@@ -5,7 +5,17 @@ from collections import defaultdict
 from collections.abc import Generator
 from contextlib import contextmanager
 from dataclasses import dataclass, replace
-from typing import Iterator, List, Literal, Optional, Tuple, Union, overload
+from typing import (
+    Any,
+    Dict,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    Union,
+    overload,
+)
 
 import numpy as np
 import pandas as pd
@@ -20,7 +30,7 @@ from kumoapi.rfm import (
 )
 from kumoapi.task import TaskType
 
-from kumoai import global_state
+from kumoai.client.rfm import RFMAPI
 from kumoai.exceptions import HTTPException
 from kumoai.experimental.rfm import LocalGraph
 from kumoai.experimental.rfm.local_graph_sampler import LocalGraphSampler
@@ -29,6 +39,7 @@ from kumoai.experimental.rfm.local_pquery_driver import (
     LocalPQueryDriver,
     date_offset_to_seconds,
 )
+from kumoai.mixin import CastMixin
 from kumoai.utils import InteractiveProgressLogger, ProgressLogger
 
 _RANDOM_SEED = 42
@@ -59,6 +70,17 @@ _SIZE_LIMIT_MSG = ("Context size exceeds the 30MB limit. {stats}\nPlease "
                    "beyond this for your use-case.")
 
 
+@dataclass(repr=False)
+class ExplainConfig(CastMixin):
+    """Configuration for explainability.
+
+    Args:
+        skip_summary: Whether to skip generating a human-readable summary of
+            the explanation.
+    """
+    skip_summary: bool = False
+
+
 @dataclass(repr=False)
 class Explanation:
     prediction: pd.DataFrame
@@ -86,6 +108,12 @@ class Explanation:
     def __repr__(self) -> str:
         return str((self.prediction, self.summary))
 
+    def _ipython_display_(self) -> None:
+        from IPython.display import Markdown, display
+
+        display(self.prediction)
+        display(Markdown(self.summary))
+
 
 class KumoRFM:
     r"""The Kumo Relational Foundation model (RFM) from the `KumoRFM: A
@@ -113,9 +141,9 @@ class KumoRFM:
 
         rfm = KumoRFM(graph)
 
-        query = ("PREDICT COUNT(transactions.*, 0, 30, days)>0 "
-                 "FOR users.user_id=0")
-        result = rfm.query(query)
+        query = ("PREDICT COUNT(orders.*, 0, 30, days)>0 "
+                 "FOR users.user_id=1")
+        result = rfm.predict(query)
 
         print(result)  # user_id  COUNT(transactions.*, 0, 30, days) > 0
                        # 1        0.85
@@ -144,9 +172,20 @@ class KumoRFM:
         self._graph_store = LocalGraphStore(graph, preprocess, verbose)
         self._graph_sampler = LocalGraphSampler(self._graph_store)
 
+        self._client: Optional[RFMAPI] = None
+
         self._batch_size: Optional[int | Literal['max']] = None
         self.num_retries: int = 0
 
+    @property
+    def _api_client(self) -> RFMAPI:
+        if self._client is not None:
+            return self._client
+
+        from kumoai.experimental.rfm import global_state
+        self._client = RFMAPI(global_state.client)
+        return self._client
+
     def __repr__(self) -> str:
         return f'{self.__class__.__name__}()'
 
@@ -208,7 +247,7 @@ class KumoRFM:
         query: str,
         indices: Union[List[str], List[float], List[int], None] = None,
         *,
-        explain: Literal[True],
+        explain: Union[Literal[True], ExplainConfig, Dict[str, Any]],
         anchor_time: Union[pd.Timestamp, Literal['entity'], None] = None,
         context_anchor_time: Union[pd.Timestamp, None] = None,
         run_mode: Union[RunMode, str] = RunMode.FAST,
@@ -226,7 +265,7 @@ class KumoRFM:
         query: str,
         indices: Union[List[str], List[float], List[int], None] = None,
         *,
-        explain: bool = False,
+        explain: Union[bool, ExplainConfig, Dict[str, Any]] = False,
         anchor_time: Union[pd.Timestamp, Literal['entity'], None] = None,
         context_anchor_time: Union[pd.Timestamp, None] = None,
         run_mode: Union[RunMode, str] = RunMode.FAST,
@@ -246,9 +285,12 @@ class KumoRFM:
                 be generated for all indices, independent of whether they
                 fulfill entity filter constraints. To pre-filter entities, use
                 :meth:`~KumoRFM.is_valid_entity`.
-            explain: If set to ``True``, will additionally explain the
-                prediction. Explainability is currently only supported for
-                single entity predictions with ``run_mode="FAST"``.
+            explain: Configuration for explainability.
+                If set to ``True``, will additionally explain the prediction.
+                Passing in an :class:`ExplainConfig` instance provides control
+                over which parts of explanation are generated.
+                Explainability is currently only supported for single entity
+                predictions with ``run_mode="FAST"``.
             anchor_time: The anchor timestamp for the prediction. If set to
                 ``None``, will use the maximum timestamp in the data.
                 If set to ``"entity"``, will use the timestamp of the entity.
@@ -272,16 +314,25 @@ class KumoRFM:
 
         Returns:
             The predictions as a :class:`pandas.DataFrame`.
-            If ``explain=True``, additionally returns a textual summary that
-            explains the prediction.
+            If ``explain`` is provided, returns an :class:`Explanation` object
+            containing the prediction, summary, and details.
         """
+        explain_config: Optional[ExplainConfig] = None
+        if explain is True:
+            explain_config = ExplainConfig()
+        elif explain is not False:
+            explain_config = ExplainConfig._cast(explain)
+
         query_def = self._parse_query(query)
+        query_str = query_def.to_string()
 
         if num_hops != 2 and num_neighbors is not None:
             warnings.warn(f"Received custom 'num_neighbors' option; ignoring "
                           f"custom 'num_hops={num_hops}' option")
 
-        if explain and run_mode in {RunMode.NORMAL, RunMode.BEST}:
+        if explain_config is not None and run_mode in {
+                RunMode.NORMAL, RunMode.BEST
+        }:
             warnings.warn(f"Explainability is currently only supported for "
                           f"run mode 'FAST' (got '{run_mode}'). Provided run "
                           f"mode has been reset. Please lower the run mode to "
@@ -298,13 +349,13 @@ class KumoRFM:
         if len(indices) == 0:
             raise ValueError("At least one entity is required")
 
-        if explain and len(indices) > 1:
+        if explain_config is not None and len(indices) > 1:
             raise ValueError(
                 f"Cannot explain predictions for more than a single entity "
                 f"(got {len(indices)})")
 
         query_repr = query_def.to_string(rich=True, exclude_predict=True)
-        if explain:
+        if explain_config is not None:
             msg = f'[bold]EXPLAIN[/bold] {query_repr}'
         else:
             msg = f'[bold]PREDICT[/bold] {query_repr}'
@@ -355,6 +406,7 @@ class KumoRFM:
                 request = RFMPredictRequest(
                     context=context,
                     run_mode=RunMode(run_mode),
+                    query=query_str,
                     use_prediction_time=use_prediction_time,
                 )
                 with warnings.catch_warnings():
@@ -378,12 +430,15 @@ class KumoRFM:
 
                 for attempt in range(self.num_retries + 1):
                     try:
-                        if explain:
-                            resp = global_state.client.rfm_api.explain(_bytes)
+                        if explain_config is not None:
+                            resp = self._api_client.explain(
+                                request=_bytes,
+                                skip_summary=explain_config.skip_summary,
+                            )
                             summary = resp.summary
                             details = resp.details
                         else:
-                            resp = global_state.client.rfm_api.predict(_bytes)
+                            resp = self._api_client.predict(_bytes)
                         df = pd.DataFrame(**resp.prediction)
 
                         # Cast 'ENTITY' to correct data type:
@@ -430,7 +485,7 @@ class KumoRFM:
         else:
             prediction = pd.concat(predictions, ignore_index=True)
 
-        if explain:
+        if explain_config is not None:
             assert len(predictions) == 1
             assert summary is not None
             assert details is not None
@@ -586,10 +641,10 @@ class KumoRFM:
 
             if len(request_bytes) > _MAX_SIZE:
                 stats_msg = Context.get_memory_stats(request_msg.context)
-                raise ValueError(_SIZE_LIMIT_MSG.format(stats_msg=stats_msg))
+                raise ValueError(_SIZE_LIMIT_MSG.format(stats=stats_msg))
 
             try:
-                resp = global_state.client.rfm_api.evaluate(request_bytes)
+                resp = self._api_client.evaluate(request_bytes)
             except HTTPException as e:
                 try:
                     msg = json.loads(e.detail)['detail']
@@ -687,7 +742,8 @@ class KumoRFM:
                 graph_definition=self._graph_def,
             )
 
-            resp = global_state.client.rfm_api.parse_query(request)
+            resp = self._api_client.parse_query(request)
+
             # TODO Expose validation warnings.
 
             if len(resp.validation_response.warnings) > 0:
@@ -991,7 +1047,6 @@ class KumoRFM:
                 train_time.astype('datetime64[ns]').astype(int).to_numpy(),
                 test_time.astype('datetime64[ns]').astype(int).to_numpy(),
             ]),
-            run_mode=run_mode,
             num_neighbors=num_neighbors,
             exclude_cols_dict=exclude_cols_dict,
         )

kumoai/experimental/rfm/sagemaker.py

@@ -0,0 +1,130 @@
+import base64
+import json
+from typing import Any, Dict, List, Tuple
+
+import boto3
+import requests
+from mypy_boto3_sagemaker_runtime.client import SageMakerRuntimeClient
+from mypy_boto3_sagemaker_runtime.type_defs import InvokeEndpointOutputTypeDef
+
+from kumoai.client import KumoClient
+from kumoai.client.endpoints import Endpoint, HTTPMethod
+from kumoai.exceptions import HTTPException
+
+
+class SageMakerResponseAdapter(requests.Response):
+    def __init__(self, sm_response: InvokeEndpointOutputTypeDef):
+        super().__init__()
+        # Read the body bytes
+        self._content = sm_response['Body'].read()
+        self.status_code = 200
+        self.headers['Content-Type'] = sm_response.get('ContentType',
+                                                       'application/json')
+        # Optionally, you can store original sm_response for debugging
+        self.sm_response = sm_response
+
+    @property
+    def text(self) -> str:
+        assert isinstance(self._content, bytes)
+        return self._content.decode('utf-8')
+
+    def json(self, **kwargs) -> dict[str, Any]:  # type: ignore
+        return json.loads(self.text, **kwargs)
+
+
+class KumoClient_SageMakerAdapter(KumoClient):
+    def __init__(self, region: str, endpoint_name: str):
+        self._client: SageMakerRuntimeClient = boto3.client(
+            service_name="sagemaker-runtime", region_name=region)
+        self._endpoint_name = endpoint_name
+
+        # Recording buffers.
+        self._recording_active = False
+        self._recorded_reqs: List[Dict[str, Any]] = []
+        self._recorded_resps: List[Dict[str, Any]] = []
+
+    def authenticate(self) -> None:
+        # TODO(siyang): call /ping to verify?
+        pass
+
+    def _request(self, endpoint: Endpoint, **kwargs: Any) -> requests.Response:
+        assert endpoint.method == HTTPMethod.POST
+        if 'json' in kwargs:
+            payload = json.dumps(kwargs.pop('json'))
+        elif 'data' in kwargs:
+            raw_payload = kwargs.pop('data')
+            assert isinstance(raw_payload, bytes)
+            payload = base64.b64encode(raw_payload).decode()
+        else:
+            raise HTTPException(400, 'Unable to send data to KumoRFM.')
+
+        request = {
+            'method': endpoint.get_path().rsplit('/')[-1],
+            'payload': payload,
+        }
+        response: InvokeEndpointOutputTypeDef = self._client.invoke_endpoint(
+            EndpointName=self._endpoint_name,
+            ContentType="application/json",
+            Body=json.dumps(request),
+        )
+
+        adapted_response = SageMakerResponseAdapter(response)
+
+        # If validation is active, store input/output
+        if self._recording_active:
+            self._recorded_reqs.append(request)
+            self._recorded_resps.append(adapted_response.json())
+
+        return adapted_response
+
+    def start_recording(self) -> None:
+        """Start recording requests/responses to/from sagemaker endpoint."""
+        assert not self._recording_active
+        self._recording_active = True
+        self._recorded_reqs.clear()
+        self._recorded_resps.clear()
+
+    def end_recording(self) -> List[Tuple[Dict[str, Any], Dict[str, Any]]]:
+        """Stop recording and return recorded requests/responses."""
+        assert self._recording_active
+        self._recording_active = False
+        recorded = list(zip(self._recorded_reqs, self._recorded_resps))
+        self._recorded_reqs.clear()
+        self._recorded_resps.clear()
+        return recorded
+
+
+class KumoClient_SageMakerProxy_Local(KumoClient):
+    def __init__(self, url: str):
+        self._client = KumoClient(url, api_key=None)
+        self._client._api_url = self._client._url
+        self._endpoint = Endpoint('/invocations', HTTPMethod.POST)
+
+    def authenticate(self) -> None:
+        try:
+            self._client._session.get(
+                self._url + '/ping',
+                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 _request(self, endpoint: Endpoint, **kwargs: Any) -> requests.Response:
+        assert endpoint.method == HTTPMethod.POST
+        if 'json' in kwargs:
+            payload = json.dumps(kwargs.pop('json'))
+        elif 'data' in kwargs:
+            raw_payload = kwargs.pop('data')
+            assert isinstance(raw_payload, bytes)
+            payload = base64.b64encode(raw_payload).decode()
+        else:
+            raise HTTPException(400, 'Unable to send data to KumoRFM.')
+        return self._client._request(
+            self._endpoint,
+            json={
+                'method': endpoint.get_path().rsplit('/')[-1],
+                'payload': payload,
+            },
+            **kwargs,
+        )

kumoai/spcs.py

@@ -54,9 +54,7 @@ def _refresh_spcs_token() -> None:
         api_key=global_state._api_key,
         spcs_token=spcs_token,
     )
-    if not client.authenticate():
-        raise ValueError("Client authentication failed. Please check if you "
-                         "have a valid API key.")
+    client.authenticate()
 
     # Update state:
     global_state.set_spcs_token(spcs_token)

kumoai/utils/progress_logger.py

@@ -103,10 +103,13 @@ class InteractiveProgressLogger(ProgressLogger):
             self._progress.update(self._task, advance=1)  # type: ignore
 
     def __enter__(self) -> Self:
+        from kumoai import in_notebook
+
         super().__enter__()
 
-        sys.stdout.write("\x1b]9;4;3\x07")
-        sys.stdout.flush()
+        if not in_notebook():  # Render progress bar in TUI.
+            sys.stdout.write("\x1b]9;4;3\x07")
+            sys.stdout.flush()
 
         if self.verbose:
             self._live = Live(
@@ -119,6 +122,8 @@ class InteractiveProgressLogger(ProgressLogger):
         return self
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        from kumoai import in_notebook
+
         super().__exit__(exc_type, exc_val, exc_tb)
 
         if exc_type is not None:
@@ -134,8 +139,9 @@ class InteractiveProgressLogger(ProgressLogger):
             self._live.stop()
             self._live = None
 
-        sys.stdout.write("\x1b]9;4;0\x07")
-        sys.stdout.flush()
+        if not in_notebook():
+            sys.stdout.write("\x1b]9;4;0\x07")
+            sys.stdout.flush()
 
     def __rich_console__(
         self,

{kumoai-2.12.0.dev202511071730.dist-info → kumoai-2.13.0.dev202511261731.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kumoai
-Version: 2.12.0.dev202511071730
+Version: 2.13.0.dev202511261731
 Summary: AI on the Modern Data Stack
 Author-email: "Kumo.AI" <hello@kumo.ai>
 License-Expression: MIT
@@ -23,11 +23,13 @@ Requires-Dist: requests>=2.28.2
 Requires-Dist: urllib3
 Requires-Dist: plotly
 Requires-Dist: typing_extensions>=4.5.0
-Requires-Dist: kumo-api==0.45.0
+Requires-Dist: kumo-api==0.46.0
 Requires-Dist: tqdm>=4.66.0
 Requires-Dist: aiohttp>=3.10.0
 Requires-Dist: pydantic>=1.10.21
 Requires-Dist: rich>=9.0.0
+Requires-Dist: mypy-boto3-sagemaker-runtime
+Requires-Dist: boto3
 Provides-Extra: doc
 Requires-Dist: sphinx; extra == "doc"
 Requires-Dist: sphinx-book-theme; extra == "doc"
@@ -38,6 +40,13 @@ Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-mock; extra == "test"
 Requires-Dist: requests-mock; extra == "test"
+Provides-Extra: test-sagemaker
+Requires-Dist: sagemaker; extra == "test-sagemaker"
+Requires-Dist: pandas==2.1.4; extra == "test-sagemaker"
+Requires-Dist: pyarrow==12.0.1; extra == "test-sagemaker"
+Provides-Extra: sagemaker
+Requires-Dist: boto3<2.0,>=1.30.0; extra == "sagemaker"
+Requires-Dist: mypy-boto3-sagemaker-runtime<2.0,>=1.34.0; extra == "sagemaker"
 Dynamic: license-file
 Dynamic: requires-dist
 

{kumoai-2.12.0.dev202511071730.dist-info → kumoai-2.13.0.dev202511261731.dist-info}/RECORD

@@ -1,27 +1,27 @@
-kumoai/__init__.py,sha256=4efagNAotP3c8mj8yyDGfVFcbgQ9l4wRC4FP-Yt0J3E,11002
+kumoai/__init__.py,sha256=qu-qohU2cQlManX1aZIlzA3ivKl52m-cSQBPSW8urUU,10837
 kumoai/_logging.py,sha256=qL4JbMQwKXri2f-SEJoFB8TY5ALG12S-nobGTNWxW-A,915
 kumoai/_singleton.py,sha256=i2BHWKpccNh5SJGDyU0IXsnYzJAYr8Xb0wz4c6LRbpo,861
-kumoai/_version.py,sha256=eTnPFx1EhbvlnxZni04U0p2axk-lywAnES4dQ8MuDTs,39
+kumoai/_version.py,sha256=P7PbPaqmt6kLq-80AyouMRr_ZBx8A7_nPBBPEXW44ag,39
 kumoai/databricks.py,sha256=ahwJz6DWLXMkndT0XwEDBxF-hoqhidFR8wBUQ4TLZ68,490
 kumoai/exceptions.py,sha256=7TMs0SC8xrU009_Pgd4QXtSF9lxJq8MtRbeX9pcQUy4,859
 kumoai/formatting.py,sha256=o3uCnLwXPhe1KI5WV9sBgRrcU7ed4rgu_pf89GL9Nc0,983
 kumoai/futures.py,sha256=J8rtZMEYFzdn5xF_x-LAiKJz3KGL6PT02f6rq_2bOJk,3836
 kumoai/jobs.py,sha256=dCi7BAdfm2tCnonYlGU4WJokJWbh3RzFfaOX2EYCIHU,2576
-kumoai/kumolib.cp310-win_amd64.pyd,sha256=GINcSpEYpFW8tP_vuwH5kWZSsWPYKnM83XHadL7hHGI,194048
+kumoai/kumolib.cp310-win_amd64.pyd,sha256=lfpQDN2Fu1tGqTxas5A9Jv3fLm-WdA_oSkhQfo6-pvg,194048
 kumoai/mixin.py,sha256=IaiB8SAI0VqOoMVzzIaUlqMt53-QPUK6OB0HikG-V9E,840
-kumoai/spcs.py,sha256=SWvfkeJvb_7sGkjSqyMBIuPbMTWCP6v0BC9HBXM1uSI,4398
+kumoai/spcs.py,sha256=KWfENrwSLruprlD-QPh63uU0N6npiNrwkeKfBk3EUyQ,4260
 kumoai/artifact_export/__init__.py,sha256=UXAQI5q92ChBzWAk8o3J6pElzYHudAzFZssQXd4o7i8,247
 kumoai/artifact_export/config.py,sha256=PRoUByzu5l-nyBKFR4vnRlq19b53ExGVy8YDCD7zMuI,8233
 kumoai/artifact_export/job.py,sha256=lOFIdPCrvhwdfvvDhQ2yzW8J4qIdYQoHZO1Rz3kJky4,3383
 kumoai/client/__init__.py,sha256=v0ISO1QD8JJhIJS6IzWz5-SL3EhtNCPeX3j1b2HBY0s,69
-kumoai/client/client.py,sha256=IoZ6WH-VIAdwpwmd5DhP4HqjQL_YpB5vaWjtaWrNECk,8801
+kumoai/client/client.py,sha256=T6Kw7-XWuAy5Dh7XU5graBl1-cTARiobycwtgxzaSE8,8731
 kumoai/client/connector.py,sha256=CO2LG5aDpCLxWNYYFRXGZs1AhYH3dRcbqBEUGwHQGzQ,4030
 kumoai/client/endpoints.py,sha256=DpEKEQ1yvL15iHZadXZKO94t-qXrYLaeV1sknX4IuPg,5532
 kumoai/client/graph.py,sha256=6MFyPYxDPfGTWeAI_84RUgWx9rVvqbLnR0Ourtgj5rg,3951
 kumoai/client/jobs.py,sha256=Y8wKiTk1I5ywc-2cxR72LaBjfhPTCVOezSCTeDpTs8Q,17521
 kumoai/client/online.py,sha256=4s_8Sv8m_k_tty4CO7RuAt0e6BDMkGvsZZ3VX8zyDb8,2798
 kumoai/client/pquery.py,sha256=0pXgQLxjoaFWDif0XRAuC_P-X3OSnXNWsiVrXej9uMk,7094
-kumoai/client/rfm.py,sha256=XCLJsSBe82fErLchpuS4Zb7fA3LBY8QxxIhrbw4_NPQ,3678
+kumoai/client/rfm.py,sha256=Gmt_dqoXekBCLiF0eQPgpoJ1cbnhnU8VbINF3U13qbQ,3838
 kumoai/client/source_table.py,sha256=mMHJtQ_yUHRI9LdHLVHxNGt83bbzmC1_d-NmXjbiTuI,2154
 kumoai/client/table.py,sha256=VhjLEMLQS1Z7zjcb2Yt3gZfiVqiD7b1gj-WNux_504A,3336
 kumoai/client/utils.py,sha256=RSD5Ia0lQQDR1drRFBJFdo2KVHfQqhJuk6m6du7Kl4E,3979
@@ -53,14 +53,15 @@ kumoai/connector/source_table.py,sha256=fnqwIKY6qYo4G0EsRzchb6FgZ-dQyU6aRaD9UAxs
 kumoai/connector/utils.py,sha256=SlkjPJS_wqfwFzIaQOHZtENQnbOz5sgLbvvvPDXE1ww,65786
 kumoai/encoder/__init__.py,sha256=8FeP6mUyCeXxr1b8kUIi5dxe5vEXQRft9tPoaV1CBqg,186
 kumoai/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-kumoai/experimental/rfm/__init__.py,sha256=21qX8kSSQ6SJ62zGZSAhESFHVrm7UgKIMqjhZiniLxk,1706
+kumoai/experimental/rfm/__init__.py,sha256=gpjpeN8PT3ZESi6kUaeyZqYnoJnysRVXDaY9hrycJA4,7020
 kumoai/experimental/rfm/authenticate.py,sha256=G89_4TMeUpr5fG_0VTzMF5sdNhaciitA1oc2loTlTmo,19321
 kumoai/experimental/rfm/local_graph.py,sha256=nZ9hDfyWg1dHFLoTEKoLt0ZJPvf9MUA1MNyfTRzJThg,30886
-kumoai/experimental/rfm/local_graph_sampler.py,sha256=ZCnILozG95EzpgMqhGTG2AF85JphLvAhj-3YPaTqoaQ,6922
+kumoai/experimental/rfm/local_graph_sampler.py,sha256=3JNpktW__nwxVKZxP4cQBgsIin7J_LNXYS7YlV36xbU,6854
 kumoai/experimental/rfm/local_graph_store.py,sha256=eUuIMFcdIRqN1kRxnqOdJpKEt-S_oyupAyHr7YuQoSU,14206
-kumoai/experimental/rfm/local_pquery_driver.py,sha256=XHxRTMRVUzKNlTItkOmW_ClEQ1xgvvwIC6MBLt7qihA,26857
+kumoai/experimental/rfm/local_pquery_driver.py,sha256=Yd_yHIrvuDj16IC1pvsqiQvZS41vvOOCRMiuDGtN6Fk,26851
 kumoai/experimental/rfm/local_table.py,sha256=5H08657TIyH7n_QnpFKr2g4BtVqdXTymmrfhSGaDmkU,20150
-kumoai/experimental/rfm/rfm.py,sha256=K9Fm6O3GWkoOCv9Bq8jSdnWvuMyPYk4lmU1WJIpLSPY,47815
+kumoai/experimental/rfm/rfm.py,sha256=MarISSPKuv6nIaGG69zFAwIagF6EA37xcSRClZrQMFc,49470
+kumoai/experimental/rfm/sagemaker.py,sha256=eebpZtASqiIGF2FpY53bbWLj6p-u5hkK4RLgBNAvEzg,4953
 kumoai/experimental/rfm/utils.py,sha256=dLx2wdyTWg7vZI_7R-I0z_lA-2aV5M8h9n3bnnLyylI,11467
 kumoai/experimental/rfm/infer/__init__.py,sha256=fPsdDr4D3hgC8snW0j3pAVpCyR-xrauuogMnTOMrfok,304
 kumoai/experimental/rfm/infer/categorical.py,sha256=bqmfrE5ZCBTcb35lA4SyAkCu3MgttAn29VBJYMBNhVg,893
@@ -90,9 +91,9 @@ kumoai/trainer/util.py,sha256=LCXkY5MNl6NbEVd2OZ0aVqF6fvr3KiCFh6pH0igAi_g,4165
 kumoai/utils/__init__.py,sha256=wAKgmwtMIGuiauW9D_GGKH95K-24Kgwmld27mm4nsro,278
 kumoai/utils/datasets.py,sha256=UyAII-oAn7x3ombuvpbSQ41aVF9SYKBjQthTD-vcT2A,3011
 kumoai/utils/forecasting.py,sha256=ZgKeUCbWLOot0giAkoigwU5du8LkrwAicFOi5hVn6wg,7624
-kumoai/utils/progress_logger.py,sha256=tzwFrUO5VuiArxx9_tSETno8JF5rnFOedX26I2yDW10,5046
-kumoai-2.12.0.dev202511071730.dist-info/licenses/LICENSE,sha256=ZUilBDp--4vbhsEr6f_Upw9rnIx09zQ3K9fXQ0rfd6w,1111
-kumoai-2.12.0.dev202511071730.dist-info/METADATA,sha256=iXxnyQJGKAUlzPdN8-dA9N4aIFXi6Y2vdzpsGTzl8w8,2112
-kumoai-2.12.0.dev202511071730.dist-info/WHEEL,sha256=KUuBC6lxAbHCKilKua8R9W_TM71_-9Sg5uEP3uDWcoU,101
-kumoai-2.12.0.dev202511071730.dist-info/top_level.txt,sha256=YjU6UcmomoDx30vEXLsOU784ED7VztQOsFApk1SFwvs,7
-kumoai-2.12.0.dev202511071730.dist-info/RECORD,,
+kumoai/utils/progress_logger.py,sha256=MZsWgHd4UZQKCXiJZgQeW-Emi_BmzlCKPLPXOL_HqBo,5239
+kumoai-2.13.0.dev202511261731.dist-info/licenses/LICENSE,sha256=ZUilBDp--4vbhsEr6f_Upw9rnIx09zQ3K9fXQ0rfd6w,1111
+kumoai-2.13.0.dev202511261731.dist-info/METADATA,sha256=KEPt_QdWVLyZhYMj3PNjwvb1gm6fXI9_FvWyMOvMrtw,2544
+kumoai-2.13.0.dev202511261731.dist-info/WHEEL,sha256=KUuBC6lxAbHCKilKua8R9W_TM71_-9Sg5uEP3uDWcoU,101
+kumoai-2.13.0.dev202511261731.dist-info/top_level.txt,sha256=YjU6UcmomoDx30vEXLsOU784ED7VztQOsFApk1SFwvs,7
+kumoai-2.13.0.dev202511261731.dist-info/RECORD,,