kumoai 2.14.0.dev202512151351__cp313-cp313-macosx_11_0_arm64.whl → 2.14.0.dev202512211732__cp313-cp313-macosx_11_0_arm64.whl

kumoai/experimental/rfm/infer/time_col.py

@@ -1,6 +1,5 @@
 import re
 import warnings
-from typing import Optional
 
 import pandas as pd
 
@@ -8,7 +7,7 @@ import pandas as pd
 def infer_time_column(
     df: pd.DataFrame,
     candidates: list[str],
-) -> Optional[str]:
+) -> str | None:
     r"""Auto-detect potential time column.
 
     Args:

kumoai/experimental/rfm/pquery/executor.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Dict, Generic, Tuple, TypeVar
+from typing import Generic, TypeVar
 
 from kumoapi.pquery import ValidatedPredictiveQuery
 from kumoapi.pquery.AST import (
@@ -21,82 +21,82 @@ class PQueryExecutor(Generic[TableData, ColumnData, IndexData], ABC):
     def execute_column(
         self,
         column: Column,
-        feat_dict: Dict[str, TableData],
+        feat_dict: dict[str, TableData],
         filter_na: bool = True,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
     def execute_aggregation(
         self,
         aggr: Aggregation,
-        feat_dict: Dict[str, TableData],
-        time_dict: Dict[str, ColumnData],
-        batch_dict: Dict[str, IndexData],
+        feat_dict: dict[str, TableData],
+        time_dict: dict[str, ColumnData],
+        batch_dict: dict[str, IndexData],
         anchor_time: ColumnData,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
     def execute_condition(
         self,
         condition: Condition,
-        feat_dict: Dict[str, TableData],
-        time_dict: Dict[str, ColumnData],
-        batch_dict: Dict[str, IndexData],
+        feat_dict: dict[str, TableData],
+        time_dict: dict[str, ColumnData],
+        batch_dict: dict[str, IndexData],
         anchor_time: ColumnData,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
     def execute_logical_operation(
         self,
         logical_operation: LogicalOperation,
-        feat_dict: Dict[str, TableData],
-        time_dict: Dict[str, ColumnData],
-        batch_dict: Dict[str, IndexData],
+        feat_dict: dict[str, TableData],
+        time_dict: dict[str, ColumnData],
+        batch_dict: dict[str, IndexData],
         anchor_time: ColumnData,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
     def execute_join(
         self,
         join: Join,
-        feat_dict: Dict[str, TableData],
-        time_dict: Dict[str, ColumnData],
-        batch_dict: Dict[str, IndexData],
+        feat_dict: dict[str, TableData],
+        time_dict: dict[str, ColumnData],
+        batch_dict: dict[str, IndexData],
         anchor_time: ColumnData,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
     def execute_filter(
         self,
         filter: Filter,
-        feat_dict: Dict[str, TableData],
-        time_dict: Dict[str, ColumnData],
-        batch_dict: Dict[str, IndexData],
+        feat_dict: dict[str, TableData],
+        time_dict: dict[str, ColumnData],
+        batch_dict: dict[str, IndexData],
         anchor_time: ColumnData,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
     def execute(
         self,
         query: ValidatedPredictiveQuery,
-        feat_dict: Dict[str, TableData],
-        time_dict: Dict[str, ColumnData],
-        batch_dict: Dict[str, IndexData],
+        feat_dict: dict[str, TableData],
+        time_dict: dict[str, ColumnData],
+        batch_dict: dict[str, IndexData],
         anchor_time: ColumnData,
         num_forecasts: int = 1,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass

kumoai/experimental/rfm/pquery/pandas_executor.py

@@ -1,5 +1,3 @@
-from typing import Dict, List, Tuple
-
 import numpy as np
 import pandas as pd
 from kumoapi.pquery import ValidatedPredictiveQuery
@@ -22,9 +20,9 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute_column(
         self,
         column: Column,
-        feat_dict: Dict[str, pd.DataFrame],
+        feat_dict: dict[str, pd.DataFrame],
         filter_na: bool = True,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         table_name, column_name = column.fqn.split(".")
         if column_name == '*':
             out = pd.Series(np.ones(len(feat_dict[table_name]), dtype='int64'))
@@ -60,7 +58,7 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
         batch: np.ndarray,
         batch_size: int,
         filter_na: bool = True,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
 
         mask = feat.notna()
         feat, batch = feat[mask], batch[mask]
@@ -104,13 +102,13 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute_aggregation(
         self,
         aggr: Aggregation,
-        feat_dict: Dict[str, pd.DataFrame],
-        time_dict: Dict[str, pd.Series],
-        batch_dict: Dict[str, np.ndarray],
+        feat_dict: dict[str, pd.DataFrame],
+        time_dict: dict[str, pd.Series],
+        batch_dict: dict[str, np.ndarray],
         anchor_time: pd.Series,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         target_table = aggr._get_target_column_name().split('.')[0]
         target_batch = batch_dict[target_table]
         target_time = time_dict[target_table]
@@ -131,8 +129,8 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
                 filter_na=True,
             )
 
-        outs: List[pd.Series] = []
-        masks: List[np.ndarray] = []
+        outs: list[pd.Series] = []
+        masks: list[np.ndarray] = []
         for _ in range(num_forecasts):
             anchor_target_time = anchor_time.iloc[target_batch]
             anchor_target_time = anchor_target_time.reset_index(drop=True)
@@ -226,13 +224,13 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute_condition(
         self,
         condition: Condition,
-        feat_dict: Dict[str, pd.DataFrame],
-        time_dict: Dict[str, pd.Series],
-        batch_dict: Dict[str, np.ndarray],
+        feat_dict: dict[str, pd.DataFrame],
+        time_dict: dict[str, pd.Series],
+        batch_dict: dict[str, np.ndarray],
         anchor_time: pd.Series,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         if num_forecasts > 1:
             raise NotImplementedError("Forecasting not yet implemented for "
                                       "non-regression tasks")
@@ -306,13 +304,13 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute_logical_operation(
         self,
         logical_operation: LogicalOperation,
-        feat_dict: Dict[str, pd.DataFrame],
-        time_dict: Dict[str, pd.Series],
-        batch_dict: Dict[str, np.ndarray],
+        feat_dict: dict[str, pd.DataFrame],
+        time_dict: dict[str, pd.Series],
+        batch_dict: dict[str, np.ndarray],
         anchor_time: pd.Series,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         if num_forecasts > 1:
             raise NotImplementedError("Forecasting not yet implemented for "
                                       "non-regression tasks")
@@ -370,13 +368,13 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute_join(
         self,
         join: Join,
-        feat_dict: Dict[str, pd.DataFrame],
-        time_dict: Dict[str, pd.Series],
-        batch_dict: Dict[str, np.ndarray],
+        feat_dict: dict[str, pd.DataFrame],
+        time_dict: dict[str, pd.Series],
+        batch_dict: dict[str, np.ndarray],
         anchor_time: pd.Series,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         if isinstance(join.rhs_target, Aggregation):
             return self.execute_aggregation(
                 aggr=join.rhs_target,
@@ -393,12 +391,12 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute_filter(
         self,
         filter: Filter,
-        feat_dict: Dict[str, pd.DataFrame],
-        time_dict: Dict[str, pd.Series],
-        batch_dict: Dict[str, np.ndarray],
+        feat_dict: dict[str, pd.DataFrame],
+        time_dict: dict[str, pd.Series],
+        batch_dict: dict[str, np.ndarray],
         anchor_time: pd.Series,
         filter_na: bool = True,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         out, mask = self.execute_column(
             column=filter.target,
             feat_dict=feat_dict,
@@ -431,12 +429,12 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute(
         self,
         query: ValidatedPredictiveQuery,
-        feat_dict: Dict[str, pd.DataFrame],
-        time_dict: Dict[str, pd.Series],
-        batch_dict: Dict[str, np.ndarray],
+        feat_dict: dict[str, pd.DataFrame],
+        time_dict: dict[str, pd.Series],
+        batch_dict: dict[str, np.ndarray],
         anchor_time: pd.Series,
         num_forecasts: int = 1,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         if isinstance(query.entity_ast, Column):
             out, mask = self.execute_column(
                 column=query.entity_ast,

kumoai/experimental/rfm/rfm.py

@@ -2,20 +2,10 @@ import json
 import time
 import warnings
 from collections import defaultdict
-from collections.abc import Generator
+from collections.abc import Generator, Iterator
 from contextlib import contextmanager
 from dataclasses import dataclass, replace
-from typing import (
-    Any,
-    Dict,
-    Iterator,
-    List,
-    Literal,
-    Optional,
-    Tuple,
-    Union,
-    overload,
-)
+from typing import Any, Literal, overload
 
 import numpy as np
 import pandas as pd
@@ -38,12 +28,13 @@ from kumoapi.rfm import (
 from kumoapi.task import TaskType
 from kumoapi.typing import AggregationType, Stype
 
+from kumoai import in_notebook, in_snowflake_notebook
 from kumoai.client.rfm import RFMAPI
 from kumoai.exceptions import HTTPException
 from kumoai.experimental.rfm import Graph
 from kumoai.experimental.rfm.base import DataBackend, Sampler
 from kumoai.mixin import CastMixin
-from kumoai.utils import InteractiveProgressLogger, ProgressLogger
+from kumoai.utils import ProgressLogger
 
 _RANDOM_SEED = 42
 
@@ -98,24 +89,41 @@ class Explanation:
     def __getitem__(self, index: Literal[1]) -> str:
         pass
 
-    def __getitem__(self, index: int) -> Union[pd.DataFrame, str]:
+    def __getitem__(self, index: int) -> pd.DataFrame | str:
         if index == 0:
             return self.prediction
         if index == 1:
             return self.summary
         raise IndexError("Index out of range")
 
-    def __iter__(self) -> Iterator[Union[pd.DataFrame, str]]:
+    def __iter__(self) -> Iterator[pd.DataFrame | str]:
         return iter((self.prediction, self.summary))
 
     def __repr__(self) -> str:
         return str((self.prediction, self.summary))
 
-    def _ipython_display_(self) -> None:
-        from IPython.display import Markdown, display
+    def print(self) -> None:
+        r"""Prints the explanation."""
+        if in_snowflake_notebook():
+            import streamlit as st
+            st.dataframe(self.prediction, hide_index=True)
+            st.markdown(self.summary)
+        elif in_notebook():
+            from IPython.display import Markdown, display
+            try:
+                if hasattr(self.prediction.style, 'hide'):
+                    display(self.prediction.hide(axis='index'))  # pandas=2
+                else:
+                    display(self.prediction.hide_index())  # pandas <1.3
+            except ImportError:
+                print(self.prediction.to_string(index=False))  # missing jinja2
+            display(Markdown(self.summary))
+        else:
+            print(self.prediction.to_string(index=False))
+            print(self.summary)
 
-        display(self.prediction)
-        display(Markdown(self.summary))
+    def _ipython_display_(self) -> None:
+        self.print()
 
 
 class KumoRFM:
@@ -162,7 +170,7 @@ class KumoRFM:
     def __init__(
         self,
         graph: Graph,
-        verbose: Union[bool, ProgressLogger] = True,
+        verbose: bool | ProgressLogger = True,
         optimize: bool = False,
     ) -> None:
         graph = graph.validate()
@@ -180,9 +188,9 @@ class KumoRFM:
         else:
             raise NotImplementedError
 
-        self._client: Optional[RFMAPI] = None
+        self._client: RFMAPI | None = None
 
-        self._batch_size: Optional[int | Literal['max']] = None
+        self._batch_size: int | Literal['max'] | None = None
         self.num_retries: int = 0
 
     @property
@@ -200,7 +208,7 @@ class KumoRFM:
     @contextmanager
     def batch_mode(
         self,
-        batch_size: Union[int, Literal['max']] = 'max',
+        batch_size: int | Literal['max'] = 'max',
         num_retries: int = 1,
     ) -> Generator[None, None, None]:
         """Context manager to predict in batches.
@@ -234,17 +242,17 @@ class KumoRFM:
     def predict(
         self,
         query: str,
-        indices: Union[List[str], List[float], List[int], None] = None,
+        indices: list[str] | list[float] | list[int] | None = None,
         *,
         explain: Literal[False] = 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,
-        num_neighbors: Optional[List[int]] = None,
+        anchor_time: pd.Timestamp | Literal['entity'] | None = None,
+        context_anchor_time: pd.Timestamp | None = None,
+        run_mode: RunMode | str = RunMode.FAST,
+        num_neighbors: list[int] | None = None,
         num_hops: int = 2,
         max_pq_iterations: int = 10,
-        random_seed: Optional[int] = _RANDOM_SEED,
-        verbose: Union[bool, ProgressLogger] = True,
+        random_seed: int | None = _RANDOM_SEED,
+        verbose: bool | ProgressLogger = True,
         use_prediction_time: bool = False,
     ) -> pd.DataFrame:
         pass
@@ -253,17 +261,17 @@ class KumoRFM:
     def predict(
         self,
         query: str,
-        indices: Union[List[str], List[float], List[int], None] = None,
+        indices: list[str] | list[float] | list[int] | None = None,
         *,
-        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,
-        num_neighbors: Optional[List[int]] = None,
+        explain: Literal[True] | ExplainConfig | dict[str, Any],
+        anchor_time: pd.Timestamp | Literal['entity'] | None = None,
+        context_anchor_time: pd.Timestamp | None = None,
+        run_mode: RunMode | str = RunMode.FAST,
+        num_neighbors: list[int] | None = None,
         num_hops: int = 2,
         max_pq_iterations: int = 10,
-        random_seed: Optional[int] = _RANDOM_SEED,
-        verbose: Union[bool, ProgressLogger] = True,
+        random_seed: int | None = _RANDOM_SEED,
+        verbose: bool | ProgressLogger = True,
         use_prediction_time: bool = False,
     ) -> Explanation:
         pass
@@ -271,19 +279,19 @@ class KumoRFM:
     def predict(
         self,
         query: str,
-        indices: Union[List[str], List[float], List[int], None] = None,
+        indices: list[str] | list[float] | list[int] | None = None,
         *,
-        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,
-        num_neighbors: Optional[List[int]] = None,
+        explain: bool | ExplainConfig | dict[str, Any] = False,
+        anchor_time: pd.Timestamp | Literal['entity'] | None = None,
+        context_anchor_time: pd.Timestamp | None = None,
+        run_mode: RunMode | str = RunMode.FAST,
+        num_neighbors: list[int] | None = None,
         num_hops: int = 2,
         max_pq_iterations: int = 10,
-        random_seed: Optional[int] = _RANDOM_SEED,
-        verbose: Union[bool, ProgressLogger] = True,
+        random_seed: int | None = _RANDOM_SEED,
+        verbose: bool | ProgressLogger = True,
         use_prediction_time: bool = False,
-    ) -> Union[pd.DataFrame, Explanation]:
+    ) -> pd.DataFrame | Explanation:
         """Returns predictions for a predictive query.
 
         Args:
@@ -325,7 +333,7 @@ class KumoRFM:
             If ``explain`` is provided, returns an :class:`Explanation` object
             containing the prediction, summary, and details.
         """
-        explain_config: Optional[ExplainConfig] = None
+        explain_config: ExplainConfig | None = None
         if explain is True:
             explain_config = ExplainConfig()
         elif explain is not False:
@@ -369,11 +377,11 @@ class KumoRFM:
             msg = f'[bold]PREDICT[/bold] {query_repr}'
 
         if not isinstance(verbose, ProgressLogger):
-            verbose = InteractiveProgressLogger(msg, verbose=verbose)
+            verbose = ProgressLogger.default(msg=msg, verbose=verbose)
 
         with verbose as logger:
 
-            batch_size: Optional[int] = None
+            batch_size: int | None = None
             if self._batch_size == 'max':
                 task_type = self._get_task_type(
                     query=query_def,
@@ -393,9 +401,9 @@ class KumoRFM:
                 logger.log(f"Splitting {len(indices):,} entities into "
                            f"{len(batches):,} batches of size {batch_size:,}")
 
-            predictions: List[pd.DataFrame] = []
-            summary: Optional[str] = None
-            details: Optional[Explanation] = None
+            predictions: list[pd.DataFrame] = []
+            summary: str | None = None
+            details: Explanation | None = None
             for i, batch in enumerate(batches):
                 # TODO Re-use the context for subsequent predictions.
                 context = self._get_context(
@@ -429,8 +437,7 @@ class KumoRFM:
                     stats = Context.get_memory_stats(request_msg.context)
                     raise ValueError(_SIZE_LIMIT_MSG.format(stats=stats))
 
-                if (isinstance(verbose, InteractiveProgressLogger) and i == 0
-                        and len(batches) > 1):
+                if i == 0 and len(batches) > 1:
                     verbose.init_progress(
                         total=len(batches),
                         description='Predicting',
@@ -469,8 +476,7 @@ class KumoRFM:
 
                         predictions.append(df)
 
-                        if (isinstance(verbose, InteractiveProgressLogger)
-                                and len(batches) > 1):
+                        if len(batches) > 1:
                             verbose.step()
 
                         break
@@ -508,9 +514,9 @@ class KumoRFM:
     def is_valid_entity(
         self,
         query: str,
-        indices: Union[List[str], List[float], List[int], None] = None,
+        indices: list[str] | list[float] | list[int] | None = None,
         *,
-        anchor_time: Union[pd.Timestamp, Literal['entity'], None] = None,
+        anchor_time: pd.Timestamp | Literal['entity'] | None = None,
     ) -> np.ndarray:
         r"""Returns a mask that denotes which entities are valid for the
         given predictive query, *i.e.*, which entities fulfill (temporal)
@@ -554,15 +560,15 @@ class KumoRFM:
         self,
         query: str,
         *,
-        metrics: Optional[List[str]] = None,
-        anchor_time: Union[pd.Timestamp, Literal['entity'], None] = None,
-        context_anchor_time: Union[pd.Timestamp, None] = None,
-        run_mode: Union[RunMode, str] = RunMode.FAST,
-        num_neighbors: Optional[List[int]] = None,
+        metrics: list[str] | None = None,
+        anchor_time: pd.Timestamp | Literal['entity'] | None = None,
+        context_anchor_time: pd.Timestamp | None = None,
+        run_mode: RunMode | str = RunMode.FAST,
+        num_neighbors: list[int] | None = None,
         num_hops: int = 2,
         max_pq_iterations: int = 10,
-        random_seed: Optional[int] = _RANDOM_SEED,
-        verbose: Union[bool, ProgressLogger] = True,
+        random_seed: int | None = _RANDOM_SEED,
+        verbose: bool | ProgressLogger = True,
         use_prediction_time: bool = False,
     ) -> pd.DataFrame:
         """Evaluates a predictive query.
@@ -610,7 +616,7 @@ class KumoRFM:
         msg = f'[bold]EVALUATE[/bold] {query_repr}'
 
         if not isinstance(verbose, ProgressLogger):
-            verbose = InteractiveProgressLogger(msg, verbose=verbose)
+            verbose = ProgressLogger.default(msg=msg, verbose=verbose)
 
         with verbose as logger:
             context = self._get_context(
@@ -669,8 +675,8 @@ class KumoRFM:
         query: str,
         size: int,
         *,
-        anchor_time: Union[pd.Timestamp, Literal['entity'], None] = None,
-        random_seed: Optional[int] = _RANDOM_SEED,
+        anchor_time: pd.Timestamp | Literal['entity'] | None = None,
+        random_seed: int | None = _RANDOM_SEED,
         max_iterations: int = 10,
     ) -> pd.DataFrame:
         """Returns the labels of a predictive query for a specified anchor
@@ -764,7 +770,7 @@ class KumoRFM:
     @staticmethod
     def _get_task_type(
         query: ValidatedPredictiveQuery,
-        edge_types: List[Tuple[str, str, str]],
+        edge_types: list[tuple[str, str, str]],
     ) -> TaskType:
         if isinstance(query.target_ast, (Condition, LogicalOperation)):
             return TaskType.BINARY_CLASSIFICATION
@@ -819,7 +825,7 @@ class KumoRFM:
         self,
         query: ValidatedPredictiveQuery,
         anchor_time: pd.Timestamp,
-        context_anchor_time: Union[pd.Timestamp, None],
+        context_anchor_time: pd.Timestamp | None,
         evaluate: bool,
     ) -> None:
 
@@ -885,16 +891,16 @@ class KumoRFM:
     def _get_context(
         self,
         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],
+        indices: list[str] | list[float] | list[int] | None,
+        anchor_time: pd.Timestamp | Literal['entity'] | None,
+        context_anchor_time: pd.Timestamp | None,
         run_mode: RunMode,
-        num_neighbors: Optional[List[int]],
+        num_neighbors: list[int] | None,
         num_hops: int,
         max_pq_iterations: int,
         evaluate: bool,
-        random_seed: Optional[int] = _RANDOM_SEED,
-        logger: Optional[ProgressLogger] = None,
+        random_seed: int | None = _RANDOM_SEED,
+        logger: ProgressLogger | None = None,
     ) -> Context:
 
         if num_neighbors is not None:
@@ -1069,7 +1075,7 @@ class KumoRFM:
                 raise NotImplementedError
             logger.log(msg)
 
-        entity_table_names: Tuple[str, ...]
+        entity_table_names: tuple[str, ...]
         if task_type.is_link_pred:
             final_aggr = query.get_final_target_aggregation()
             assert final_aggr is not None
@@ -1127,7 +1133,7 @@ class KumoRFM:
 
     @staticmethod
     def _validate_metrics(
-        metrics: List[str],
+        metrics: list[str],
         task_type: TaskType,
     ) -> None:
 
@@ -1184,7 +1190,7 @@ class KumoRFM:
                                  f"'https://github.com/kumo-ai/kumo-rfm'.")
 
 
-def format_value(value: Union[int, float]) -> str:
+def format_value(value: int | float) -> str:
     if value == int(value):
         return f'{int(value):,}'
     if abs(value) >= 1000:

kumoai/experimental/rfm/sagemaker.py

@@ -1,6 +1,6 @@
 import base64
 import json
-from typing import Any, Dict, List, Tuple
+from typing import Any
 
 import requests
 
@@ -48,8 +48,8 @@ class KumoClient_SageMakerAdapter(KumoClient):
 
         # Recording buffers.
         self._recording_active = False
-        self._recorded_reqs: List[Dict[str, Any]] = []
-        self._recorded_resps: List[Dict[str, Any]] = []
+        self._recorded_reqs: list[dict[str, Any]] = []
+        self._recorded_resps: list[dict[str, Any]] = []
 
     def authenticate(self) -> None:
         # TODO(siyang): call /ping to verify?
@@ -92,7 +92,7 @@ class KumoClient_SageMakerAdapter(KumoClient):
         self._recorded_reqs.clear()
         self._recorded_resps.clear()
 
-    def end_recording(self) -> List[Tuple[Dict[str, Any], Dict[str, Any]]]:
+    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