arize-phoenix 0.0.18__py3-none-any.whl → 0.0.19__py3-none-any.whl

phoenix/server/api/types/Dimension.py

@@ -2,17 +2,21 @@ from typing import List, Optional
 
 import strawberry
 from strawberry.types import Info
+from typing_extensions import Annotated
 
-from phoenix.core.model_schema import REFERENCE, ScalarDimension
+from phoenix.core.model_schema import PRIMARY, REFERENCE, ScalarDimension
 
 from ..context import Context
 from ..input_types.Granularity import Granularity
 from ..input_types.TimeRange import TimeRange
 from .DataQualityMetric import DataQualityMetric
+from .DatasetRole import DatasetRole
 from .DimensionDataType import DimensionDataType
+from .DimensionShape import DimensionShape
 from .DimensionType import DimensionType
 from .node import Node
 from .ScalarDriftMetricEnum import ScalarDriftMetric
+from .Segments import DatasetValues, Segments
 from .TimeSeries import (
     DataQualityTimeSeries,
     DriftTimeSeries,
@@ -31,6 +35,9 @@ class Dimension(Node):
     dataType: DimensionDataType = strawberry.field(
         description="The data type of the column. Categorical or numeric."
     )
+    shape: DimensionShape = strawberry.field(
+        description="Whether the dimension data is continuous or discrete."
+    )
     dimension: strawberry.Private[ScalarDimension]
 
     @strawberry.field
@@ -50,8 +57,17 @@ class Dimension(Node):
         model = info.context.model
         if model[REFERENCE].empty:
             return None
-        time_range, granularity = ensure_timeseries_parameters(model, time_range)
-        data = get_drift_timeseries_data(self.dimension, metric, time_range, granularity)
+        dataset = model[PRIMARY]
+        time_range, granularity = ensure_timeseries_parameters(
+            dataset,
+            time_range,
+        )
+        data = get_drift_timeseries_data(
+            self.dimension,
+            metric,
+            time_range,
+            granularity,
+        )
         return data[0].value if len(data) else None
 
     @strawberry.field
@@ -60,9 +76,27 @@ class Dimension(Node):
         info: Info[Context, None],
         metric: DataQualityMetric,
         time_range: Optional[TimeRange] = None,
+        dataset_role: Annotated[
+            Optional[DatasetRole],
+            strawberry.argument(
+                description="The dataset (primary or reference) to query",
+            ),
+        ] = DatasetRole.primary,
     ) -> Optional[float]:
-        time_range, granularity = ensure_timeseries_parameters(info.context.model, time_range)
-        data = get_data_quality_timeseries_data(self.dimension, metric, time_range, granularity)
+        if dataset_role is None:
+            dataset_role = DatasetRole.primary
+        dataset = info.context.model[dataset_role.value]
+        time_range, granularity = ensure_timeseries_parameters(
+            dataset,
+            time_range,
+        )
+        data = get_data_quality_timeseries_data(
+            self.dimension,
+            metric,
+            time_range,
+            granularity,
+            dataset_role,
+        )
         return data[0].value if len(data) else None
 
     @strawberry.field(
@@ -89,17 +123,34 @@ class Dimension(Node):
         metric: DataQualityMetric,
         time_range: TimeRange,
         granularity: Granularity,
+        dataset_role: Annotated[
+            Optional[DatasetRole],
+            strawberry.argument(
+                description="The dataset (primary or reference) to query",
+            ),
+        ] = DatasetRole.primary,
     ) -> DataQualityTimeSeries:
+        if dataset_role is None:
+            dataset_role = DatasetRole.primary
+        dataset = info.context.model[dataset_role.value]
         time_range, granularity = ensure_timeseries_parameters(
-            info.context.model, time_range, granularity
+            dataset,
+            time_range,
+            granularity,
         )
         return DataQualityTimeSeries(
-            data=get_data_quality_timeseries_data(self.dimension, metric, time_range, granularity)
+            data=get_data_quality_timeseries_data(
+                self.dimension,
+                metric,
+                time_range,
+                granularity,
+                dataset_role,
+            )
         )
 
     @strawberry.field(
         description=(
-            "Returns the time series of the specified metric for data within a time range. Data"
+            "The time series of the specified metric for data within a time range. Data"
             " points are generated starting at the end time and are separated by the sampling"
             " interval. Each data point is labeled by the end instant and contains data from their"
             " respective evaluation windows."
@@ -115,11 +166,32 @@ class Dimension(Node):
         model = info.context.model
         if model[REFERENCE].empty:
             return DriftTimeSeries(data=[])
-        time_range, granularity = ensure_timeseries_parameters(model, time_range, granularity)
+        dataset = model[PRIMARY]
+        time_range, granularity = ensure_timeseries_parameters(
+            dataset,
+            time_range,
+            granularity,
+        )
         return DriftTimeSeries(
-            data=get_drift_timeseries_data(self.dimension, metric, time_range, granularity)
+            data=get_drift_timeseries_data(
+                self.dimension,
+                metric,
+                time_range,
+                granularity,
+            )
         )
 
+    @strawberry.field(
+        description="Returns the segments across both datasets and returns the counts per segment",
+    )  # type: ignore
+    def segments_comparison(
+        self,
+        primary_time_range: Optional[TimeRange] = strawberry.UNSET,
+    ) -> Segments:
+        # TODO: Implement binning across primary and reference
+
+        return Segments(segments=[], total_counts=DatasetValues(primary_value=0, reference_value=0))
+
 
 def to_gql_dimension(id_attr: int, dimension: ScalarDimension) -> Dimension:
     """
@@ -128,7 +200,8 @@ def to_gql_dimension(id_attr: int, dimension: ScalarDimension) -> Dimension:
     return Dimension(
         id_attr=id_attr,
         name=dimension.name,
-        type=DimensionType.from_(dimension),
-        dataType=DimensionDataType.from_(dimension),
+        type=DimensionType.from_dimension(dimension),
+        dataType=DimensionDataType.from_dimension(dimension),
         dimension=dimension,
+        shape=DimensionShape.from_dimension(dimension),
     )

phoenix/server/api/types/DimensionDataType.py

@@ -11,8 +11,8 @@ class DimensionDataType(Enum):
     numeric = "numeric"
 
     @classmethod
-    def from_(cls, dim: Dimension) -> "DimensionDataType":
-        data_type = dim.data_type
+    def from_dimension(cls, dimension: Dimension) -> "DimensionDataType":
+        data_type = dimension.data_type
         if data_type in (CONTINUOUS,):
             return cls.numeric
         return cls.categorical

phoenix/server/api/types/DimensionShape.py

@@ -0,0 +1,21 @@
+from enum import Enum
+
+import strawberry
+
+from phoenix.core.model_schema import CONTINUOUS, Dimension
+
+
+@strawberry.enum
+class DimensionShape(Enum):
+    continuous = "continuous"
+    discrete = "discrete"
+
+    @classmethod
+    def from_dimension(cls, dim: Dimension) -> "DimensionShape":
+        data_type = dim.data_type
+        if data_type in (CONTINUOUS,):
+            return cls.continuous
+
+        # For now we assume all non-continuous data is discrete
+        # E.g. floats are the only dimension data type that is continuous
+        return cls.discrete

phoenix/server/api/types/DimensionType.py

@@ -21,7 +21,7 @@ class DimensionType(Enum):
     actual = "actual"
 
     @classmethod
-    def from_(cls, dim: Dimension) -> "DimensionType":
+    def from_dimension(cls, dim: Dimension) -> "DimensionType":
         role = dim.role
         if role in (FEATURE,):
             return cls.feature

phoenix/server/api/types/EmbeddingDimension.py

@@ -20,7 +20,6 @@ from phoenix.core.model_schema import (
     PRIMARY,
     REFERENCE,
     Dataset,
-    DatasetRole,
     EventId,
 )
 from phoenix.metrics.timeseries import row_interval_from_sorted_time_index
@@ -29,6 +28,7 @@ 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 phoenix.server.api.types.DatasetRole import DatasetRole
 from phoenix.server.api.types.VectorDriftMetricEnum import VectorDriftMetric
 
 from ..input_types.Granularity import Granularity
@@ -83,8 +83,17 @@ class EmbeddingDimension(Node):
         model = info.context.model
         if model[REFERENCE].empty:
             return None
-        time_range, granularity = ensure_timeseries_parameters(model, time_range)
-        data = get_drift_timeseries_data(self.dimension, metric, time_range, granularity)
+        dataset = model[PRIMARY]
+        time_range, granularity = ensure_timeseries_parameters(
+            dataset,
+            time_range,
+        )
+        data = get_drift_timeseries_data(
+            self.dimension,
+            metric,
+            time_range,
+            granularity,
+        )
         return data[0].value if len(data) else None
 
     @strawberry.field(
@@ -101,12 +110,29 @@ class EmbeddingDimension(Node):
         metric: DataQualityMetric,
         time_range: TimeRange,
         granularity: Granularity,
+        dataset_role: Annotated[
+            Optional[DatasetRole],
+            strawberry.argument(
+                description="The dataset (primary or reference) to query",
+            ),
+        ] = DatasetRole.primary,
     ) -> DataQualityTimeSeries:
+        if dataset_role is None:
+            dataset_role = DatasetRole.primary
+        dataset = info.context.model[dataset_role.value]
         time_range, granularity = ensure_timeseries_parameters(
-            info.context.model, time_range, granularity
+            dataset,
+            time_range,
+            granularity,
         )
         return DataQualityTimeSeries(
-            data=get_data_quality_timeseries_data(self.dimension, metric, time_range, granularity)
+            data=get_data_quality_timeseries_data(
+                self.dimension,
+                metric,
+                time_range,
+                granularity,
+                dataset_role,
+            )
         )
 
     @strawberry.field(
@@ -129,9 +155,19 @@ class EmbeddingDimension(Node):
         model = info.context.model
         if model[REFERENCE].empty:
             return DriftTimeSeries(data=[])
-        time_range, granularity = ensure_timeseries_parameters(model, time_range, granularity)
+        dataset = model[PRIMARY]
+        time_range, granularity = ensure_timeseries_parameters(
+            dataset,
+            time_range,
+            granularity,
+        )
         return DriftTimeSeries(
-            data=get_drift_timeseries_data(self.dimension, metric, time_range, granularity)
+            data=get_drift_timeseries_data(
+                self.dimension,
+                metric,
+                time_range,
+                granularity,
+            )
         )
 
     @strawberry.field
@@ -226,7 +262,7 @@ class EmbeddingDimension(Node):
             ),
         ).generate(data, n_components=n_components)
 
-        points: Dict[DatasetRole, List[UMAPPoint]] = defaultdict(list)
+        points: Dict[ms.DatasetRole, List[UMAPPoint]] = defaultdict(list)
         for event_id, vector in vectors.items():
             row_id = event_id.row_id
             dataset_id = event_id.dataset_id

phoenix/server/api/types/NumericRange.py

@@ -0,0 +1,10 @@
+import strawberry
+
+
+@strawberry.type
+class NumericRange:
+    """A numeric range to denote a bin or domain"""
+
+    start: float
+    end: float
+    # TODO consider denoting right open or closed

phoenix/server/api/types/Segments.py

@@ -0,0 +1,44 @@
+from typing import List, Optional
+
+import strawberry
+
+from .NumericRange import NumericRange
+
+
+@strawberry.type
+class NominalBin:
+    """A bin that contains a discrete value"""
+
+    name: str
+
+
+@strawberry.type
+class IntervalBin:
+    """A bin that contains a discrete value"""
+
+    # TODO figure out the empty case
+    range: NumericRange
+
+
+@strawberry.type
+class DatasetValues:
+    """Numeric values per dataset role"""
+
+    primary_value: Optional[float]
+    reference_value: Optional[float]
+
+
+@strawberry.type
+class Segment:
+    """A segment of the parent's data, split out using a heuristic"""
+
+    bin: strawberry.union("Bin", types=(NominalBin, IntervalBin))  # type: ignore
+    counts: DatasetValues
+    # TODO add support for a "z" metric list
+    # values: List[Optional[float]]
+
+
+@strawberry.type
+class Segments:
+    segments: List[Segment]
+    total_counts: DatasetValues

phoenix/server/api/types/TimeSeries.py

@@ -5,7 +5,7 @@ from typing import Iterable, List, Optional, Tuple, Union, cast
 import pandas as pd
 import strawberry
 
-from phoenix.core.model_schema import CONTINUOUS, PRIMARY, REFERENCE, Dimension, Model
+from phoenix.core.model_schema import CONTINUOUS, REFERENCE, Dataset, Dimension
 from phoenix.metrics import Metric, binning
 from phoenix.metrics.mixins import DriftOperator
 from phoenix.metrics.timeseries import timeseries
@@ -14,6 +14,7 @@ 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.DatasetRole import DatasetRole
 from phoenix.server.api.types.ScalarDriftMetricEnum import ScalarDriftMetric
 from phoenix.server.api.types.VectorDriftMetricEnum import VectorDriftMetric
 
@@ -63,10 +64,11 @@ def _get_timeseries_data(
     metric: Union[ScalarDriftMetric, VectorDriftMetric, DataQualityMetric],
     time_range: TimeRange,
     granularity: Granularity,
+    dataset_role: DatasetRole,
 ) -> List[TimeSeriesDataPoint]:
     if not (metric_cls := METRICS.get(metric.value, None)):
         raise NotImplementedError(f"Metric {metric} is not implemented.")
-    data = dimension[PRIMARY]
+    data = dimension[dataset_role.value]
     metric_instance = metric_cls(operand_column_name=dimension.name)
     if issubclass(metric_cls, DriftOperator):
         ref_data = dimension[REFERENCE]
@@ -105,8 +107,15 @@ def get_data_quality_timeseries_data(
     metric: DataQualityMetric,
     time_range: TimeRange,
     granularity: Granularity,
+    dataset_role: DatasetRole,
 ) -> List[TimeSeriesDataPoint]:
-    return _get_timeseries_data(dimension, metric, time_range, granularity)
+    return _get_timeseries_data(
+        dimension,
+        metric,
+        time_range,
+        granularity,
+        dataset_role,
+    )
 
 
 @strawberry.type
@@ -120,16 +129,22 @@ def get_drift_timeseries_data(
     time_range: TimeRange,
     granularity: Granularity,
 ) -> List[TimeSeriesDataPoint]:
-    return _get_timeseries_data(dimension, metric, time_range, granularity)
+    return _get_timeseries_data(
+        dimension,
+        metric,
+        time_range,
+        granularity,
+        DatasetRole.primary,
+    )
 
 
 def ensure_timeseries_parameters(
-    model: Model,
+    dataset: Dataset,
     time_range: Optional[TimeRange] = None,
     granularity: Optional[Granularity] = None,
 ) -> Tuple[TimeRange, Granularity]:
     if time_range is None:
-        start, end = model[PRIMARY].time_range
+        start, end = dataset.time_range
         time_range = TimeRange(start=start, end=end)
     if granularity is None:
         total_minutes = int((time_range.end - time_range.start).total_seconds()) // 60

phoenix/server/main.py

@@ -11,7 +11,7 @@ import uvicorn
 import phoenix.config as config
 from phoenix.core.model_schema_adapter import create_model_from_datasets
 from phoenix.datasets import Dataset
-from phoenix.datasets.fixtures import FIXTURES, download_fixture_if_missing
+from phoenix.datasets.fixtures import FIXTURES, get_datasets
 from phoenix.server.app import create_app
 
 logger = logging.getLogger(__name__)
@@ -48,6 +48,7 @@ if __name__ == "__main__":
     parser = ArgumentParser()
     parser.add_argument("--export_path")
     parser.add_argument("--port", type=int, default=config.PORT)
+    parser.add_argument("--no-internet", action="store_true")
     parser.add_argument("--debug", action="store_false")  # TODO: Disable before public launch
     subparsers = parser.add_subparsers(dest="command", required=True)
     datasets_parser = subparsers.add_parser("datasets")
@@ -70,7 +71,10 @@ if __name__ == "__main__":
     else:
         fixture_name = args.fixture
         primary_only = args.primary_only
-        primary_dataset, reference_dataset = download_fixture_if_missing(fixture_name)
+        primary_dataset, reference_dataset = get_datasets(
+            fixture_name,
+            args.no_internet,
+        )
         if primary_only:
             reference_dataset_name = None
             reference_dataset = None