kumoai 2.12.0.dev202511021731__cp310-cp310-macosx_11_0_arm64.whl → 2.13.0.dev202511191731__cp310-cp310-macosx_11_0_arm64.whl

kumoai/experimental/rfm/rfm.py

@@ -5,23 +5,32 @@ 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, Union, overload
+from typing import (
+    Any,
+    Dict,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    Union,
+    overload,
+)
 
 import numpy as np
 import pandas as pd
 from kumoapi.model_plan import RunMode
-from kumoapi.pquery import QueryType
+from kumoapi.pquery import QueryType, ValidatedPredictiveQuery
 from kumoapi.rfm import Context
 from kumoapi.rfm import Explanation as ExplanationConfig
 from kumoapi.rfm import (
-    PQueryDefinition,
     RFMEvaluateRequest,
+    RFMParseQueryRequest,
     RFMPredictRequest,
-    RFMValidateQueryRequest,
 )
 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
@@ -30,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
@@ -60,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
@@ -87,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
@@ -114,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
@@ -147,6 +174,8 @@ class KumoRFM:
 
         self._batch_size: Optional[int | Literal['max']] = None
         self.num_retries: int = 0
+        from kumoai.experimental.rfm import global_state
+        self._api_client = RFMAPI(global_state.client)
 
     def __repr__(self) -> str:
         return f'{self.__class__.__name__}()'
@@ -209,7 +238,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,
@@ -227,7 +256,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,
@@ -247,9 +276,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.
@@ -273,42 +305,48 @@ 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 "
                           f"suppress this warning.")
 
         if indices is None:
-            if query_def.entity.ids is None:
+            if query_def.rfm_entity_ids is None:
                 raise ValueError("Cannot find entities to predict for. Please "
                                  "pass them via `predict(query, indices=...)`")
-            indices = query_def.entity.ids.value
+            indices = query_def.get_rfm_entity_id_list()
         else:
-            query_def = replace(
-                query_def,
-                entity=replace(query_def.entity, ids=None),
-            )
+            query_def = replace(query_def, rfm_entity_ids=None)
 
         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}'
@@ -320,8 +358,8 @@ class KumoRFM:
 
             batch_size: Optional[int] = None
             if self._batch_size == 'max':
-                task_type = query_def.get_task_type(
-                    stypes=self._graph_store.stype_dict,
+                task_type = LocalPQueryDriver.get_task_type(
+                    query_def,
                     edge_types=self._graph_store.edge_types,
                 )
                 batch_size = _MAX_PRED_SIZE[task_type]
@@ -359,6 +397,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():
@@ -382,17 +421,20 @@ 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:
                         if 'ENTITY' in df:
-                            entity = query_def.entity.pkey.table_name
+                            entity = query_def.entity_table
                             pkey_map = self._graph_store.pkey_map_dict[entity]
                             df['ENTITY'] = df['ENTITY'].astype(
                                 type(pkey_map.index[0]))
@@ -434,7 +476,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
@@ -468,11 +510,11 @@ class KumoRFM:
         query_def = self._parse_query(query)
 
         if indices is None:
-            if query_def.entity.ids is None:
+            if query_def.rfm_entity_ids is None:
                 raise ValueError("Cannot find entities to predict for. Please "
                                  "pass them via "
                                  "`is_valid_entity(query, indices=...)`")
-            indices = query_def.entity.ids.value
+            indices = query_def.get_rfm_entity_id_list()
 
         if len(indices) == 0:
             raise ValueError("At least one entity is required")
@@ -484,14 +526,13 @@ class KumoRFM:
             self._validate_time(query_def, anchor_time, None, False)
         else:
             assert anchor_time == 'entity'
-            if (query_def.entity.pkey.table_name
-                    not in self._graph_store.time_dict):
+            if (query_def.entity_table not in self._graph_store.time_dict):
                 raise ValueError(f"Anchor time 'entity' requires the entity "
-                                 f"table '{query_def.entity.pkey.table_name}' "
-                                 f"to have a time column")
+                                 f"table '{query_def.entity_table}' "
+                                 f"to have a time column.")
 
         node = self._graph_store.get_node_id(
-            table_name=query_def.entity.pkey.table_name,
+            table_name=query_def.entity_table,
             pkey=pd.Series(indices),
         )
         query_driver = LocalPQueryDriver(self._graph_store, query_def)
@@ -547,10 +588,10 @@ class KumoRFM:
             warnings.warn(f"Received custom 'num_neighbors' option; ignoring "
                           f"custom 'num_hops={num_hops}' option")
 
-        if query_def.entity.ids is not None:
+        if query_def.rfm_entity_ids is not None:
             query_def = replace(
                 query_def,
-                entity=replace(query_def.entity, ids=None),
+                rfm_entity_ids=None,
             )
 
         query_repr = query_def.to_string(rich=True, exclude_predict=True)
@@ -591,10 +632,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']
@@ -639,18 +680,19 @@ class KumoRFM:
 
         if anchor_time is None:
             anchor_time = self._graph_store.max_time
-            anchor_time = anchor_time - (query_def.target.end_offset *
-                                         query_def.num_forecasts)
+            if query_def.target_ast.date_offset_range is not None:
+                anchor_time = anchor_time - (
+                    query_def.target_ast.date_offset_range.end_date_offset *
+                    query_def.num_forecasts)
 
         assert anchor_time is not None
         if isinstance(anchor_time, pd.Timestamp):
             self._validate_time(query_def, anchor_time, None, evaluate=True)
         else:
             assert anchor_time == 'entity'
-            if (query_def.entity.pkey.table_name
-                    not in self._graph_store.time_dict):
+            if (query_def.entity_table not in self._graph_store.time_dict):
                 raise ValueError(f"Anchor time 'entity' requires the entity "
-                                 f"table '{query_def.entity.pkey.table_name}' "
+                                 f"table '{query_def.entity_table}' "
                                  f"to have a time column")
 
         query_driver = LocalPQueryDriver(self._graph_store, query_def,
@@ -665,7 +707,7 @@ class KumoRFM:
         )
 
         entity = self._graph_store.pkey_map_dict[
-            query_def.entity.pkey.table_name].index[node]
+            query_def.entity_table].index[node]
 
         return pd.DataFrame({
             'ENTITY': entity,
@@ -675,8 +717,8 @@ class KumoRFM:
 
     # Helpers #################################################################
 
-    def _parse_query(self, query: str) -> PQueryDefinition:
-        if isinstance(query, PQueryDefinition):
+    def _parse_query(self, query: str) -> ValidatedPredictiveQuery:
+        if isinstance(query, ValidatedPredictiveQuery):
             return query
 
         if isinstance(query, str) and query.strip()[:9].lower() == 'evaluate ':
@@ -686,12 +728,13 @@ class KumoRFM:
                              "predictions or evaluations.")
 
         try:
-            request = RFMValidateQueryRequest(
+            request = RFMParseQueryRequest(
                 query=query,
                 graph_definition=self._graph_def,
             )
 
-            resp = global_state.client.rfm_api.validate_query(request)
+            resp = self._api_client.parse_query(request)
+
             # TODO Expose validation warnings.
 
             if len(resp.validation_response.warnings) > 0:
@@ -702,7 +745,7 @@ class KumoRFM:
                 warnings.warn(f"Encountered the following warnings during "
                               f"parsing:\n{msg}")
 
-            return resp.query_definition
+            return resp.query
         except HTTPException as e:
             try:
                 msg = json.loads(e.detail)['detail']
@@ -713,7 +756,7 @@ class KumoRFM:
 
     def _validate_time(
         self,
-        query: PQueryDefinition,
+        query: ValidatedPredictiveQuery,
         anchor_time: pd.Timestamp,
         context_anchor_time: Union[pd.Timestamp, None],
         evaluate: bool,
@@ -736,6 +779,11 @@ class KumoRFM:
                              f"only contains data back to "
                              f"'{self._graph_store.min_time}'.")
 
+        if query.target_ast.date_offset_range is not None:
+            end_offset = query.target_ast.date_offset_range.end_date_offset
+        else:
+            end_offset = pd.DateOffset(0)
+        forecast_end_offset = end_offset * query.num_forecasts
         if (context_anchor_time is not None
                 and context_anchor_time > anchor_time):
             warnings.warn(f"Context anchor timestamp "
@@ -744,19 +792,18 @@ class KumoRFM:
                           f"(got '{anchor_time}'). Please make sure this is "
                           f"intended.")
         elif (query.query_type == QueryType.TEMPORAL
-              and context_anchor_time is not None and context_anchor_time +
-              query.target.end_offset * query.num_forecasts > anchor_time):
+              and context_anchor_time is not None
+              and context_anchor_time + forecast_end_offset > anchor_time):
             warnings.warn(f"Aggregation for context examples at timestamp "
                           f"'{context_anchor_time}' will leak information "
                           f"from the prediction anchor timestamp "
                           f"'{anchor_time}'. Please make sure this is "
                           f"intended.")
 
-        elif (context_anchor_time is not None and context_anchor_time -
-              query.target.end_offset * query.num_forecasts
+        elif (context_anchor_time is not None
+              and context_anchor_time - forecast_end_offset
               < self._graph_store.min_time):
-            _time = context_anchor_time - (query.target.end_offset *
-                                           query.num_forecasts)
+            _time = context_anchor_time - forecast_end_offset
             warnings.warn(f"Context anchor timestamp is too early or "
                           f"aggregation time range is too large. To form "
                           f"proper input data, we would need data back to "
@@ -769,8 +816,7 @@ class KumoRFM:
                           f"latest timestamp '{self._graph_store.max_time}' "
                           f"in the data. Please make sure this is intended.")
 
-        max_eval_time = (self._graph_store.max_time -
-                         query.target.end_offset * query.num_forecasts)
+        max_eval_time = self._graph_store.max_time - forecast_end_offset
         if evaluate and anchor_time > max_eval_time:
             raise ValueError(
                 f"Anchor timestamp for evaluation is after the latest "
@@ -778,7 +824,7 @@ class KumoRFM:
 
     def _get_context(
         self,
-        query: PQueryDefinition,
+        query: ValidatedPredictiveQuery,
         indices: Union[List[str], List[float], List[int], None],
         anchor_time: Union[pd.Timestamp, Literal['entity'], None],
         context_anchor_time: Union[pd.Timestamp, None],
@@ -806,8 +852,8 @@ class KumoRFM:
                              f"must go beyond this for your use-case.")
 
         query_driver = LocalPQueryDriver(self._graph_store, query, random_seed)
-        task_type = query.get_task_type(
-            stypes=self._graph_store.stype_dict,
+        task_type = LocalPQueryDriver.get_task_type(
+            query,
             edge_types=self._graph_store.edge_types,
         )
 
@@ -839,11 +885,15 @@ class KumoRFM:
             else:
                 num_neighbors = [64, 64, 8, 8, 4, 4][:num_hops]
 
+        if query.target_ast.date_offset_range is None:
+            end_offset = pd.DateOffset(0)
+        else:
+            end_offset = query.target_ast.date_offset_range.end_date_offset
+        forecast_end_offset = end_offset * query.num_forecasts
         if anchor_time is None:
             anchor_time = self._graph_store.max_time
             if evaluate:
-                anchor_time = anchor_time - (query.target.end_offset *
-                                             query.num_forecasts)
+                anchor_time = anchor_time - forecast_end_offset
             if logger is not None:
                 assert isinstance(anchor_time, pd.Timestamp)
                 if anchor_time == pd.Timestamp.min:
@@ -858,15 +908,14 @@ class KumoRFM:
         assert anchor_time is not None
         if isinstance(anchor_time, pd.Timestamp):
             if context_anchor_time is None:
-                context_anchor_time = anchor_time - (query.target.end_offset *
-                                                     query.num_forecasts)
+                context_anchor_time = anchor_time - forecast_end_offset
             self._validate_time(query, anchor_time, context_anchor_time,
                                 evaluate)
         else:
             assert anchor_time == 'entity'
-            if query.entity.pkey.table_name not in self._graph_store.time_dict:
+            if query.entity_table not in self._graph_store.time_dict:
                 raise ValueError(f"Anchor time 'entity' requires the entity "
-                                 f"table '{query.entity.pkey.table_name}' to "
+                                 f"table '{query.entity_table}' to "
                                  f"have a time column")
             if context_anchor_time is not None:
                 warnings.warn("Ignoring option 'context_anchor_time' for "
@@ -917,7 +966,7 @@ class KumoRFM:
                                  f"in batches")
 
             test_node = self._graph_store.get_node_id(
-                table_name=query.entity.pkey.table_name,
+                table_name=query.entity_table,
                 pkey=pd.Series(indices),
             )
 
@@ -925,8 +974,7 @@ class KumoRFM:
                 test_time = pd.Series(anchor_time).repeat(
                     len(test_node)).reset_index(drop=True)
             else:
-                time = self._graph_store.time_dict[
-                    query.entity.pkey.table_name]
+                time = self._graph_store.time_dict[query.entity_table]
                 time = time[test_node] * 1000**3
                 test_time = pd.Series(time, dtype='datetime64[ns]')
 
@@ -959,12 +1007,23 @@ class KumoRFM:
                 raise NotImplementedError
             logger.log(msg)
 
-        entity_table_names = query.get_entity_table_names(
-            self._graph_store.edge_types)
+        entity_table_names: Tuple[str, ...]
+        if task_type.is_link_pred:
+            final_aggr = query.get_final_target_aggregation()
+            assert final_aggr is not None
+            edge_fkey = final_aggr._get_target_column_name()
+            for edge_type in self._graph_store.edge_types:
+                if edge_fkey == f'{edge_type[0]}.{edge_type[1]}':
+                    entity_table_names = (
+                        query.entity_table,
+                        edge_type[2],
+                    )
+        else:
+            entity_table_names = (query.entity_table, )
 
         # Exclude the entity anchor time from the feature set to prevent
         # running out-of-distribution between in-context and test examples:
-        exclude_cols_dict = query.exclude_cols_dict
+        exclude_cols_dict = query.get_exclude_cols_dict()
         if anchor_time == 'entity':
             if entity_table_names[0] not in exclude_cols_dict:
                 exclude_cols_dict[entity_table_names[0]] = []
@@ -979,7 +1038,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,
         )
@@ -993,7 +1051,7 @@ class KumoRFM:
 
         step_size: Optional[int] = None
         if query.query_type == QueryType.TEMPORAL:
-            step_size = date_offset_to_seconds(query.target.end_offset)
+            step_size = date_offset_to_seconds(end_offset)
 
         return Context(
             task_type=task_type,

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)