kumoai 2.12.0.dev202511071730__cp310-cp310-win_amd64.whl → 2.13.0.dev202512021731__cp310-cp310-win_amd64.whl

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,15 +30,16 @@ 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 import Graph
 from kumoai.experimental.rfm.local_graph_sampler import LocalGraphSampler
 from kumoai.experimental.rfm.local_graph_store import LocalGraphStore
 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
@@ -95,17 +123,17 @@ class KumoRFM:
     :class:`KumoRFM` is a foundation model to generate predictions for any
     relational dataset without training.
     The model is pre-trained and the class provides an interface to query the
-    model from a :class:`LocalGraph` object.
+    model from a :class:`Graph` object.
 
     .. code-block:: python
 
-        from kumoai.experimental.rfm import LocalGraph, KumoRFM
+        from kumoai.experimental.rfm import Graph, KumoRFM
 
         df_users = pd.DataFrame(...)
         df_items = pd.DataFrame(...)
         df_orders = pd.DataFrame(...)
 
-        graph = LocalGraph.from_data({
+        graph = Graph.from_data({
             'users': df_users,
             'items': df_items,
             'orders': df_orders,
@@ -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
@@ -135,7 +163,7 @@ class KumoRFM:
     """
     def __init__(
         self,
-        graph: LocalGraph,
+        graph: Graph,
         preprocess: bool = False,
         verbose: Union[bool, ProgressLogger] = True,
     ) -> None:
@@ -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,138 @@
+import base64
+import json
+from typing import Any, Dict, List, Tuple
+
+import requests
+
+from kumoai.client import KumoClient
+from kumoai.client.endpoints import Endpoint, HTTPMethod
+from kumoai.exceptions import HTTPException
+
+try:
+    # isort: off
+    from mypy_boto3_sagemaker_runtime.client import SageMakerRuntimeClient
+    from mypy_boto3_sagemaker_runtime.type_defs import (
+        InvokeEndpointOutputTypeDef, )
+    # isort: on
+except ImportError:
+    SageMakerRuntimeClient = Any
+    InvokeEndpointOutputTypeDef = Any
+
+
+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):
+        import boto3
+        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/experimental/rfm/utils.py

@@ -1,127 +1,8 @@
 import re
 import warnings
-from typing import Any, Dict, Optional
+from typing import Optional
 
-import numpy as np
 import pandas as pd
-import pyarrow as pa
-from kumoapi.typing import Dtype, Stype
-
-from kumoai.experimental.rfm.infer import (
-    contains_categorical,
-    contains_id,
-    contains_multicategorical,
-    contains_timestamp,
-)
-
-# Mapping from pandas/numpy dtypes to Kumo Dtypes
-PANDAS_TO_DTYPE: Dict[Any, Dtype] = {
-    np.dtype('bool'): Dtype.bool,
-    pd.BooleanDtype(): Dtype.bool,
-    pa.bool_(): Dtype.bool,
-    np.dtype('byte'): Dtype.int,
-    pd.UInt8Dtype(): Dtype.int,
-    np.dtype('int16'): Dtype.int,
-    pd.Int16Dtype(): Dtype.int,
-    np.dtype('int32'): Dtype.int,
-    pd.Int32Dtype(): Dtype.int,
-    np.dtype('int64'): Dtype.int,
-    pd.Int64Dtype(): Dtype.int,
-    np.dtype('float32'): Dtype.float,
-    pd.Float32Dtype(): Dtype.float,
-    np.dtype('float64'): Dtype.float,
-    pd.Float64Dtype(): Dtype.float,
-    np.dtype('object'): Dtype.string,
-    pd.StringDtype(storage='python'): Dtype.string,
-    pd.StringDtype(storage='pyarrow'): Dtype.string,
-    pa.string(): Dtype.string,
-    pa.binary(): Dtype.binary,
-    np.dtype('datetime64[ns]'): Dtype.date,
-    np.dtype('timedelta64[ns]'): Dtype.timedelta,
-    pa.list_(pa.float32()): Dtype.floatlist,
-    pa.list_(pa.int64()): Dtype.intlist,
-    pa.list_(pa.string()): Dtype.stringlist,
-}
-
-
-def to_dtype(ser: pd.Series) -> Dtype:
-    """Extracts the :class:`Dtype` from a :class:`pandas.Series`.
-
-    Args:
-        ser: A :class:`pandas.Series` to analyze.
-
-    Returns:
-        The data type.
-    """
-    if pd.api.types.is_datetime64_any_dtype(ser.dtype):
-        return Dtype.date
-
-    if isinstance(ser.dtype, pd.CategoricalDtype):
-        return Dtype.string
-
-    if pd.api.types.is_object_dtype(ser.dtype):
-        index = ser.iloc[:1000].first_valid_index()
-        if index is not None and pd.api.types.is_list_like(ser[index]):
-            pos = ser.index.get_loc(index)
-            assert isinstance(pos, int)
-            ser = ser.iloc[pos:pos + 1000].dropna()
-
-            if not ser.map(pd.api.types.is_list_like).all():
-                raise ValueError("Data contains a mix of list-like and "
-                                 "non-list-like values")
-
-            ser = ser[ser.map(lambda x: not isinstance(x, list) or len(x) > 0)]
-
-            dtypes = ser.apply(lambda x: PANDAS_TO_DTYPE.get(
-                np.array(x).dtype, Dtype.string)).unique().tolist()
-
-            invalid_dtypes = set(dtypes) - {
-                Dtype.string,
-                Dtype.int,
-                Dtype.float,
-            }
-            if len(invalid_dtypes) > 0:
-                raise ValueError(f"Data contains unsupported list data types: "
-                                 f"{list(invalid_dtypes)}")
-
-            if Dtype.string in dtypes:
-                return Dtype.stringlist
-
-            if dtypes == [Dtype.int]:
-                return Dtype.intlist
-
-            return Dtype.floatlist
-
-    if ser.dtype not in PANDAS_TO_DTYPE:
-        raise ValueError(f"Unsupported data type '{ser.dtype}'")
-
-    return PANDAS_TO_DTYPE[ser.dtype]
-
-
-def infer_stype(ser: pd.Series, column_name: str, dtype: Dtype) -> Stype:
-    r"""Infers the semantic type of a column.
-
-    Args:
-        ser: A :class:`pandas.Series` to analyze.
-        column_name: The name of the column (used for pattern matching).
-        dtype: The data type.
-
-    Returns:
-        The semantic type.
-    """
-    if contains_id(ser, column_name, dtype):
-        return Stype.ID
-
-    if contains_timestamp(ser, column_name, dtype):
-        return Stype.timestamp
-
-    if contains_multicategorical(ser, column_name, dtype):
-        return Stype.multicategorical
-
-    if contains_categorical(ser, column_name, dtype):
-        return Stype.categorical
-
-    return dtype.default_stype
 
 
 def detect_primary_key(

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/testing/decorators.py

@@ -25,7 +25,7 @@ def onlyFullTest(func: Callable) -> Callable:
 def has_package(package: str) -> bool:
     r"""Returns ``True`` in case ``package`` is installed."""
     req = Requirement(package)
-    if importlib.util.find_spec(req.name) is None:
+    if importlib.util.find_spec(req.name) is None:  # type: ignore
         return False
 
     try:

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.dev202512021731.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kumoai
-Version: 2.12.0.dev202511071730
+Version: 2.13.0.dev202512021731
 Summary: AI on the Modern Data Stack
 Author-email: "Kumo.AI" <hello@kumo.ai>
 License-Expression: MIT
@@ -23,7 +23,7 @@ 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.48.0
 Requires-Dist: tqdm>=4.66.0
 Requires-Dist: aiohttp>=3.10.0
 Requires-Dist: pydantic>=1.10.21
@@ -38,6 +38,15 @@ Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-mock; extra == "test"
 Requires-Dist: requests-mock; extra == "test"
+Provides-Extra: sqlite
+Requires-Dist: adbc_driver_sqlite; extra == "sqlite"
+Provides-Extra: snowflake
+Requires-Dist: snowflake-connector-python; extra == "snowflake"
+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"
+Provides-Extra: test-sagemaker
+Requires-Dist: sagemaker<3.0; extra == "test-sagemaker"
 Dynamic: license-file
 Dynamic: requires-dist