arize-phoenix 4.4.4rc5__py3-none-any.whl → 4.5.0__py3-none-any.whl

phoenix/core/model_schema.py

@@ -48,7 +48,7 @@ from pandas.core.dtypes.common import (
 from typing_extensions import TypeAlias, TypeGuard
 from wrapt import ObjectProxy
 
-from phoenix.config import GENERATED_INFERENCES_NAME_PREFIX
+from phoenix.config import GENERATED_DATASET_NAME_PREFIX
 from phoenix.datetime_utils import floor_to_minute
 
 
@@ -185,7 +185,7 @@ class RetrievalEmbedding(Embedding):
                 yield value
 
 
-class InferencesRole(Enum):
+class DatasetRole(Enum):
     """A dataframe's role in a Model: primary or reference (as
     baseline for drift).
     """
@@ -194,8 +194,8 @@ class InferencesRole(Enum):
     REFERENCE = auto()
 
 
-PRIMARY = InferencesRole.PRIMARY
-REFERENCE = InferencesRole.REFERENCE
+PRIMARY = DatasetRole.PRIMARY
+REFERENCE = DatasetRole.REFERENCE
 
 
 @dataclass(frozen=True, repr=False, eq=False)
@@ -381,7 +381,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: InferencesRole) -> "pd.Series[Any]":
+    def __getitem__(self, df_role: DatasetRole) -> "pd.Series[Any]":
         if self._model is None:
             return pd.Series(dtype=object)
         model = cast(Model, self._model)
@@ -416,7 +416,7 @@ class ScalarDimension(Dimension):
         if self._model is None or self.data_type is CONTINUOUS:
             return ()
         model = cast(Model, self._model)
-        return model.dimension_categories_from_all_inferences(self.name)
+        return model.dimension_categories_from_all_datasets(self.name)
 
 
 @dataclass(frozen=True)
@@ -582,7 +582,7 @@ class EventId(NamedTuple):
     """Identifies an event."""
 
     row_id: int = 0
-    inferences_id: InferencesRole = PRIMARY
+    dataset_id: DatasetRole = PRIMARY
 
     def __str__(self) -> str:
         return ":".join(map(str, self))
@@ -625,7 +625,7 @@ class Events(ModelData):
         self,
         df: pd.DataFrame,
         /,
-        role: InferencesRole,
+        role: DatasetRole,
         **kwargs: Any,
     ) -> None:
         super().__init__(df, **kwargs)
@@ -676,7 +676,7 @@ class Events(ModelData):
         return super().__getitem__(key)
 
 
-class Inferences(Events):
+class Dataset(Events):
     """pd.DataFrame wrapped with extra functions and metadata."""
 
     def __init__(
@@ -701,13 +701,13 @@ class Inferences(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_INFERENCES_NAME_PREFIX):
+        if ds_name.startswith(GENERATED_DATASET_NAME_PREFIX):
             # The generated names are UUIDs so use the role as the name
-            return "primary" if self.role is InferencesRole.PRIMARY else "reference"
+            return "primary" if self.role is DatasetRole.PRIMARY else "reference"
         return ds_name
 
     @property
-    def role(self) -> InferencesRole:
+    def role(self) -> DatasetRole:
         return self._self_role
 
     @property
@@ -746,14 +746,14 @@ class Model:
     a column of NaNs.
     """
 
-    _inference_sets: Dict[InferencesRole, Inferences]
+    _datasets: Dict[DatasetRole, Dataset]
     _dimensions: Dict[Name, Dimension]
     _dim_names_by_role: Dict[DimensionRole, List[Name]]
-    _original_columns_by_role: Dict[InferencesRole, "pd.Index[Any]"]
+    _original_columns_by_role: Dict[DatasetRole, "pd.Index[Any]"]
     _default_timestamps_factory: _ConstantValueSeriesFactory
     _nan_series_factory: _ConstantValueSeriesFactory
-    _dimension_categories_from_all_inferences: _Cache[Name, Tuple[str, ...]]
-    _dimension_min_max_from_all_inferences: _Cache[Name, Tuple[float, float]]
+    _dimension_categories_from_all_datasets: _Cache[Name, Tuple[str, ...]]
+    _dimension_min_max_from_all_datasets: _Cache[Name, Tuple[float, float]]
 
     def __init__(
         self,
@@ -769,12 +769,12 @@ class Model:
         # memoization
         object.__setattr__(
             self,
-            "_dimension_categories_from_all_inferences",
+            "_dimension_categories_from_all_datasets",
             _Cache[Name, "pd.Series[Any]"](),
         )
         object.__setattr__(
             self,
-            "_dimension_min_max_from_all_inferences",
+            "_dimension_min_max_from_all_datasets",
             _Cache[Name, Tuple[float, float]](),
         )
 
@@ -785,21 +785,21 @@ class Model:
         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)
-        inference_sets = starmap(
-            self._new_inferences,
-            zip(padded_dfs, padded_df_names, InferencesRole),
+        datasets = starmap(
+            self._new_dataset,
+            zip(padded_dfs, padded_df_names, DatasetRole),
         )
-        # Store inferences by role.
+        # Store datasets by role.
         object.__setattr__(
             self,
-            "_inference_sets",
-            {inferences.role: inferences for inferences in inference_sets},
+            "_datasets",
+            {dataset.role: dataset for dataset in datasets},
         )
         # Preserve originals, useful for exporting.
         object.__setattr__(
             self,
             "_original_columns_by_role",
-            {role: inferences.columns for role, inferences in self._inference_sets.items()},
+            {role: dataset.columns for role, dataset in self._datasets.items()},
         )
 
         object.__setattr__(
@@ -828,7 +828,7 @@ class Model:
                 (name, self._new_dimension(name, role=FEATURE))
                 for name in _get_omitted_column_names(
                     self._dimensions.values(),
-                    self._inference_sets.values(),
+                    self._datasets.values(),
                 )
             )
 
@@ -849,7 +849,7 @@ class Model:
                 data_type=(
                     _guess_data_type(
                         dataset.loc[:, dim.name]
-                        for dataset in self._inference_sets.values()
+                        for dataset in self._datasets.values()
                         if dim.name in dataset.columns
                     )
                 ),
@@ -859,9 +859,9 @@ class Model:
         # Add TIMESTAMP if missing.
         # If needed, normalize the timestamps values.
         # If needed, sort the dataframes by time.
-        for inferences_role, dataset in list(self._inference_sets.items()):
+        for dataset_role, dataset in list(self._datasets.items()):
             df = dataset.__wrapped__
-            df_original_columns = self._original_columns_by_role[inferences_role]
+            df_original_columns = self._original_columns_by_role[dataset_role]
 
             # PREDICTION_ID
             dim_pred_id = self._dimensions.get(
@@ -897,20 +897,20 @@ class Model:
             df = df.set_index(dim_time.name, drop=False)
 
             # Update dataset since its dataframe may have changed.
-            self._inference_sets[inferences_role] = self._new_inferences(
-                df, name=dataset.name, role=inferences_role
+            self._datasets[dataset_role] = self._new_dataset(
+                df, name=dataset.name, role=dataset_role
             )
 
     @cached_property
     def is_empty(self) -> bool:
         """Returns True if the model has no data."""
-        return not any(map(len, self._inference_sets.values()))
+        return not any(map(len, self._datasets.values()))
 
     def export_rows_as_parquet_file(
         self,
-        row_numbers: Mapping[InferencesRole, Iterable[int]],
+        row_numbers: Mapping[DatasetRole, Iterable[int]],
         parquet_file: BinaryIO,
-        cluster_ids: Optional[Mapping[InferencesRole, Mapping[int, str]]] = None,
+        cluster_ids: Optional[Mapping[DatasetRole, Mapping[int, str]]] = None,
     ) -> None:
         """
         Given row numbers, exports dataframe subset into parquet file.
@@ -921,31 +921,29 @@ class Model:
 
         Parameters
         ----------
-        row_numbers: Mapping[InferencesRole, Iterable[int]]
+        row_numbers: Mapping[DatasetRole, Iterable[int]]
             mapping of dataset role to list of row numbers
         parquet_file: file handle
             output parquet file handle
-        cluster_ids: Optional[Mapping[InferencesRole, Mapping[int, str]]]
-            mapping of inferences role to mapping of row number to cluster id.
+        cluster_ids: Optional[Mapping[DatasetRole, Mapping[int, str]]]
+            mapping of dataset 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_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]
+        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]
             columns = [
                 df.columns.get_loc(column_name)
-                for column_name in self._original_columns_by_role[inferences_role]
+                for column_name in self._original_columns_by_role[dataset_role]
             ]
             rows = pd.Series(sorted(set(numbers)))
             filtered_df = df.iloc[rows, columns].reset_index(drop=True)
-            if model_has_multiple_inference_sets:
+            if model_has_multiple_datasets:
                 filtered_df["__phoenix_dataset_name__"] = df.display_name
-            if cluster_ids and (ids := cluster_ids.get(inferences_role)):
+            if cluster_ids and (ids := cluster_ids.get(dataset_role)):
                 filtered_df["__phoenix_cluster_id__"] = rows.apply(ids.get)
             export_dataframes.append(filtered_df)
         pd.concat(export_dataframes).to_parquet(
@@ -979,24 +977,24 @@ class Model:
             if not dim.is_dummy and isinstance(dim, EmbeddingDimension)
         )
 
-    def dimension_categories_from_all_inferences(
+    def dimension_categories_from_all_datasets(
         self,
         dimension_name: Name,
     ) -> Tuple[str, ...]:
         dim = self[dimension_name]
         if dim.data_type is CONTINUOUS:
             return cast(Tuple[str, ...], ())
-        with self._dimension_categories_from_all_inferences() as cache:
+        with self._dimension_categories_from_all_datasets() as cache:
             try:
                 return cache[dimension_name]
             except KeyError:
                 pass
         categories_by_dataset = (
-            pd.Series(dim[role].unique()).dropna().astype(str) for role in InferencesRole
+            pd.Series(dim[role].unique()).dropna().astype(str) for role in DatasetRole
         )
         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_inferences() as cache:
+        with self._dimension_categories_from_all_datasets() as cache:
             cache[dimension_name] = ans
         return ans
 
@@ -1007,24 +1005,24 @@ class Model:
         dim = self[dimension_name]
         if dim.data_type is not CONTINUOUS:
             return (np.nan, np.nan)
-        with self._dimension_min_max_from_all_inferences() as cache:
+        with self._dimension_min_max_from_all_datasets() as cache:
             try:
                 return cache[dimension_name]
             except KeyError:
                 pass
-        min_max_by_df = (_agg_min_max(dim[df_role]) for df_role in InferencesRole)
+        min_max_by_df = (_agg_min_max(dim[df_role]) for df_role in DatasetRole)
         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_inferences() as cache:
+        with self._dimension_min_max_from_all_datasets() as cache:
             cache[dimension_name] = ans
         return ans
 
     @overload
-    def __getitem__(self, key: Type[Inferences]) -> Iterator[Inferences]: ...
+    def __getitem__(self, key: Type[Dataset]) -> Iterator[Dataset]: ...
 
     @overload
-    def __getitem__(self, key: InferencesRole) -> Inferences: ...
+    def __getitem__(self, key: DatasetRole) -> Dataset: ...
 
     @overload
     def __getitem__(self, key: ColumnKey) -> Dimension: ...
@@ -1051,10 +1049,10 @@ class Model:
     ) -> Iterator[Dimension]: ...
 
     def __getitem__(self, key: Any) -> Any:
-        if key is Inferences:
-            return self._inference_sets.values()
-        if isinstance(key, InferencesRole):
-            return self._inference_sets[key]
+        if key is Dataset:
+            return self._datasets.values()
+        if isinstance(key, DatasetRole):
+            return self._datasets[key]
         if _is_column_key(key):
             return self._get_dim(key)
         if _is_multi_dimension_key(key):
@@ -1154,17 +1152,17 @@ class Model:
             )
         raise ValueError(f"invalid argument: {repr(obj)}")
 
-    def _new_inferences(
+    def _new_dataset(
         self,
         df: pd.DataFrame,
         /,
         name: str,
-        role: InferencesRole,
-    ) -> Inferences:
-        """Creates a new Inferences, setting the model weak reference to the
+        role: DatasetRole,
+    ) -> Dataset:
+        """Creates a new Dataset, setting the model weak reference to the
         `self` Model instance.
         """
-        return Inferences(df, name=name, role=role, _model=proxy(self))
+        return Dataset(df, name=name, role=role, _model=proxy(self))
 
 
 @dataclass(frozen=True)
@@ -1346,7 +1344,7 @@ def _series_uuid(length: int) -> "pd.Series[str]":
 
 
 def _raise_if_too_many_dataframes(given: int) -> None:
-    limit = len(InferencesRole)
+    limit = len(DatasetRole)
     if not 0 < given <= limit:
         raise ValueError(f"expected between 1 to {limit} dataframes, but {given} were given")
 

phoenix/core/model_schema_adapter.py

@@ -10,21 +10,21 @@ from phoenix import EmbeddingColumnNames, Inferences
 from phoenix.core.model import _get_embedding_dimensions
 from phoenix.core.model_schema import Embedding, Model, RetrievalEmbedding, Schema
 from phoenix.inferences.schema import RetrievalEmbeddingColumnNames
-from phoenix.inferences.schema import Schema as InferencesSchema
+from phoenix.inferences.schema import Schema as DatasetSchema
 
-InferencesName: TypeAlias = str
+DatasetName: TypeAlias = str
 ColumnName: TypeAlias = str
 DisplayName: TypeAlias = str
 
 
-def create_model_from_inferences(*inference_sets: Optional[Inferences]) -> Model:
+def create_model_from_datasets(*datasets: Optional[Inferences]) -> Model:
     # TODO: move this validation into model_schema.Model.
-    if len(inference_sets) > 1 and inference_sets[0] is not None:
+    if len(datasets) > 1 and datasets[0] is not None:
         # Check that for each embedding dimension all vectors
-        # have the same length between inferences.
-        _ = _get_embedding_dimensions(inference_sets[0], inference_sets[1])
+        # have the same length between datasets.
+        _ = _get_embedding_dimensions(datasets[0], datasets[1])
 
-    named_dataframes: List[Tuple[InferencesName, pd.DataFrame]] = []
+    named_dataframes: List[Tuple[DatasetName, pd.DataFrame]] = []
     prediction_ids: List[ColumnName] = []
     timestamps: List[ColumnName] = []
     prediction_labels: List[ColumnName] = []
@@ -37,35 +37,33 @@ def create_model_from_inferences(*inference_sets: Optional[Inferences]) -> Model
     prompts: List[EmbeddingColumnNames] = []
     responses: List[Union[str, EmbeddingColumnNames]] = []
 
-    for inferences in filter(_is_inferences, inference_sets):
-        df = inferences.dataframe
+    for dataset in filter(_is_dataset, datasets):
+        df = dataset.dataframe
         # Coerce string column names at run time.
         df = df.set_axis(
             map(str, df.columns),
             axis=1,
         )
-        named_dataframes.append((inferences.name, df))
-        inferences_schema = (
-            inferences.schema if inferences.schema is not None else InferencesSchema()
-        )
+        named_dataframes.append((dataset.name, df))
+        dataset_schema = dataset.schema if dataset.schema is not None else DatasetSchema()
         for display_name, embedding in (
-            inferences_schema.embedding_feature_column_names or {}
+            dataset_schema.embedding_feature_column_names or {}
         ).items():
             if display_name not in embeddings:
                 embeddings[display_name] = embedding
-        if inferences_schema.prompt_column_names is not None:
-            prompts.append(inferences_schema.prompt_column_names)
-        if inferences_schema.response_column_names is not None:
-            responses.append(inferences_schema.response_column_names)
+        if dataset_schema.prompt_column_names is not None:
+            prompts.append(dataset_schema.prompt_column_names)
+        if dataset_schema.response_column_names is not None:
+            responses.append(dataset_schema.response_column_names)
         for source, sink in (
-            ([inferences_schema.prediction_id_column_name], prediction_ids),
-            ([inferences_schema.timestamp_column_name], timestamps),
-            ([inferences_schema.prediction_label_column_name], prediction_labels),
-            ([inferences_schema.prediction_score_column_name], prediction_scores),
-            ([inferences_schema.actual_label_column_name], actual_labels),
-            ([inferences_schema.actual_score_column_name], actual_scores),
-            (inferences_schema.feature_column_names or (), features),
-            (inferences_schema.tag_column_names or (), tags),
+            ([dataset_schema.prediction_id_column_name], prediction_ids),
+            ([dataset_schema.timestamp_column_name], timestamps),
+            ([dataset_schema.prediction_label_column_name], prediction_labels),
+            ([dataset_schema.prediction_score_column_name], prediction_scores),
+            ([dataset_schema.actual_label_column_name], actual_labels),
+            ([dataset_schema.actual_score_column_name], actual_scores),
+            (dataset_schema.feature_column_names or (), features),
+            (dataset_schema.tag_column_names or (), tags),
         ):
             # Coerce None to "" to simplify type checks.
             sink.extend(map(lambda s: "" if s is None else str(s), source))
@@ -134,7 +132,7 @@ def create_model_from_inferences(*inference_sets: Optional[Inferences]) -> Model
     )
 
 
-def _is_inferences(obj: Optional[Inferences]) -> TypeGuard[Inferences]:
+def _is_dataset(obj: Optional[Inferences]) -> TypeGuard[Inferences]:
     return type(obj) is Inferences
 
 

phoenix/db/bulk_inserter.py

@@ -1,6 +1,5 @@
 import asyncio
 import logging
-from asyncio import Queue
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from itertools import islice
@@ -15,7 +14,6 @@ from typing import (
     Optional,
     Set,
     Tuple,
-    cast,
 )
 
 from cachetools import LRUCache
@@ -24,11 +22,10 @@ from typing_extensions import TypeAlias
 
 import phoenix.trace.v1 as pb
 from phoenix.db.insertion.evaluation import (
-    EvaluationInsertionEvent,
+    EvaluationInsertionResult,
     InsertEvaluationError,
     insert_evaluation,
 )
-from phoenix.db.insertion.helpers import DataManipulation, DataManipulationEvent
 from phoenix.db.insertion.span import SpanInsertionEvent, insert_span
 from phoenix.server.api.dataloaders import CacheForDataLoaders
 from phoenix.trace.schemas import Span
@@ -49,29 +46,23 @@ class BulkInserter:
         db: Callable[[], AsyncContextManager[AsyncSession]],
         *,
         cache_for_dataloaders: Optional[CacheForDataLoaders] = None,
-        initial_batch_of_operations: Iterable[DataManipulation] = (),
         initial_batch_of_spans: Optional[Iterable[Tuple[Span, str]]] = None,
         initial_batch_of_evaluations: Optional[Iterable[pb.Evaluation]] = None,
         sleep: float = 0.1,
-        max_ops_per_transaction: int = 1000,
-        max_queue_size: int = 1000,
+        max_num_per_transaction: int = 1000,
         enable_prometheus: bool = False,
     ) -> None:
         """
         :param db: A function to initiate a new database session.
         :param initial_batch_of_spans: Initial batch of spans to insert.
         :param sleep: The time to sleep between bulk insertions
-        :param max_ops_per_transaction: The maximum number of operations to dequeue from
-        the operations queue for each transaction.
-        :param max_queue_size: The maximum length of the operations queue.
-        :param enable_prometheus: Whether Prometheus is enabled.
+        :param max_num_per_transaction: The maximum number of items to insert in a single
+        transaction. Multiple transactions will be used if there are more items in the batch.
         """
         self._db = db
         self._running = False
         self._sleep = sleep
-        self._max_ops_per_transaction = max_ops_per_transaction
-        self._operations: Optional[Queue[DataManipulation]] = None
-        self._max_queue_size = max_queue_size
+        self._max_num_per_transaction = max_num_per_transaction
         self._spans: List[Tuple[Span, str]] = (
             [] if initial_batch_of_spans is None else list(initial_batch_of_spans)
         )
@@ -90,58 +81,27 @@ class BulkInserter:
 
     async def __aenter__(
         self,
-    ) -> Tuple[
-        Callable[[Span, str], Awaitable[None]],
-        Callable[[pb.Evaluation], Awaitable[None]],
-        Callable[[DataManipulation], None],
-    ]:
+    ) -> Tuple[Callable[[Span, str], Awaitable[None]], Callable[[pb.Evaluation], Awaitable[None]]]:
         self._running = True
-        self._operations = Queue(maxsize=self._max_queue_size)
         self._task = asyncio.create_task(self._bulk_insert())
-        return (
-            self._queue_span,
-            self._queue_evaluation,
-            self._enqueue_operation,
-        )
+        return self._queue_span, self._queue_evaluation
 
     async def __aexit__(self, *args: Any) -> None:
-        self._operations = None
         self._running = False
 
-    def _enqueue_operation(self, operation: DataManipulation) -> None:
-        cast("Queue[DataManipulation]", self._operations).put_nowait(operation)
-
     async def _queue_span(self, span: Span, project_name: str) -> None:
         self._spans.append((span, project_name))
 
     async def _queue_evaluation(self, evaluation: pb.Evaluation) -> None:
         self._evaluations.append(evaluation)
 
-    async def _process_events(self, events: Iterable[Optional[DataManipulationEvent]]) -> None: ...
-
     async def _bulk_insert(self) -> None:
-        assert isinstance(self._operations, Queue)
         spans_buffer, evaluations_buffer = None, None
         # start first insert immediately if the inserter has not run recently
-        while self._running or not self._operations.empty() or self._spans or self._evaluations:
-            if self._operations.empty() and not (self._spans or self._evaluations):
+        while self._spans or self._evaluations or self._running:
+            if not (self._spans or self._evaluations):
                 await asyncio.sleep(self._sleep)
                 continue
-            ops_remaining, events = self._max_ops_per_transaction, []
-            async with self._db() as session:
-                while ops_remaining and not self._operations.empty():
-                    ops_remaining -= 1
-                    op = await self._operations.get()
-                    try:
-                        async with session.begin_nested():
-                            events.append(await op(session))
-                    except Exception as e:
-                        if self._enable_prometheus:
-                            from phoenix.server.prometheus import BULK_LOADER_EXCEPTIONS
-
-                            BULK_LOADER_EXCEPTIONS.inc()
-                        logger.exception(str(e))
-            await self._process_events(events)
             # It's important to grab the buffers at the same time so there's
             # no race condition, since an eval insertion will fail if the span
             # it references doesn't exist. Grabbing the eval buffer later may
@@ -170,11 +130,11 @@ class BulkInserter:
 
     async def _insert_spans(self, spans: List[Tuple[Span, str]]) -> TransactionResult:
         transaction_result = TransactionResult()
-        for i in range(0, len(spans), self._max_ops_per_transaction):
+        for i in range(0, len(spans), self._max_num_per_transaction):
             try:
                 start = perf_counter()
                 async with self._db() as session:
-                    for span, project_name in islice(spans, i, i + self._max_ops_per_transaction):
+                    for span, project_name in islice(spans, i, i + self._max_num_per_transaction):
                         if self._enable_prometheus:
                             from phoenix.server.prometheus import BULK_LOADER_SPAN_INSERTIONS
 
@@ -209,16 +169,16 @@ class BulkInserter:
 
     async def _insert_evaluations(self, evaluations: List[pb.Evaluation]) -> TransactionResult:
         transaction_result = TransactionResult()
-        for i in range(0, len(evaluations), self._max_ops_per_transaction):
+        for i in range(0, len(evaluations), self._max_num_per_transaction):
             try:
                 start = perf_counter()
                 async with self._db() as session:
-                    for evaluation in islice(evaluations, i, i + self._max_ops_per_transaction):
+                    for evaluation in islice(evaluations, i, i + self._max_num_per_transaction):
                         if self._enable_prometheus:
                             from phoenix.server.prometheus import BULK_LOADER_EVALUATION_INSERTIONS
 
                             BULK_LOADER_EVALUATION_INSERTIONS.inc()
-                        result: Optional[EvaluationInsertionEvent] = None
+                        result: Optional[EvaluationInsertionResult] = None
                         try:
                             async with session.begin_nested():
                                 result = await insert_evaluation(session, evaluation)

phoenix/db/insertion/evaluation.py

@@ -15,24 +15,24 @@ class InsertEvaluationError(PhoenixException):
     pass
 
 
-class EvaluationInsertionEvent(NamedTuple):
+class EvaluationInsertionResult(NamedTuple):
     project_rowid: int
     evaluation_name: str
 
 
-class SpanEvaluationInsertionEvent(EvaluationInsertionEvent): ...
+class SpanEvaluationInsertionEvent(EvaluationInsertionResult): ...
 
 
-class TraceEvaluationInsertionEvent(EvaluationInsertionEvent): ...
+class TraceEvaluationInsertionEvent(EvaluationInsertionResult): ...
 
 
-class DocumentEvaluationInsertionEvent(EvaluationInsertionEvent): ...
+class DocumentEvaluationInsertionEvent(EvaluationInsertionResult): ...
 
 
 async def insert_evaluation(
     session: AsyncSession,
     evaluation: pb.Evaluation,
-) -> Optional[EvaluationInsertionEvent]:
+) -> Optional[EvaluationInsertionResult]:
     evaluation_name = evaluation.name
     result = evaluation.result
     label = result.label.value if result.HasField("label") else None
@@ -160,7 +160,7 @@ async def _insert_document_evaluation(
     label: Optional[str],
     score: Optional[float],
     explanation: Optional[str],
-) -> EvaluationInsertionEvent:
+) -> EvaluationInsertionResult:
     dialect = SupportedSQLDialect(session.bind.dialect.name)
     stmt = (
         select(

phoenix/db/insertion/helpers.py

@@ -1,25 +1,14 @@
-from abc import ABC
 from enum import Enum, auto
-from typing import Any, Awaitable, Callable, Mapping, Optional, Sequence
+from typing import Any, Mapping, Optional, Sequence
 
 from sqlalchemy import Insert, insert
 from sqlalchemy.dialects.postgresql import insert as insert_postgresql
 from sqlalchemy.dialects.sqlite import insert as insert_sqlite
-from sqlalchemy.ext.asyncio import AsyncSession
-from typing_extensions import TypeAlias, assert_never
+from typing_extensions import assert_never
 
 from phoenix.db.helpers import SupportedSQLDialect
 
 
-class DataManipulationEvent(ABC):
-    """
-    Execution of DML (Data Manipulation Language) statements.
-    """
-
-
-DataManipulation: TypeAlias = Callable[[AsyncSession], Awaitable[Optional[DataManipulationEvent]]]
-
-
 class OnConflict(Enum):
     DO_NOTHING = auto()
     DO_UPDATE = auto()