arize-phoenix 5.5.2__py3-none-any.whl → 5.7.0__py3-none-any.whl

phoenix/inferences/inferences.py

@@ -5,7 +5,7 @@ from copy import deepcopy
 from dataclasses import dataclass, fields, replace
 from enum import Enum
 from itertools import groupby
-from typing import Any, Dict, List, Optional, Set, Tuple, Union
+from typing import Any, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -154,7 +154,7 @@ class Inferences:
     @deprecated("Inferences.from_open_inference is deprecated and will be removed.")
     def from_open_inference(cls, dataframe: DataFrame) -> "Inferences":
         schema = Schema()
-        column_renaming: Dict[str, str] = {}
+        column_renaming: dict[str, str] = {}
         for group_name, group in groupby(
             sorted(
                 map(_parse_open_inference_column_name, dataframe.columns),
@@ -351,7 +351,7 @@ def _parse_open_inference_column_name(column_name: str) -> _OpenInferenceColumnN
     raise ValueError(f"Invalid format for column name: {column_name}")
 
 
-def _parse_dataframe_and_schema(dataframe: DataFrame, schema: Schema) -> Tuple[DataFrame, Schema]:
+def _parse_dataframe_and_schema(dataframe: DataFrame, schema: Schema) -> tuple[DataFrame, Schema]:
     """
     Parses a dataframe according to a schema, infers feature columns names when
     they are not explicitly provided, and removes excluded column names from
@@ -364,12 +364,12 @@ def _parse_dataframe_and_schema(dataframe: DataFrame, schema: Schema) -> Tuple[D
     names present in the dataframe but not included in any other schema fields.
     """
 
-    unseen_excluded_column_names: Set[str] = (
+    unseen_excluded_column_names: set[str] = (
         set(schema.excluded_column_names) if schema.excluded_column_names is not None else set()
     )
-    unseen_column_names: Set[str] = set(dataframe.columns.to_list())
-    column_name_to_include: Dict[str, bool] = {}
-    schema_patch: Dict[SchemaFieldName, SchemaFieldValue] = {}
+    unseen_column_names: set[str] = set(dataframe.columns.to_list())
+    column_name_to_include: dict[str, bool] = {}
+    schema_patch: dict[SchemaFieldName, SchemaFieldValue] = {}
 
     for schema_field_name in SINGLE_COLUMN_SCHEMA_FIELD_NAMES:
         _check_single_column_schema_field_for_excluded_columns(
@@ -434,10 +434,10 @@ def _parse_dataframe_and_schema(dataframe: DataFrame, schema: Schema) -> Tuple[D
 def _check_single_column_schema_field_for_excluded_columns(
     schema: Schema,
     schema_field_name: str,
-    unseen_excluded_column_names: Set[str],
-    schema_patch: Dict[SchemaFieldName, SchemaFieldValue],
-    column_name_to_include: Dict[str, bool],
-    unseen_column_names: Set[str],
+    unseen_excluded_column_names: set[str],
+    schema_patch: dict[SchemaFieldName, SchemaFieldValue],
+    column_name_to_include: dict[str, bool],
+    unseen_column_names: set[str],
 ) -> None:
     """
     Checks single-column schema fields for excluded column names.
@@ -455,18 +455,18 @@ def _check_single_column_schema_field_for_excluded_columns(
 def _check_multi_column_schema_field_for_excluded_columns(
     schema: Schema,
     schema_field_name: str,
-    unseen_excluded_column_names: Set[str],
-    schema_patch: Dict[SchemaFieldName, SchemaFieldValue],
-    column_name_to_include: Dict[str, bool],
-    unseen_column_names: Set[str],
+    unseen_excluded_column_names: set[str],
+    schema_patch: dict[SchemaFieldName, SchemaFieldValue],
+    column_name_to_include: dict[str, bool],
+    unseen_column_names: set[str],
 ) -> None:
     """
     Checks multi-column schema fields for excluded columns names.
     """
-    column_names: Optional[List[str]] = getattr(schema, schema_field_name)
+    column_names: Optional[list[str]] = getattr(schema, schema_field_name)
     if column_names:
-        included_column_names: List[str] = []
-        excluded_column_names: List[str] = []
+        included_column_names: list[str] = []
+        excluded_column_names: list[str] = []
         for column_name in column_names:
             is_included_column = column_name not in unseen_excluded_column_names
             column_name_to_include[column_name] = is_included_column
@@ -482,10 +482,10 @@ def _check_multi_column_schema_field_for_excluded_columns(
 
 def _check_embedding_features_schema_field_for_excluded_columns(
     embedding_features: EmbeddingFeatures,
-    unseen_excluded_column_names: Set[str],
-    schema_patch: Dict[SchemaFieldName, SchemaFieldValue],
-    column_name_to_include: Dict[str, bool],
-    unseen_column_names: Set[str],
+    unseen_excluded_column_names: set[str],
+    schema_patch: dict[SchemaFieldName, SchemaFieldValue],
+    column_name_to_include: dict[str, bool],
+    unseen_column_names: set[str],
 ) -> None:
     """
     Check embedding features for excluded column names.
@@ -527,8 +527,8 @@ def _check_embedding_features_schema_field_for_excluded_columns(
 
 def _check_embedding_column_names_for_excluded_columns(
     embedding_column_name_mapping: EmbeddingColumnNames,
-    column_name_to_include: Dict[str, bool],
-    unseen_column_names: Set[str],
+    column_name_to_include: dict[str, bool],
+    unseen_column_names: set[str],
 ) -> None:
     """
     Check embedding column names for excluded column names.
@@ -542,10 +542,10 @@ def _check_embedding_column_names_for_excluded_columns(
 
 def _discover_feature_columns(
     dataframe: DataFrame,
-    unseen_excluded_column_names: Set[str],
-    schema_patch: Dict[SchemaFieldName, SchemaFieldValue],
-    column_name_to_include: Dict[str, bool],
-    unseen_column_names: Set[str],
+    unseen_excluded_column_names: set[str],
+    schema_patch: dict[SchemaFieldName, SchemaFieldValue],
+    column_name_to_include: dict[str, bool],
+    unseen_column_names: set[str],
 ) -> None:
     """
     Adds unseen and un-excluded columns as features, with the exception of "prediction_id"
@@ -559,10 +559,10 @@ def _discover_feature_columns(
         else:
             unseen_excluded_column_names.discard(column_name)
             logger.debug(f"excluded feature: {column_name}")
-    original_column_positions: List[int] = dataframe.columns.get_indexer(
+    original_column_positions: list[int] = dataframe.columns.get_indexer(
         discovered_feature_column_names
     )  # type: ignore
-    feature_column_name_to_position: Dict[str, int] = dict(
+    feature_column_name_to_position: dict[str, int] = dict(
         zip(discovered_feature_column_names, original_column_positions)
     )
     discovered_feature_column_names.sort(key=lambda col: feature_column_name_to_position[col])
@@ -575,16 +575,16 @@ def _discover_feature_columns(
 def _create_and_normalize_dataframe_and_schema(
     dataframe: DataFrame,
     schema: Schema,
-    schema_patch: Dict[SchemaFieldName, SchemaFieldValue],
-    column_name_to_include: Dict[str, bool],
-) -> Tuple[DataFrame, Schema]:
+    schema_patch: dict[SchemaFieldName, SchemaFieldValue],
+    column_name_to_include: dict[str, bool],
+) -> tuple[DataFrame, Schema]:
     """
     Creates new dataframe and schema objects to reflect excluded column names
     and discovered features. This also normalizes dataframe columns to ensure a
     standard set of columns (i.e. timestamp and prediction_id) and datatypes for
     those columns.
     """
-    included_column_names: List[str] = []
+    included_column_names: list[str] = []
     for column_name in dataframe.columns:
         if column_name_to_include.get(str(column_name), False):
             included_column_names.append(str(column_name))
@@ -648,7 +648,7 @@ def _normalize_timestamps(
     dataframe: DataFrame,
     schema: Schema,
     default_timestamp: Timestamp,
-) -> Tuple[DataFrame, Schema]:
+) -> tuple[DataFrame, Schema]:
     """
     Ensures that the dataframe has a timestamp column and the schema has a timestamp field. If the
     input dataframe contains a Unix or datetime timestamp or ISO8601 timestamp strings column, it
@@ -686,7 +686,7 @@ def _get_schema_from_unknown_schema_param(schemaLike: SchemaLike) -> Schema:
         if not isinstance(schemaLike, ArizeSchema):
             raise ValueError("Unknown schema passed to Dataset. Please pass a phoenix Schema")
 
-        embedding_feature_column_names: Dict[str, EmbeddingColumnNames] = {}
+        embedding_feature_column_names: dict[str, EmbeddingColumnNames] = {}
         if schemaLike.embedding_feature_column_names is not None:
             for (
                 embedding_name,
@@ -734,7 +734,7 @@ def _get_schema_from_unknown_schema_param(schemaLike: SchemaLike) -> Schema:
         )
 
 
-def _add_prediction_id(num_rows: int) -> List[str]:
+def _add_prediction_id(num_rows: int) -> list[str]:
     return [str(uuid.uuid4()) for _ in range(num_rows)]
 
 

phoenix/inferences/schema.py

@@ -1,13 +1,14 @@
 import json
+from collections.abc import Mapping
 from dataclasses import asdict, dataclass, replace
-from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
+from typing import Any, Optional, Union
 
-EmbeddingFeatures = Dict[str, "EmbeddingColumnNames"]
+EmbeddingFeatures = dict[str, "EmbeddingColumnNames"]
 SchemaFieldName = str
-SchemaFieldValue = Union[Optional[str], Optional[List[str]], Optional[EmbeddingFeatures]]
+SchemaFieldValue = Union[Optional[str], Optional[list[str]], Optional[EmbeddingFeatures]]
 
-MULTI_COLUMN_SCHEMA_FIELD_NAMES: Tuple[str, ...] = ("feature_column_names", "tag_column_names")
-SINGLE_COLUMN_SCHEMA_FIELD_NAMES: Tuple[str, ...] = (
+MULTI_COLUMN_SCHEMA_FIELD_NAMES: tuple[str, ...] = ("feature_column_names", "tag_column_names")
+SINGLE_COLUMN_SCHEMA_FIELD_NAMES: tuple[str, ...] = (
     "prediction_id_column_name",
     "timestamp_column_name",
     "prediction_label_column_name",
@@ -19,7 +20,7 @@ LLM_SCHEMA_FIELD_NAMES = ["prompt_column_names", "response_column_names"]
 
 
 @dataclass(frozen=True)
-class EmbeddingColumnNames(Dict[str, Any]):
+class EmbeddingColumnNames(dict[str, Any]):
     """
     A dataclass to hold the column names for the embedding features.
     An embedding feature is a feature that is represented by a vector.
@@ -80,8 +81,8 @@ class Schema:
     prediction_id_column_name: Optional[str] = None
     id_column_name: Optional[str] = None  # Syntax sugar for prediction_id_column_name
     timestamp_column_name: Optional[str] = None
-    feature_column_names: Optional[List[str]] = None
-    tag_column_names: Optional[List[str]] = None
+    feature_column_names: Optional[list[str]] = None
+    tag_column_names: Optional[list[str]] = None
     prediction_label_column_name: Optional[str] = None
     prediction_score_column_name: Optional[str] = None
     actual_label_column_name: Optional[str] = None
@@ -91,7 +92,7 @@ class Schema:
     # document_column_names is used explicitly when the schema is used to capture a corpus
     document_column_names: Optional[EmbeddingColumnNames] = None
     embedding_feature_column_names: Optional[EmbeddingFeatures] = None
-    excluded_column_names: Optional[List[str]] = None
+    excluded_column_names: Optional[list[str]] = None
 
     def __post_init__(self) -> None:
         # re-map document_column_names to be in the prompt_column_names position
@@ -107,7 +108,7 @@ class Schema:
     def replace(self, **changes: Any) -> "Schema":
         return replace(self, **changes)
 
-    def asdict(self) -> Dict[str, str]:
+    def asdict(self) -> dict[str, str]:
         return asdict(self)
 
     def to_json(self) -> str:

phoenix/inferences/validation.py

@@ -1,5 +1,4 @@
 import math
-from typing import List
 
 import numpy as np
 from pandas import DataFrame, Series
@@ -11,8 +10,8 @@ from .schema import EmbeddingColumnNames, Schema
 RESERVED_EMBEDDING_NAMES = ("prompt", "response")
 
 
-def _check_valid_schema(schema: Schema) -> List[err.ValidationError]:
-    errs: List[str] = []
+def _check_valid_schema(schema: Schema) -> list[err.ValidationError]:
+    errs: list[str] = []
     if schema.excluded_column_names is None:
         return []
 
@@ -34,7 +33,7 @@ def _check_valid_schema(schema: Schema) -> List[err.ValidationError]:
     return []
 
 
-def validate_inferences_inputs(dataframe: DataFrame, schema: Schema) -> List[err.ValidationError]:
+def validate_inferences_inputs(dataframe: DataFrame, schema: Schema) -> list[err.ValidationError]:
     errors = _check_missing_columns(dataframe, schema)
     if errors:
         return errors
@@ -53,12 +52,12 @@ def validate_inferences_inputs(dataframe: DataFrame, schema: Schema) -> List[err
     return []
 
 
-def _check_valid_embedding_data(dataframe: DataFrame, schema: Schema) -> List[err.ValidationError]:
+def _check_valid_embedding_data(dataframe: DataFrame, schema: Schema) -> list[err.ValidationError]:
     embedding_col_names = schema.embedding_feature_column_names
     if embedding_col_names is None:
         return []
 
-    embedding_errors: List[err.ValidationError] = []
+    embedding_errors: list[err.ValidationError] = []
     for embedding_name, column_names in embedding_col_names.items():
         if embedding_name in RESERVED_EMBEDDING_NAMES:
             embedding_errors += _validate_reserved_embedding_name(embedding_name, schema)
@@ -71,8 +70,8 @@ def _check_valid_embedding_data(dataframe: DataFrame, schema: Schema) -> List[er
 
 def _check_valid_prompt_response_data(
     dataframe: DataFrame, schema: Schema
-) -> List[err.ValidationError]:
-    prompt_response_errors: List[err.ValidationError] = []
+) -> list[err.ValidationError]:
+    prompt_response_errors: list[err.ValidationError] = []
 
     prompt_response_column_names = {
         "prompt": schema.prompt_column_names,
@@ -89,7 +88,7 @@ def _check_valid_prompt_response_data(
 
 def _validate_reserved_embedding_name(
     embedding_name: str, schema: Schema
-) -> List[err.ValidationError]:
+) -> list[err.ValidationError]:
     if embedding_name == "prompt" and schema.prompt_column_names is not None:
         return [err.InvalidEmbeddingReservedName(embedding_name, "schema.prompt_column_names")]
     elif embedding_name == "response" and schema.response_column_names is not None:
@@ -99,9 +98,9 @@ def _validate_reserved_embedding_name(
 
 def _validate_embedding_vector(
     dataframe: DataFrame, name: str, vector_column_name: str
-) -> List[err.ValidationError]:
+) -> list[err.ValidationError]:
     vector_column = dataframe[vector_column_name]
-    errors: List[err.ValidationError] = []
+    errors: list[err.ValidationError] = []
     vector_length = None
 
     for vector in vector_column:
@@ -156,8 +155,8 @@ def _validate_embedding_vector(
     return errors
 
 
-def _check_column_types(dataframe: DataFrame, schema: Schema) -> List[err.ValidationError]:
-    wrong_type_cols: List[str] = []
+def _check_column_types(dataframe: DataFrame, schema: Schema) -> list[err.ValidationError]:
+    wrong_type_cols: list[str] = []
     if schema.prediction_id_column_name is not None:
         if not (
             is_numeric_dtype(dataframe.dtypes[schema.prediction_id_column_name])
@@ -172,7 +171,7 @@ def _check_column_types(dataframe: DataFrame, schema: Schema) -> List[err.Valida
     return []
 
 
-def _check_missing_columns(dataframe: DataFrame, schema: Schema) -> List[err.ValidationError]:
+def _check_missing_columns(dataframe: DataFrame, schema: Schema) -> list[err.ValidationError]:
     # converting to a set first makes the checks run a lot faster
     existing_columns = set(dataframe.columns)
     missing_columns = []

phoenix/logging/_formatter.py

@@ -1,7 +1,7 @@
 import datetime as dt
 import json
 import logging
-from typing import Dict, Optional
+from typing import Optional
 
 LOG_RECORD_BUILTIN_ATTRS = {
     "args",
@@ -34,7 +34,7 @@ class PhoenixJSONFormatter(logging.Formatter):
     def __init__(
         self,
         *,
-        fmt_keys: Optional[Dict[str, str]] = None,
+        fmt_keys: Optional[dict[str, str]] = None,
     ):
         super().__init__()
         self.fmt_keys = fmt_keys if fmt_keys is not None else {}
@@ -43,7 +43,7 @@ class PhoenixJSONFormatter(logging.Formatter):
         message = self._prepare_log_dict(record)
         return json.dumps(message, default=str)
 
-    def _prepare_log_dict(self, record: logging.LogRecord) -> Dict[str, str]:
+    def _prepare_log_dict(self, record: logging.LogRecord) -> dict[str, str]:
         always_fields = {
             "message": record.getMessage(),
             "timestamp": dt.datetime.fromtimestamp(record.created, tz=dt.timezone.utc).isoformat(),

phoenix/metrics/__init__.py

@@ -1,8 +1,9 @@
 import logging
 import warnings
 from abc import ABC, abstractmethod
+from collections.abc import Iterable, Mapping
 from dataclasses import dataclass
-from typing import Any, Iterable, List, Mapping, Optional, Union
+from typing import Any, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -36,13 +37,13 @@ class Metric(ABC):
     def calc(self, dataframe: pd.DataFrame) -> Any: ...
 
     @abstractmethod
-    def operands(self) -> List[Column]: ...
+    def operands(self) -> list[Column]: ...
 
     def __call__(
         self,
         df: pd.DataFrame,
         /,
-        subset_rows: Optional[Union[slice, List[int]]] = None,
+        subset_rows: Optional[Union[slice, list[int]]] = None,
     ) -> Any:
         """
         Computes the metric on a dataframe.
@@ -51,7 +52,7 @@ class Metric(ABC):
         ----------
         df: pandas DataFrame
             The dataframe input to the metric.
-        subset_rows: Optional[Union[slice, List[int]]] = None
+        subset_rows: Optional[Union[slice, list[int]]] = None
             Optionally specifying a subset of rows for the computation.
             Can be a list or slice (e.g. `slice(100, 200)`) of integers.
         """

phoenix/metrics/binning.py

@@ -1,8 +1,9 @@
 import warnings
 from abc import ABC, abstractmethod
+from collections.abc import Iterable, Sequence
 from dataclasses import dataclass
 from functools import partial
-from typing import Any, Iterable, Optional, Sequence, cast
+from typing import Any, Optional, cast
 
 import numpy as np
 import pandas as pd

phoenix/metrics/metrics.py

@@ -1,8 +1,9 @@
 import math
 import warnings
+from collections.abc import Callable
 from dataclasses import dataclass, field
 from functools import cached_property
-from typing import Callable, Union, cast
+from typing import Union, cast
 
 import numpy as np
 import numpy.typing as npt

phoenix/metrics/mixins.py

@@ -7,10 +7,11 @@ on cooperative multiple inheritance and method resolution order in Python.
 import collections
 import inspect
 from abc import ABC, abstractmethod
+from collections.abc import Callable
 from dataclasses import dataclass, field, fields, replace
 from functools import cached_property
 from itertools import repeat
-from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Mapping, Optional
+from typing import TYPE_CHECKING, Any, Iterator, Mapping, Optional
 
 import numpy as np
 import pandas as pd
@@ -42,7 +43,7 @@ class VectorOperator(ABC):
 
 @dataclass(frozen=True)
 class NullaryOperator(Metric, ABC):
-    def operands(self) -> List[Column]:
+    def operands(self) -> list[Column]:
         return []
 
 
@@ -55,7 +56,7 @@ class UnaryOperator(Metric, ABC):
 
     operand: Column = Column()
 
-    def operands(self) -> List[Column]:
+    def operands(self) -> list[Column]:
         return [self.operand]
 
 
@@ -98,10 +99,10 @@ class EvaluationMetricKeywordParameters(_BaseMapping):
         return sum(1 for _ in self)
 
     @property
-    def columns(self) -> List[Column]:
+    def columns(self) -> list[Column]:
         return [v for v in self.values() if isinstance(v, Column)]
 
-    def __call__(self, df: pd.DataFrame) -> Dict[str, Any]:
+    def __call__(self, df: pd.DataFrame) -> dict[str, Any]:
         return {k: v(df) if isinstance(v, Column) else v for k, v in self.items()}
 
 
@@ -142,7 +143,7 @@ class EvaluationMetric(Metric, ABC):
                 ),
             )
 
-    def operands(self) -> List[Column]:
+    def operands(self) -> list[Column]:
         return [self.actual, self.predicted] + self.parameters.columns
 
     def calc(self, df: pd.DataFrame) -> float:

phoenix/metrics/retrieval_metrics.py

@@ -1,5 +1,6 @@
+from collections.abc import Iterable
 from dataclasses import dataclass, field
-from typing import Iterable, Optional, cast
+from typing import Optional, cast
 
 import numpy as np
 import pandas as pd

phoenix/metrics/timeseries.py

@@ -1,7 +1,8 @@
+from collections.abc import Callable, Iterable, Iterator
 from datetime import datetime, timedelta, timezone
 from functools import partial
 from itertools import accumulate, repeat
-from typing import Callable, Iterable, Iterator, Tuple, cast
+from typing import cast
 
 import pandas as pd
 from typing_extensions import TypeAlias
@@ -41,12 +42,12 @@ def row_interval_from_sorted_time_index(
     time_index: pd.DatetimeIndex,
     time_start: datetime,
     time_stop: datetime,
-) -> Tuple[StartIndex, StopIndex]:
+) -> tuple[StartIndex, StopIndex]:
     """
     Returns end exclusive time slice from sorted index.
     """
     return cast(
-        Tuple[StartIndex, StopIndex],
+        tuple[StartIndex, StopIndex],
         time_index.searchsorted((time_start, time_stop)),
     )
 
@@ -86,7 +87,7 @@ def _groupers(
     end_time: datetime,
     evaluation_window: timedelta,
     sampling_interval: timedelta,
-) -> Iterator[Tuple[StartTime, EndTime, pd.Grouper]]:
+) -> Iterator[tuple[StartTime, EndTime, pd.Grouper]]:
     """
     Yields pandas.Groupers from time series parameters.
     """

phoenix/metrics/wrappers.py

@@ -18,7 +18,7 @@ from abc import ABC
 from enum import Enum
 from inspect import Signature
 from itertools import chain, islice
-from typing import Any, Dict, List, Tuple, cast
+from typing import Any, cast
 
 import numpy as np
 import pandas as pd
@@ -157,7 +157,7 @@ def _coerce_dtype_if_necessary(
 def _eliminate_missing_values_from_all_series(
     *args: Any,
     **kwargs: Any,
-) -> Tuple[List[Any], Dict[str, Any]]:
+) -> tuple[list[Any], dict[str, Any]]:
     positional_arguments = list(args)
     keyword_arguments = dict(kwargs)
     all_series = [

phoenix/pointcloud/clustering.py

@@ -1,12 +1,11 @@
 from dataclasses import asdict, dataclass
-from typing import List, Set
 
 import numpy as np
 import numpy.typing as npt
 from typing_extensions import TypeAlias
 
 RowIndex: TypeAlias = int
-RawCluster: TypeAlias = Set[RowIndex]
+RawCluster: TypeAlias = set[RowIndex]
 Matrix: TypeAlias = npt.NDArray[np.float64]
 
 
@@ -16,11 +15,11 @@ class Hdbscan:
     min_samples: float = 1
     cluster_selection_epsilon: float = 0.0
 
-    def find_clusters(self, mat: Matrix) -> List[RawCluster]:
+    def find_clusters(self, mat: Matrix) -> list[RawCluster]:
         from fast_hdbscan import HDBSCAN
 
         cluster_ids: npt.NDArray[np.int_] = HDBSCAN(**asdict(self)).fit_predict(mat)
-        ans: List[RawCluster] = [set() for _ in range(np.max(cluster_ids) + 1)]
+        ans: list[RawCluster] = [set() for _ in range(np.max(cluster_ids) + 1)]
         for row_idx, cluster_id in enumerate(cluster_ids):
             if cluster_id > -1:
                 ans[cluster_id].add(row_idx)

phoenix/pointcloud/pointcloud.py

@@ -1,9 +1,9 @@
+from collections.abc import Hashable, Mapping
 from dataclasses import dataclass
-from typing import Dict, List, Mapping, Protocol, Set, Tuple
+from typing import Protocol, TypeVar
 
 import numpy as np
 import numpy.typing as npt
-from strawberry import ID
 from typing_extensions import TypeAlias
 
 from phoenix.pointcloud.clustering import RawCluster
@@ -12,13 +12,15 @@ Vector: TypeAlias = npt.NDArray[np.float64]
 Matrix: TypeAlias = npt.NDArray[np.float64]
 RowIndex: TypeAlias = int
 
+_IdType = TypeVar("_IdType", bound=Hashable)
+
 
 class DimensionalityReducer(Protocol):
     def project(self, mat: Matrix, n_components: int) -> Matrix: ...
 
 
 class ClustersFinder(Protocol):
-    def find_clusters(self, mat: Matrix) -> List[RawCluster]: ...
+    def find_clusters(self, mat: Matrix) -> list[RawCluster]: ...
 
 
 @dataclass(frozen=True)
@@ -28,9 +30,9 @@ class PointCloud:
 
     def generate(
         self,
-        data: Mapping[ID, Vector],
+        data: Mapping[_IdType, Vector],
         n_components: int = 3,
-    ) -> Tuple[Dict[ID, Vector], Dict[str, Set[ID]]]:
+    ) -> tuple[dict[_IdType, Vector], dict[str, set[_IdType]]]:
         """
         Given a set of vectors, projects them onto lower dimensions, and
         finds clusters among the projections.

phoenix/pointcloud/umap_parameters.py

@@ -1,5 +1,6 @@
+from collections.abc import Mapping
 from dataclasses import dataclass
-from typing import Any, Mapping, Optional
+from typing import Any, Optional
 
 DEFAULT_MIN_DIST = 0.0
 DEFAULT_N_NEIGHBORS = 30