arize-phoenix 0.0.2rc3__py3-none-any.whl → 0.0.2rc5__py3-none-any.whl

phoenix/server/api/input_types/Granularity.py

@@ -40,6 +40,8 @@ class Granularity:
 def to_timestamps(
     time_range: TimeRange, granularity: Granularity
 ) -> Generator[datetime, None, None]:
+    if not granularity.sampling_interval_minutes:
+        return
     yield from (
         takewhile(
             lambda t: time_range.start < t,  # type: ignore

phoenix/server/api/interceptor.py

@@ -0,0 +1,28 @@
+import math
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class Interceptor(ABC):
+    """an abstract class making use of the descriptor protocol
+    see https://docs.python.org/3/howto/descriptor.html"""
+
+    private_name: str
+
+    def __set_name__(self, owner: Any, name: str) -> None:
+        self.private_name = "_" + name
+
+    def __get__(self, instance: Any, owner: Any) -> Any:
+        return self if instance is None else getattr(instance, self.private_name)
+
+    @abstractmethod
+    def __set__(self, instance: Any, value: Any) -> None:
+        ...
+
+
+class NoneIfNan(Interceptor):
+    """descriptor that converts NaN and Inf to None because NaN can't be
+    serialized to JSON by the graphql object"""
+
+    def __set__(self, instance: Any, value: float) -> None:
+        setattr(instance, self.private_name, value if math.isfinite(value) else None)

phoenix/server/api/types/Dimension.py

@@ -1,23 +1,18 @@
-import math
-from datetime import timedelta
 from typing import List, Optional
 
 import strawberry
 from strawberry.types import Info
 
 from phoenix.core import Dimension as CoreDimension
-from phoenix.metrics.mixins import UnaryOperator
-from phoenix.metrics.timeseries import timeseries
 from phoenix.server.api.context import Context
 
-from ..input_types.Granularity import Granularity, to_timestamps
+from ..input_types.Granularity import Granularity
 from ..input_types.TimeRange import TimeRange
-from . import METRICS
 from .DataQualityMetric import DataQualityMetric
-from .DataQualityTimeSeries import DataQualityTimeSeries, to_gql_timeseries
 from .DimensionDataType import DimensionDataType
 from .DimensionType import DimensionType
 from .node import Node
+from .TimeSeries import DataQualityTimeSeries
 
 
 @strawberry.type
@@ -32,15 +27,22 @@ class Dimension(Node):
     )
 
     @strawberry.field
-    async def dataQualityMetric(
-        self, metric: DataQualityMetric, info: Info[Context, None]
+    async def data_quality_metric(
+        self,
+        info: Info[Context, None],
+        metric: DataQualityMetric,
+        time_range: Optional[TimeRange] = None,
     ) -> Optional[float]:
-        metric_cls = METRICS.get(metric.value, None)
-        if not metric_cls or not issubclass(metric_cls, UnaryOperator):
-            raise NotImplementedError(f"Metric {metric} is not implemented.")
-        df = info.context.model.primary_dataset.dataframe
-        _, ans = metric_cls(self.name)(df)
-        return None if math.isnan(ans) else ans
+        if len(
+            data := DataQualityTimeSeries(
+                self.name,
+                info.context.model,
+                metric,
+                time_range,
+            ).data
+        ):
+            return data.pop().value
+        return None
 
     @strawberry.field(
         description=(
@@ -70,24 +72,12 @@ class Dimension(Node):
         time_range: TimeRange,
         granularity: Granularity,
     ) -> DataQualityTimeSeries:
-        dimension_name = self.name
-        metric_cls = METRICS.get(metric.value, None)
-        if not metric_cls or not issubclass(metric_cls, UnaryOperator):
-            raise NotImplementedError(f"Metric {metric} is not implemented.")
-        dataset = info.context.model.primary_dataset
-        metric_instance = metric_cls(dimension_name)
-        return dataset.dataframe.pipe(
-            timeseries(
-                start_time=time_range.start,
-                end_time=time_range.end,
-                evaluation_window=timedelta(minutes=granularity.evaluation_window_minutes),
-                sampling_interval=timedelta(minutes=granularity.sampling_interval_minutes),
-            ),
-            metrics=(metric_instance,),
-        ).pipe(
-            to_gql_timeseries,
-            metric=metric_instance,
-            timestamps=to_timestamps(time_range, granularity),
+        return DataQualityTimeSeries(
+            self.name,
+            info.context.model,
+            metric,
+            time_range,
+            granularity,
         )
 
 

phoenix/server/api/types/EmbeddingDimension.py

@@ -1,7 +1,7 @@
 from collections import defaultdict
 from datetime import datetime, timedelta
 from itertools import chain
-from typing import Any, List, Mapping, Optional
+from typing import Any, Optional
 
 import numpy as np
 import numpy.typing as npt
@@ -16,26 +16,21 @@ from phoenix.datasets import Dataset
 from phoenix.datasets.dataset import DatasetType
 from phoenix.datasets.errors import SchemaError
 from phoenix.datasets.event import EventId
-from phoenix.metrics.embeddings import euclidean_distance
-from phoenix.metrics.mixins import UnaryOperator
-from phoenix.metrics.timeseries import row_interval_from_sorted_time_index, timeseries
+from phoenix.metrics.timeseries import row_interval_from_sorted_time_index
 from phoenix.pointcloud.clustering import Hdbscan
 from phoenix.pointcloud.pointcloud import PointCloud
 from phoenix.pointcloud.projectors import Umap
 from phoenix.server.api.context import Context
 from phoenix.server.api.input_types.TimeRange import TimeRange
 
-from ..input_types.Granularity import Granularity, to_timestamps
-from . import METRICS
+from ..input_types.Granularity import Granularity
 from .DataQualityMetric import DataQualityMetric
-from .DataQualityTimeSeries import DataQualityTimeSeries, to_gql_timeseries
 from .DriftMetric import DriftMetric
-from .DriftTimeSeries import DriftTimeSeries
 from .EmbeddingMetadata import EmbeddingMetadata
 from .EventMetadata import EventMetadata
 from .node import Node
-from .TimeSeries import TimeSeriesDataPoint
-from .UMAPPoints import Cluster, UMAPPoint, UMAPPoints, to_gql_coordinates
+from .TimeSeries import DataQualityTimeSeries, DriftTimeSeries
+from .UMAPPoints import UMAPPoint, UMAPPoints, to_gql_clusters, to_gql_coordinates
 
 # Default UMAP hyperparameters
 DEFAULT_N_COMPONENTS = 3
@@ -43,17 +38,6 @@ DEFAULT_MIN_DIST = 0
 DEFAULT_N_NEIGHBORS = 30
 DEFAULT_N_SAMPLES = 500
 
-
-def to_gql_clusters(clusters: Mapping[EventId, int]) -> List[Cluster]:
-    clusteredEvents = defaultdict(list)
-    for event_id, cluster_id in clusters.items():
-        clusteredEvents[ID(str(cluster_id))].append(ID(str(event_id)))
-    return [
-        Cluster(id=cluster_id, point_ids=event_ids)
-        for cluster_id, event_ids in clusteredEvents.items()
-    ]
-
-
 DRIFT_EVAL_WINDOW_NUM_INTERVALS = 72
 EVAL_INTERVAL_LENGTH = timedelta(hours=1)
 
@@ -66,7 +50,10 @@ class EmbeddingDimension(Node):
 
     @strawberry.field
     def drift_metric(
-        self, metric: DriftMetric, time_range: TimeRange, info: Info[Context, None]
+        self,
+        info: Info[Context, None],
+        metric: DriftMetric,
+        time_range: Optional[TimeRange] = None,
     ) -> Optional[float]:
         """
         Computes a drift metric between all reference data and the primary data
@@ -75,29 +62,16 @@ class EmbeddingDimension(Node):
         exists, if no primary data exists in the input time range, or if the
         input time range is invalid.
         """
-        model = info.context.model
-        primary_dataset = model.primary_dataset
-        reference_dataset = model.reference_dataset
-        if reference_dataset is None or not time_range.is_valid():
-            return None
-        embedding_feature_name = self.name
-        reference_embeddings_column = reference_dataset.get_embedding_vector_column(
-            embedding_feature_name
-        )
-        reference_embeddings = _to_array(reference_embeddings_column)
-        primary_embeddings = _get_embeddings_array_for_time_range(
-            dataset=primary_dataset,
-            embedding_feature_name=embedding_feature_name,
-            start=time_range.start,
-            end=time_range.end,
-        )
-        if primary_embeddings is None:
-            return None
-        primary_centroid = _compute_mean_vector(primary_embeddings)
-        reference_centroid = _compute_mean_vector(reference_embeddings)
-        if metric is DriftMetric.euclideanDistance:
-            return euclidean_distance(primary_centroid, reference_centroid)
-        raise NotImplementedError(f'Metric "{metric}" has not been implemented.')
+        if len(
+            data := DriftTimeSeries(
+                str(info.context.model.primary_dataset.get_embedding_vector_column(self.name).name),
+                info.context.model,
+                metric,
+                time_range,
+            ).data
+        ):
+            return data.pop().value
+        return None
 
     @strawberry.field(
         description=(
@@ -114,37 +88,22 @@ class EmbeddingDimension(Node):
         time_range: TimeRange,
         granularity: Granularity,
     ) -> DataQualityTimeSeries:
-        metric_cls = METRICS.get(metric.value, None)
-        if not metric_cls or not issubclass(metric_cls, UnaryOperator):
-            raise NotImplementedError(f"Metric {metric} is not implemented.")
-        dataset = info.context.model.primary_dataset
-        metric_instance = metric_cls(dataset.get_embedding_vector_column(self.name).name)
-        return dataset.dataframe.pipe(
-            timeseries(
-                start_time=time_range.start,
-                end_time=time_range.end,
-                evaluation_window=timedelta(minutes=granularity.evaluation_window_minutes),
-                sampling_interval=timedelta(minutes=granularity.sampling_interval_minutes),
-            ),
-            metrics=(metric_instance,),
-        ).pipe(
-            to_gql_timeseries,
-            metric=metric_instance,
-            timestamps=to_timestamps(time_range, granularity),
+        return DataQualityTimeSeries(
+            str(info.context.model.primary_dataset.get_embedding_vector_column(self.name).name),
+            info.context.model,
+            metric,
+            time_range,
+            granularity,
         )
 
     @strawberry.field
     def drift_time_series(
         self,
-        metric: DriftMetric,
-        time_range: Annotated[
-            TimeRange,
-            strawberry.argument(
-                description="The time range of the primary dataset",
-            ),
-        ],
         info: Info[Context, None],
-    ) -> Optional[DriftTimeSeries]:
+        metric: DriftMetric,
+        time_range: TimeRange,
+        granularity: Granularity,
+    ) -> DriftTimeSeries:
         """
         Computes a drift time-series between the primary and reference datasets.
         The output drift time-series contains one data point for each whole hour
@@ -156,46 +115,13 @@ class EmbeddingDimension(Node):
         Returns None if no reference dataset exists or if the input time range
         is invalid.
         """
-        model = info.context.model
-        primary_dataset = model.primary_dataset
-        reference_dataset = model.reference_dataset
-        if reference_dataset is None or not time_range.is_valid():
-            return None
-        embedding_feature_name = self.name
-        reference_embeddings_column = reference_dataset.get_embedding_vector_column(
-            embedding_feature_name
+        return DriftTimeSeries(
+            str(info.context.model.primary_dataset.get_embedding_vector_column(self.name).name),
+            info.context.model,
+            metric,
+            time_range,
+            granularity,
         )
-        reference_embeddings = _to_array(reference_embeddings_column)
-        reference_centroid = _compute_mean_vector(reference_embeddings)
-        time_series_data_points = []
-        if metric is DriftMetric.euclideanDistance:
-            eval_window_end = time_range.start
-            while eval_window_end < time_range.end:
-                eval_window_start = (
-                    eval_window_end - DRIFT_EVAL_WINDOW_NUM_INTERVALS * EVAL_INTERVAL_LENGTH
-                )
-                primary_embeddings = _get_embeddings_array_for_time_range(
-                    dataset=primary_dataset,
-                    embedding_feature_name=embedding_feature_name,
-                    start=eval_window_start,
-                    end=eval_window_end,
-                )
-                distance: Optional[float] = None
-                if primary_embeddings is not None:
-                    primary_centroid = _compute_mean_vector(primary_embeddings)
-                    distance = euclidean_distance(
-                        reference_centroid,
-                        primary_centroid,
-                    )
-                time_series_data_points.append(
-                    TimeSeriesDataPoint(
-                        timestamp=eval_window_end,
-                        value=distance,
-                    )
-                )
-                eval_window_end += EVAL_INTERVAL_LENGTH
-            return DriftTimeSeries(data=time_series_data_points)
-        raise NotImplementedError(f'Metric "{metric}" has not been implemented.')
 
     @strawberry.field
     def UMAPPoints(
@@ -276,7 +202,7 @@ class EmbeddingDimension(Node):
         min_dist = DEFAULT_MIN_DIST if min_dist is None else min_dist
         n_neighbors = DEFAULT_N_NEIGHBORS if n_neighbors is None else n_neighbors
 
-        vectors, clusters = PointCloud(
+        vectors, cluster_membership = PointCloud(
             dimensionalityReducer=Umap(n_neighbors=n_neighbors, min_dist=min_dist),
             clustersFinder=Hdbscan(),
         ).generate(data, n_components=n_components)
@@ -341,10 +267,12 @@ class EmbeddingDimension(Node):
                 )
             )
 
+        has_reference_data = datasets[DatasetType.REFERENCE] is not None
+
         return UMAPPoints(
             data=points[DatasetType.PRIMARY],
             reference_data=points[DatasetType.REFERENCE],
-            clusters=to_gql_clusters(clusters),
+            clusters=to_gql_clusters(cluster_membership, has_reference_data=has_reference_data),
         )
 
 

phoenix/server/api/types/TimeSeries.py

@@ -1,9 +1,21 @@
-from datetime import datetime
+from datetime import datetime, timedelta
 from functools import total_ordering
-from typing import List, Optional
+from typing import Iterable, List, Optional, Union, cast
 
+import pandas as pd
 import strawberry
 
+from phoenix.core.model import Model
+from phoenix.metrics import Metric
+from phoenix.metrics.mixins import DriftOperator
+from phoenix.metrics.timeseries import timeseries
+from phoenix.server.api.input_types.Granularity import Granularity, to_timestamps
+from phoenix.server.api.input_types.TimeRange import TimeRange
+from phoenix.server.api.interceptor import NoneIfNan
+from phoenix.server.api.types import METRICS
+from phoenix.server.api.types.DataQualityMetric import DataQualityMetric
+from phoenix.server.api.types.DriftMetric import DriftMetric
+
 
 @strawberry.type
 @total_ordering
@@ -14,14 +26,116 @@ class TimeSeriesDataPoint:
     timestamp: datetime
 
     """The value of the data point"""
-    value: Optional[float]
+    value: Optional[float] = strawberry.field(default=NoneIfNan())
 
     def __lt__(self, other: "TimeSeriesDataPoint") -> bool:
         return self.timestamp < other.timestamp
 
 
+def to_gql_datapoints(
+    df: pd.DataFrame, metric: Metric, timestamps: Iterable[datetime]
+) -> List[TimeSeriesDataPoint]:
+    data = []
+    for timestamp in timestamps:
+        try:
+            row = df.iloc[cast(int, df.index.get_loc(timestamp)), :].to_dict()
+        except KeyError:
+            row = {}
+        data.append(
+            TimeSeriesDataPoint(
+                timestamp=timestamp,
+                value=metric.get_value(row),
+            )
+        )
+    return sorted(data)
+
+
 @strawberry.interface
 class TimeSeries:
     """A collection of data points over time"""
 
     data: List[TimeSeriesDataPoint]
+
+    def __init__(
+        self,
+        column_name: str,
+        model: Model,
+        metric: Union[DriftMetric, DataQualityMetric],
+        time_range: Optional[TimeRange] = None,
+        granularity: Optional[Granularity] = None,
+    ):
+        if not (metric_cls := METRICS.get(metric.value, None)):
+            raise NotImplementedError(f"Metric {metric} is not implemented.")
+        dataset = model.primary_dataset
+        metric_instance = metric_cls(column_name=column_name)
+        if (
+            issubclass(metric_cls, DriftOperator)
+            and (ref_dataset := model.reference_dataset) is not None
+        ):
+            metric_instance.reference_data = ref_dataset.dataframe
+        if time_range is None:
+            time_range = TimeRange(
+                start=dataset.start_time,
+                end=dataset.end_time,
+            )
+        if granularity is None:
+            total_minutes = int((time_range.end - time_range.start).total_seconds()) // 60
+            granularity = Granularity(
+                evaluation_window_minutes=total_minutes,
+                sampling_interval_minutes=total_minutes,
+            )
+        self.data = dataset.dataframe.pipe(
+            timeseries(
+                start_time=time_range.start,
+                end_time=time_range.end,
+                evaluation_window=timedelta(minutes=granularity.evaluation_window_minutes),
+                sampling_interval=timedelta(minutes=granularity.sampling_interval_minutes),
+            ),
+            metrics=(metric_instance,),
+        ).pipe(
+            to_gql_datapoints,
+            metric=metric_instance,
+            timestamps=to_timestamps(time_range, granularity),
+        )
+
+
+@strawberry.type
+class DataQualityTimeSeries(TimeSeries):
+    """A time series of data quality metrics"""
+
+    def __init__(
+        self,
+        column_name: str,
+        model: Model,
+        metric: DataQualityMetric,
+        time_range: Optional[TimeRange] = None,
+        granularity: Optional[Granularity] = None,
+    ):
+        super().__init__(
+            column_name,
+            model,
+            metric,
+            time_range,
+            granularity,
+        )
+
+
+@strawberry.type
+class DriftTimeSeries(TimeSeries):
+    """A time series of drift metrics"""
+
+    def __init__(
+        self,
+        column_name: str,
+        model: Model,
+        metric: DriftMetric,
+        time_range: Optional[TimeRange] = None,
+        granularity: Optional[Granularity] = None,
+    ):
+        super().__init__(
+            column_name,
+            model,
+            metric,
+            time_range,
+            granularity,
+        )

phoenix/server/api/types/UMAPPoints.py

@@ -1,4 +1,4 @@
-from typing import List, Union
+from typing import Dict, List, Optional, Set, Union
 
 import numpy as np
 import numpy.typing as npt
@@ -6,11 +6,70 @@ import strawberry
 from strawberry.scalars import ID
 from typing_extensions import TypeAlias
 
+from phoenix.core.embedding_dimension import calculate_drift_ratio
+from phoenix.datasets.event import EventId
+from phoenix.server.api.interceptor import NoneIfNan
+
 from .EmbeddingMetadata import EmbeddingMetadata
 from .EventMetadata import EventMetadata
 
-EventId: TypeAlias = ID
 ClusterId: TypeAlias = ID
+PointId: TypeAlias = ID
+
+
+@strawberry.type
+class Cluster:
+    """A grouping of points in a UMAP plot"""
+
+    """The ID of the cluster"""
+    id: ClusterId
+
+    """A list of points that belong to the cluster"""
+    point_ids: List[PointId]
+
+    """A list of points that belong to the cluster"""
+    drift_ratio: Optional[float] = strawberry.field(
+        description="ratio of primary points over reference points",
+        default=NoneIfNan(),
+    )
+
+
+def to_gql_clusters(
+    cluster_membership: Dict[EventId, int],
+    has_reference_data: bool,
+) -> List[Cluster]:
+    """
+    Converts a dictionary of event IDs to cluster IDs to a list of clusters for the graphQL response
+
+    Parameters
+    ----------
+    cluster_membership: Dict[EventId, int]
+        A dictionary of event IDs to cluster IDs
+    has_reference_data: bool
+        Whether or not the model has reference data
+        Used to determine if drift ratio should be calculated
+    """
+
+    clusters: Dict[int, Set[EventId]] = {}
+    for event_id, cluster_id in cluster_membership.items():
+        if cluster_id in clusters:
+            clusters[cluster_id].add(event_id)
+        else:
+            clusters[cluster_id] = {event_id}
+
+    gql_clusters: List[Cluster] = []
+    for cluster_id, cluster_events in clusters.items():
+        gql_clusters.append(
+            Cluster(
+                id=ID(str(cluster_id)),
+                point_ids=[ID(str(event)) for event in cluster_events],
+                drift_ratio=calculate_drift_ratio(cluster_events)
+                if has_reference_data
+                else float("nan"),
+            )
+        )
+
+    return gql_clusters
 
 
 @strawberry.type
@@ -39,7 +98,7 @@ class UMAPPoint:
     """point and metadata for a UMAP plot"""
 
     """A unique ID for the the point"""
-    id: EventId
+    id: PointId
 
     """The coordinates of the point. Can be two or three dimensional"""
     coordinates: Union[Point2D, Point3D]
@@ -51,17 +110,6 @@ class UMAPPoint:
     event_metadata: EventMetadata
 
 
-@strawberry.type
-class Cluster:
-    """A grouping of points in a UMAP plot"""
-
-    """The ID of the cluster"""
-    id: ClusterId
-
-    """A list of points that belong to the cluster"""
-    point_ids: List[EventId]
-
-
 @strawberry.type
 class UMAPPoints:
     data: List[UMAPPoint]

phoenix/server/main.py

@@ -8,12 +8,12 @@ from typing import Optional
 import uvicorn
 
 import phoenix.config as config
-from phoenix.server.app import create_app
-from phoenix.server.fixtures import (
+from phoenix.datasets.fixtures import (
     FIXTURES,
     download_fixture_if_missing,
     get_dataset_names_from_fixture_name,
 )
+from phoenix.server.app import create_app
 
 logger = logging.getLogger(__name__)
 
@@ -66,7 +66,7 @@ if __name__ == "__main__":
         primary_dataset_name, reference_dataset_name = get_dataset_names_from_fixture_name(
             fixture_name
         )
-        print(f'🌎 Downloading fixture "{fixture_name}" if missing')
+        print(f'🌎 Initializing fixture: "{fixture_name}"')
         download_fixture_if_missing(fixture_name)
 
     print(f"1️⃣ primary dataset: {primary_dataset_name}")