arize-phoenix 3.16.0__py3-none-any.whl → 7.7.0__py3-none-any.whl

phoenix/core/model_schema.py

@@ -4,6 +4,7 @@ import re
 import threading
 from abc import ABC, abstractmethod
 from collections import defaultdict
+from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass, field, fields, replace
 from datetime import datetime, timedelta, timezone
@@ -14,19 +15,9 @@ from random import random
 from typing import (
     Any,
     BinaryIO,
-    Callable,
-    Dict,
     Generic,
-    Hashable,
-    Iterable,
-    Iterator,
-    List,
-    Mapping,
     NamedTuple,
     Optional,
-    Sequence,
-    Tuple,
-    Type,
     TypeVar,
     Union,
     cast,
@@ -48,7 +39,7 @@ from pandas.core.dtypes.common import (
 from typing_extensions import TypeAlias, TypeGuard
 from wrapt import ObjectProxy
 
-from phoenix.config import GENERATED_DATASET_NAME_PREFIX
+from phoenix.config import GENERATED_INFERENCES_NAME_PREFIX
 from phoenix.datetime_utils import floor_to_minute
 
 
@@ -185,7 +176,7 @@ class RetrievalEmbedding(Embedding):
                 yield value
 
 
-class DatasetRole(Enum):
+class InferencesRole(Enum):
     """A dataframe's role in a Model: primary or reference (as
     baseline for drift).
     """
@@ -194,8 +185,8 @@ class DatasetRole(Enum):
     REFERENCE = auto()
 
 
-PRIMARY = DatasetRole.PRIMARY
-REFERENCE = DatasetRole.REFERENCE
+PRIMARY = InferencesRole.PRIMARY
+REFERENCE = InferencesRole.REFERENCE
 
 
 @dataclass(frozen=True, repr=False, eq=False)
@@ -265,7 +256,7 @@ class _Cache(Generic[_Key, _Value]):
     2
     """
 
-    _cache: Dict[_Key, _Value] = field(
+    _cache: dict[_Key, _Value] = field(
         init=False,
         default_factory=dict,
     )
@@ -275,7 +266,7 @@ class _Cache(Generic[_Key, _Value]):
     )
 
     @contextmanager
-    def __call__(self) -> Iterator[Dict[_Key, _Value]]:
+    def __call__(self) -> Iterator[dict[_Key, _Value]]:
         with self._lock:
             yield self._cache
 
@@ -381,7 +372,7 @@ class Dimension(Column, ABC):
             # But we really want the role to be specified for a Dimension.
             raise ValueError("role must be assigned")
 
-    def __getitem__(self, df_role: DatasetRole) -> "pd.Series[Any]":
+    def __getitem__(self, df_role: InferencesRole) -> "pd.Series[Any]":
         if self._model is None:
             return pd.Series(dtype=object)
         model = cast(Model, self._model)
@@ -405,18 +396,18 @@ class Dimension(Column, ABC):
 @dataclass(frozen=True)
 class ScalarDimension(Dimension):
     @property
-    def min_max(self) -> Tuple[Any, Any]:
+    def min_max(self) -> tuple[Any, Any]:
         if self._model is None:
             return np.nan, np.nan
         model = cast(Model, self._model)
         return model.dimension_min_max_from_all_df(self.name)
 
     @property
-    def categories(self) -> Tuple[str, ...]:
+    def categories(self) -> tuple[str, ...]:
         if self._model is None or self.data_type is CONTINUOUS:
             return ()
         model = cast(Model, self._model)
-        return model.dimension_categories_from_all_datasets(self.name)
+        return model.dimension_categories_from_all_inferences(self.name)
 
 
 @dataclass(frozen=True)
@@ -502,7 +493,7 @@ class RetrievalEmbeddingDimension(EmbeddingDimension):
 Name: TypeAlias = str
 ColumnKey: TypeAlias = Union[Name, Column, SingularDimensionalRole]
 MultiDimensionKey: TypeAlias = Union[MultiDimensionalRole, Sequence[DimensionRole]]
-RowNumbering: TypeAlias = Union[int, List[int]]
+RowNumbering: TypeAlias = Union[int, list[int]]
 
 
 def _is_column_key(key: Any) -> TypeGuard[ColumnKey]:
@@ -525,7 +516,7 @@ def _is_multi_dimension_key(
 
 def _is_dimension_type_filter(
     key: Any,
-) -> TypeGuard[Tuple[MultiDimensionKey, Union[Type[ScalarDimension], Type[EmbeddingDimension]]]]:
+) -> TypeGuard[tuple[MultiDimensionKey, Union[type[ScalarDimension], type[EmbeddingDimension]]]]:
     return (
         isinstance(key, tuple)
         and len(key) == 2
@@ -534,7 +525,7 @@ def _is_dimension_type_filter(
     )
 
 
-def _is_named_df(obj: Any) -> TypeGuard[Tuple[Name, pd.DataFrame]]:
+def _is_named_df(obj: Any) -> TypeGuard[tuple[Name, pd.DataFrame]]:
     return (
         isinstance(obj, tuple)
         and len(obj) == 2
@@ -582,7 +573,7 @@ class EventId(NamedTuple):
     """Identifies an event."""
 
     row_id: int = 0
-    dataset_id: DatasetRole = PRIMARY
+    inferences_id: InferencesRole = PRIMARY
 
     def __str__(self) -> str:
         return ":".join(map(str, self))
@@ -625,7 +616,7 @@ class Events(ModelData):
         self,
         df: pd.DataFrame,
         /,
-        role: DatasetRole,
+        role: InferencesRole,
         **kwargs: Any,
     ) -> None:
         super().__init__(df, **kwargs)
@@ -664,7 +655,7 @@ class Events(ModelData):
     def __getitem__(self, key: ColumnKey) -> "pd.Series[Any]": ...
 
     @overload
-    def __getitem__(self, key: List[RowId]) -> "Events": ...
+    def __getitem__(self, key: list[RowId]) -> "Events": ...
 
     def __getitem__(self, key: Any) -> Any:
         if isinstance(key, list):
@@ -676,7 +667,7 @@ class Events(ModelData):
         return super().__getitem__(key)
 
 
-class Dataset(Events):
+class Inferences(Events):
     """pd.DataFrame wrapped with extra functions and metadata."""
 
     def __init__(
@@ -701,13 +692,13 @@ class Dataset(Events):
         friendly. Falls back to the role of the dataset if no name is provided.
         """
         ds_name = self._self_name
-        if ds_name.startswith(GENERATED_DATASET_NAME_PREFIX):
+        if ds_name.startswith(GENERATED_INFERENCES_NAME_PREFIX):
             # The generated names are UUIDs so use the role as the name
-            return "primary" if self.role is DatasetRole.PRIMARY else "reference"
+            return "primary" if self.role is InferencesRole.PRIMARY else "reference"
         return ds_name
 
     @property
-    def role(self) -> DatasetRole:
+    def role(self) -> InferencesRole:
         return self._self_role
 
     @property
@@ -715,14 +706,14 @@ class Dataset(Events):
         return len(self) == 0
 
     @cached_property
-    def primary_key(self) -> pd.Index:
-        return pd.Index(self[PREDICTION_ID])
+    def primary_key(self) -> "pd.Index[Any]":
+        return cast("pd.Index[Any]", pd.Index(self[PREDICTION_ID]))
 
     @overload
     def __getitem__(self, key: ColumnKey) -> "pd.Series[Any]": ...
 
     @overload
-    def __getitem__(self, key: List[RowId]) -> Events: ...
+    def __getitem__(self, key: list[RowId]) -> Events: ...
 
     def __getitem__(self, key: Any) -> Any:
         if isinstance(key, list):
@@ -746,19 +737,19 @@ class Model:
     a column of NaNs.
     """
 
-    _datasets: Dict[DatasetRole, Dataset]
-    _dimensions: Dict[Name, Dimension]
-    _dim_names_by_role: Dict[DimensionRole, List[Name]]
-    _original_columns_by_role: Dict[DatasetRole, pd.Index]
+    _inference_sets: dict[InferencesRole, Inferences]
+    _dimensions: dict[Name, Dimension]
+    _dim_names_by_role: dict[DimensionRole, list[Name]]
+    _original_columns_by_role: dict[InferencesRole, "pd.Index[Any]"]
     _default_timestamps_factory: _ConstantValueSeriesFactory
     _nan_series_factory: _ConstantValueSeriesFactory
-    _dimension_categories_from_all_datasets: _Cache[Name, Tuple[str, ...]]
-    _dimension_min_max_from_all_datasets: _Cache[Name, Tuple[float, float]]
+    _dimension_categories_from_all_inferences: _Cache[Name, tuple[str, ...]]
+    _dimension_min_max_from_all_inferences: _Cache[Name, tuple[float, float]]
 
     def __init__(
         self,
         dimensions: Iterable[Dimension],
-        dataframes: Iterable[Union[pd.DataFrame, Tuple[Name, pd.DataFrame]]],
+        dataframes: Iterable[Union[pd.DataFrame, tuple[Name, pd.DataFrame]]],
         /,
         treat_omitted_columns_as_features: bool = True,
         timestamps_already_normalized: bool = False,
@@ -769,37 +760,37 @@ class Model:
         # memoization
         object.__setattr__(
             self,
-            "_dimension_categories_from_all_datasets",
+            "_dimension_categories_from_all_inferences",
             _Cache[Name, "pd.Series[Any]"](),
         )
         object.__setattr__(
             self,
-            "_dimension_min_max_from_all_datasets",
-            _Cache[Name, Tuple[float, float]](),
+            "_dimension_min_max_from_all_inferences",
+            _Cache[Name, tuple[float, float]](),
         )
 
         df_names, dfs = cast(
-            Tuple[Iterable[Name], Iterable[pd.DataFrame]],
+            tuple[Iterable[Name], Iterable[pd.DataFrame]],
             zip(*_coerce_tuple(dataframes)),
         )
         str_col_dfs = _coerce_str_column_names(dfs)
         padded_dfs = _add_padding(str_col_dfs, pd.DataFrame)
         padded_df_names = _add_padding(df_names, _rand_str)
-        datasets = starmap(
-            self._new_dataset,
-            zip(padded_dfs, padded_df_names, DatasetRole),
+        inference_sets = starmap(
+            self._new_inferences,
+            zip(padded_dfs, padded_df_names, InferencesRole),
         )
-        # Store datasets by role.
+        # Store inferences by role.
         object.__setattr__(
             self,
-            "_datasets",
-            {dataset.role: dataset for dataset in datasets},
+            "_inference_sets",
+            {inferences.role: inferences for inferences in inference_sets},
         )
         # Preserve originals, useful for exporting.
         object.__setattr__(
             self,
             "_original_columns_by_role",
-            {role: dataset.columns for role, dataset in self._datasets.items()},
+            {role: inferences.columns for role, inferences in self._inference_sets.items()},
         )
 
         object.__setattr__(
@@ -828,7 +819,7 @@ class Model:
                 (name, self._new_dimension(name, role=FEATURE))
                 for name in _get_omitted_column_names(
                     self._dimensions.values(),
-                    self._datasets.values(),
+                    self._inference_sets.values(),
                 )
             )
 
@@ -849,7 +840,7 @@ class Model:
                 data_type=(
                     _guess_data_type(
                         dataset.loc[:, dim.name]
-                        for dataset in self._datasets.values()
+                        for dataset in self._inference_sets.values()
                         if dim.name in dataset.columns
                     )
                 ),
@@ -859,9 +850,9 @@ class Model:
         # Add TIMESTAMP if missing.
         # If needed, normalize the timestamps values.
         # If needed, sort the dataframes by time.
-        for dataset_role, dataset in list(self._datasets.items()):
+        for inferences_role, dataset in list(self._inference_sets.items()):
             df = dataset.__wrapped__
-            df_original_columns = self._original_columns_by_role[dataset_role]
+            df_original_columns = self._original_columns_by_role[inferences_role]
 
             # PREDICTION_ID
             dim_pred_id = self._dimensions.get(
@@ -897,20 +888,20 @@ class Model:
             df = df.set_index(dim_time.name, drop=False)
 
             # Update dataset since its dataframe may have changed.
-            self._datasets[dataset_role] = self._new_dataset(
-                df, name=dataset.name, role=dataset_role
+            self._inference_sets[inferences_role] = self._new_inferences(
+                df, name=dataset.name, role=inferences_role
             )
 
     @cached_property
     def is_empty(self) -> bool:
         """Returns True if the model has no data."""
-        return not any(map(len, self._datasets.values()))
+        return not any(map(len, self._inference_sets.values()))
 
     def export_rows_as_parquet_file(
         self,
-        row_numbers: Mapping[DatasetRole, Iterable[int]],
+        row_numbers: Mapping[InferencesRole, Iterable[int]],
         parquet_file: BinaryIO,
-        cluster_ids: Optional[Mapping[DatasetRole, Mapping[int, str]]] = None,
+        cluster_ids: Optional[Mapping[InferencesRole, Mapping[int, str]]] = None,
     ) -> None:
         """
         Given row numbers, exports dataframe subset into parquet file.
@@ -921,29 +912,31 @@ class Model:
 
         Parameters
         ----------
-        row_numbers: Mapping[DatasetRole, Iterable[int]]
+        row_numbers: Mapping[InferencesRole, Iterable[int]]
             mapping of dataset role to list of row numbers
         parquet_file: file handle
             output parquet file handle
-        cluster_ids: Optional[Mapping[DatasetRole, Mapping[int, str]]]
-            mapping of dataset role to mapping of row number to cluster id.
+        cluster_ids: Optional[Mapping[InferencesRole, Mapping[int, str]]]
+            mapping of inferences role to mapping of row number to cluster id.
             If cluster_ids is non-empty, a new column is inserted to the
             dataframe containing the cluster IDs of each row in the exported
             data. The name of the added column name is `__phoenix_cluster_id__`.
         """
         export_dataframes = [pd.DataFrame()]
-        model_has_multiple_datasets = sum(not df.empty for df in self._datasets.values()) > 1
-        for dataset_role, numbers in row_numbers.items():
-            df = self._datasets[dataset_role]
+        model_has_multiple_inference_sets = (
+            sum(not df.empty for df in self._inference_sets.values()) > 1
+        )
+        for inferences_role, numbers in row_numbers.items():
+            df = self._inference_sets[inferences_role]
             columns = [
                 df.columns.get_loc(column_name)
-                for column_name in self._original_columns_by_role[dataset_role]
+                for column_name in self._original_columns_by_role[inferences_role]
             ]
             rows = pd.Series(sorted(set(numbers)))
             filtered_df = df.iloc[rows, columns].reset_index(drop=True)
-            if model_has_multiple_datasets:
+            if model_has_multiple_inference_sets:
                 filtered_df["__phoenix_dataset_name__"] = df.display_name
-            if cluster_ids and (ids := cluster_ids.get(dataset_role)):
+            if cluster_ids and (ids := cluster_ids.get(inferences_role)):
                 filtered_df["__phoenix_cluster_id__"] = rows.apply(ids.get)
             export_dataframes.append(filtered_df)
         pd.concat(export_dataframes).to_parquet(
@@ -954,7 +947,7 @@ class Model:
         )
 
     @cached_property
-    def scalar_dimensions(self) -> Tuple[ScalarDimension, ...]:
+    def scalar_dimensions(self) -> tuple[ScalarDimension, ...]:
         """Put these in a cached sequence because currently their positions
         in this list also determine their node IDs in graphql.
         """
@@ -967,7 +960,7 @@ class Model:
         )
 
     @cached_property
-    def embedding_dimensions(self) -> Tuple[EmbeddingDimension, ...]:
+    def embedding_dimensions(self) -> tuple[EmbeddingDimension, ...]:
         """Put these in a cached sequence because currently their positions
         in this list also determine their node IDs in graphql.
         """
@@ -977,52 +970,52 @@ class Model:
             if not dim.is_dummy and isinstance(dim, EmbeddingDimension)
         )
 
-    def dimension_categories_from_all_datasets(
+    def dimension_categories_from_all_inferences(
         self,
         dimension_name: Name,
-    ) -> Tuple[str, ...]:
+    ) -> tuple[str, ...]:
         dim = self[dimension_name]
         if dim.data_type is CONTINUOUS:
-            return cast(Tuple[str, ...], ())
-        with self._dimension_categories_from_all_datasets() as cache:
+            return cast(tuple[str, ...], ())
+        with self._dimension_categories_from_all_inferences() as cache:
             try:
                 return cache[dimension_name]
             except KeyError:
                 pass
         categories_by_dataset = (
-            pd.Series(dim[role].unique()).dropna().astype(str) for role in DatasetRole
+            pd.Series(dim[role].unique()).dropna().astype(str) for role in InferencesRole
         )
         all_values_combined = chain.from_iterable(categories_by_dataset)
         ans = tuple(np.sort(pd.Series(all_values_combined).unique()))
-        with self._dimension_categories_from_all_datasets() as cache:
+        with self._dimension_categories_from_all_inferences() as cache:
             cache[dimension_name] = ans
         return ans
 
     def dimension_min_max_from_all_df(
         self,
         dimension_name: Name,
-    ) -> Tuple[float, float]:
+    ) -> tuple[float, float]:
         dim = self[dimension_name]
         if dim.data_type is not CONTINUOUS:
             return (np.nan, np.nan)
-        with self._dimension_min_max_from_all_datasets() as cache:
+        with self._dimension_min_max_from_all_inferences() as cache:
             try:
                 return cache[dimension_name]
             except KeyError:
                 pass
-        min_max_by_df = (_agg_min_max(dim[df_role]) for df_role in DatasetRole)
+        min_max_by_df = (_agg_min_max(dim[df_role]) for df_role in InferencesRole)
         all_values_combined = chain.from_iterable(min_max_by_df)
         min_max = _agg_min_max(pd.Series(all_values_combined))
         ans = (min_max.min(), min_max.max())
-        with self._dimension_min_max_from_all_datasets() as cache:
+        with self._dimension_min_max_from_all_inferences() as cache:
             cache[dimension_name] = ans
         return ans
 
     @overload
-    def __getitem__(self, key: Type[Dataset]) -> Iterator[Dataset]: ...
+    def __getitem__(self, key: type[Inferences]) -> Iterator[Inferences]: ...
 
     @overload
-    def __getitem__(self, key: DatasetRole) -> Dataset: ...
+    def __getitem__(self, key: InferencesRole) -> Inferences: ...
 
     @overload
     def __getitem__(self, key: ColumnKey) -> Dimension: ...
@@ -1031,28 +1024,28 @@ class Model:
     def __getitem__(self, key: MultiDimensionKey) -> Iterator[Dimension]: ...
 
     @overload
-    def __getitem__(self, key: Type[ScalarDimension]) -> Iterator[ScalarDimension]: ...
+    def __getitem__(self, key: type[ScalarDimension]) -> Iterator[ScalarDimension]: ...
 
     @overload
-    def __getitem__(self, key: Type[EmbeddingDimension]) -> Iterator[EmbeddingDimension]: ...
+    def __getitem__(self, key: type[EmbeddingDimension]) -> Iterator[EmbeddingDimension]: ...
 
     @overload
-    def __getitem__(self, key: Type[Dimension]) -> Iterator[Dimension]: ...
+    def __getitem__(self, key: type[Dimension]) -> Iterator[Dimension]: ...
 
     @overload
     def __getitem__(
         self,
-        key: Tuple[
+        key: tuple[
             MultiDimensionKey,
-            Union[Type[ScalarDimension], Type[EmbeddingDimension]],
+            Union[type[ScalarDimension], type[EmbeddingDimension]],
         ],
     ) -> Iterator[Dimension]: ...
 
     def __getitem__(self, key: Any) -> Any:
-        if key is Dataset:
-            return self._datasets.values()
-        if isinstance(key, DatasetRole):
-            return self._datasets[key]
+        if key is Inferences:
+            return self._inference_sets.values()
+        if isinstance(key, InferencesRole):
+            return self._inference_sets[key]
         if _is_column_key(key):
             return self._get_dim(key)
         if _is_multi_dimension_key(key):
@@ -1092,9 +1085,9 @@ class Model:
 
     def _get_multi_dims_by_type(
         self,
-        key: Tuple[
+        key: tuple[
             MultiDimensionKey,
-            Union[Type[ScalarDimension], Type[EmbeddingDimension]],
+            Union[type[ScalarDimension], type[EmbeddingDimension]],
         ],
     ) -> Iterator[Dimension]:
         return filter(lambda dim: type(dim) is key[1], self[key[0]])
@@ -1103,7 +1096,7 @@ class Model:
     def _new_dimension(
         self,
         obj: DimensionRole,
-        cls: Type[Dimension] = ScalarDimension,
+        cls: type[Dimension] = ScalarDimension,
         **kwargs: Any,
     ) -> Dimension: ...
 
@@ -1111,7 +1104,7 @@ class Model:
     def _new_dimension(
         self,
         obj: Name,
-        cls: Type[Dimension] = ScalarDimension,
+        cls: type[Dimension] = ScalarDimension,
         **kwargs: Any,
     ) -> Dimension: ...
 
@@ -1123,7 +1116,7 @@ class Model:
     ) -> Dimension: ...
 
     def _new_dimension(
-        self, obj: Any, cls: Type[Dimension] = ScalarDimension, **kwargs: Any
+        self, obj: Any, cls: type[Dimension] = ScalarDimension, **kwargs: Any
     ) -> Dimension:
         """Creates a new Dimension or copies an existing one, setting the
         model weak reference to the `self` Model instance, and sharing the
@@ -1152,17 +1145,17 @@ class Model:
             )
         raise ValueError(f"invalid argument: {repr(obj)}")
 
-    def _new_dataset(
+    def _new_inferences(
         self,
         df: pd.DataFrame,
         /,
         name: str,
-        role: DatasetRole,
-    ) -> Dataset:
-        """Creates a new Dataset, setting the model weak reference to the
+        role: InferencesRole,
+    ) -> Inferences:
+        """Creates a new Inferences, setting the model weak reference to the
         `self` Model instance.
         """
-        return Dataset(df, name=name, role=role, _model=proxy(self))
+        return Inferences(df, name=name, role=role, _model=proxy(self))
 
 
 @dataclass(frozen=True)
@@ -1173,13 +1166,13 @@ class Schema(SchemaSpec):
     prediction_score: Optional[str] = None
     actual_label: Optional[str] = None
     actual_score: Optional[str] = None
-    prompt: Optional[RetrievalEmbedding] = None
+    prompt: Optional[Embedding] = None
     response: Optional[Union[str, Embedding]] = None
     features: Iterable[Union[str, CompositeDimensionSpec]] = field(default_factory=list)
     tags: Iterable[Union[str, CompositeDimensionSpec]] = field(default_factory=list)
 
     # internal attribute not exposed to users
-    _dimensions: List[Dimension] = field(
+    _dimensions: list[Dimension] = field(
         init=False, repr=False, hash=False, compare=False, default_factory=list
     )
 
@@ -1246,7 +1239,7 @@ class Schema(SchemaSpec):
 
     def __call__(
         self,
-        *dataframes: Union[pd.DataFrame, Tuple[Name, pd.DataFrame]],
+        *dataframes: Union[pd.DataFrame, tuple[Name, pd.DataFrame]],
         **kwargs: Any,
     ) -> Model:
         """Dimensions are the "baton" that Schema hands over to Model."""
@@ -1300,7 +1293,7 @@ def _get_omitted_column_names(
 
 def _group_names_by_dim_role(
     dimensions: Iterable[Dimension],
-) -> Iterator[Tuple[DimensionRole, List[str]]]:
+) -> Iterator[tuple[DimensionRole, list[str]]]:
     return (
         (role, [dim.name for dim in dims])
         for role, dims in groupby(
@@ -1344,7 +1337,7 @@ def _series_uuid(length: int) -> "pd.Series[str]":
 
 
 def _raise_if_too_many_dataframes(given: int) -> None:
-    limit = len(DatasetRole)
+    limit = len(InferencesRole)
     if not 0 < given <= limit:
         raise ValueError(f"expected between 1 to {limit} dataframes, but {given} were given")
 
@@ -1354,8 +1347,8 @@ def _coerce_str(obj: Optional[str]) -> str:
 
 
 def _coerce_tuple(
-    dataframes: Iterable[Union[pd.DataFrame, Tuple[Name, pd.DataFrame]]],
-) -> Iterator[Tuple[Name, pd.DataFrame]]:
+    dataframes: Iterable[Union[pd.DataFrame, tuple[Name, pd.DataFrame]]],
+) -> Iterator[tuple[Name, pd.DataFrame]]:
     for dataframe in dataframes:
         if isinstance(dataframe, pd.DataFrame):
             yield (_rand_str(), dataframe)