arize-phoenix 4.5.0__py3-none-any.whl → 4.6.1__py3-none-any.whl

phoenix/server/api/types/Dimension.py

@@ -4,6 +4,7 @@ from typing import Any, Dict, List, Optional
 import pandas as pd
 import strawberry
 from strawberry import UNSET
+from strawberry.relay import Node, NodeID
 from strawberry.types import Info
 from typing_extensions import Annotated
 
@@ -17,12 +18,11 @@ 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 .DatasetValues import DatasetValues
 from .DimensionDataType import DimensionDataType
 from .DimensionShape import DimensionShape
 from .DimensionType import DimensionType
-from .node import Node
+from .InferencesRole import InferencesRole
 from .ScalarDriftMetricEnum import ScalarDriftMetric
 from .Segments import (
     GqlBinFactory,
@@ -40,6 +40,7 @@ from .TimeSeries import (
 
 @strawberry.type
 class Dimension(Node):
+    id_attr: NodeID[int]
     name: str = strawberry.field(description="The name of the dimension (a.k.a. the column name)")
     type: DimensionType = strawberry.field(
         description="Whether the dimension represents a feature, tag, prediction, or actual."
@@ -62,16 +63,16 @@ class Dimension(Node):
         """
         Computes a drift metric between all reference data and the primary data
         belonging to the input time range (inclusive of the time range start and
-        exclusive of the time range end). Returns None if no reference dataset
-        exists, if no primary data exists in the input time range, or if the
+        exclusive of the time range end). Returns None if no reference inferences
+        exist, if no primary data exists in the input time range, or if the
         input time range is invalid.
         """
         model = info.context.model
         if model[REFERENCE].empty:
             return None
-        dataset = model[PRIMARY]
+        inferences = model[PRIMARY]
         time_range, granularity = ensure_timeseries_parameters(
-            dataset,
+            inferences,
             time_range,
         )
         data = get_drift_timeseries_data(
@@ -92,18 +93,18 @@ class Dimension(Node):
         info: Info[Context, None],
         metric: DataQualityMetric,
         time_range: Optional[TimeRange] = UNSET,
-        dataset_role: Annotated[
-            Optional[DatasetRole],
+        inferences_role: Annotated[
+            Optional[InferencesRole],
             strawberry.argument(
-                description="The dataset (primary or reference) to query",
+                description="The inferences (primary or reference) to query",
             ),
-        ] = DatasetRole.primary,
+        ] = InferencesRole.primary,
     ) -> Optional[float]:
-        if not isinstance(dataset_role, DatasetRole):
-            dataset_role = DatasetRole.primary
-        dataset = info.context.model[dataset_role.value]
+        if not isinstance(inferences_role, InferencesRole):
+            inferences_role = InferencesRole.primary
+        inferences = info.context.model[inferences_role.value]
         time_range, granularity = ensure_timeseries_parameters(
-            dataset,
+            inferences,
             time_range,
         )
         data = get_data_quality_timeseries_data(
@@ -111,7 +112,7 @@ class Dimension(Node):
             metric,
             time_range,
             granularity,
-            dataset_role,
+            inferences_role,
         )
         return data[0].value if len(data) else None
 
@@ -139,18 +140,18 @@ class Dimension(Node):
         metric: DataQualityMetric,
         time_range: TimeRange,
         granularity: Granularity,
-        dataset_role: Annotated[
-            Optional[DatasetRole],
+        inferences_role: Annotated[
+            Optional[InferencesRole],
             strawberry.argument(
-                description="The dataset (primary or reference) to query",
+                description="The inferences (primary or reference) to query",
             ),
-        ] = DatasetRole.primary,
+        ] = InferencesRole.primary,
     ) -> DataQualityTimeSeries:
-        if not isinstance(dataset_role, DatasetRole):
-            dataset_role = DatasetRole.primary
-        dataset = info.context.model[dataset_role.value]
+        if not isinstance(inferences_role, InferencesRole):
+            inferences_role = InferencesRole.primary
+        inferences = info.context.model[inferences_role.value]
         time_range, granularity = ensure_timeseries_parameters(
-            dataset,
+            inferences,
             time_range,
             granularity,
         )
@@ -160,7 +161,7 @@ class Dimension(Node):
                 metric,
                 time_range,
                 granularity,
-                dataset_role,
+                inferences_role,
             )
         )
 
@@ -182,9 +183,9 @@ class Dimension(Node):
         model = info.context.model
         if model[REFERENCE].empty:
             return DriftTimeSeries(data=[])
-        dataset = model[PRIMARY]
+        inferences = model[PRIMARY]
         time_range, granularity = ensure_timeseries_parameters(
-            dataset,
+            inferences,
             time_range,
             granularity,
         )
@@ -202,7 +203,7 @@ class Dimension(Node):
         )
 
     @strawberry.field(
-        description="Returns the segments across both datasets and returns the counts per segment",
+        description="The segments across both inference sets and returns the counts per segment",
     )  # type: ignore
     def segments_comparison(
         self,
@@ -249,8 +250,8 @@ class Dimension(Node):
         if isinstance(binning_method, binning.IntervalBinning) and binning_method.bins is not None:
             all_bins = all_bins.union(binning_method.bins)
         for bin in all_bins:
-            values: Dict[ms.DatasetRole, Any] = defaultdict(lambda: None)
-            for role in ms.DatasetRole:
+            values: Dict[ms.InferencesRole, Any] = defaultdict(lambda: None)
+            for role in ms.InferencesRole:
                 if model[role].empty:
                     continue
                 try:

phoenix/server/api/types/EmbeddingDimension.py

@@ -8,6 +8,7 @@ import numpy.typing as npt
 import pandas as pd
 import strawberry
 from strawberry import UNSET
+from strawberry.relay import GlobalID, Node, NodeID
 from strawberry.scalars import ID
 from strawberry.types import Info
 from typing_extensions import Annotated
@@ -22,7 +23,7 @@ from phoenix.core.model_schema import (
     PRIMARY,
     PROMPT,
     REFERENCE,
-    Dataset,
+    Inferences,
 )
 from phoenix.metrics.timeseries import row_interval_from_sorted_time_index
 from phoenix.pointcloud.clustering import Hdbscan
@@ -31,7 +32,7 @@ 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.Cluster import to_gql_clusters
-from phoenix.server.api.types.DatasetRole import AncillaryDatasetRole, DatasetRole
+from phoenix.server.api.types.InferencesRole import AncillaryInferencesRole, InferencesRole
 from phoenix.server.api.types.VectorDriftMetricEnum import VectorDriftMetric
 
 from ..input_types.Granularity import Granularity
@@ -39,7 +40,6 @@ from .DataQualityMetric import DataQualityMetric
 from .EmbeddingMetadata import EmbeddingMetadata
 from .Event import create_event_id, unpack_event_id
 from .EventMetadata import EventMetadata
-from .node import GlobalID, Node
 from .Retrieval import Retrieval
 from .TimeSeries import (
     DataQualityTimeSeries,
@@ -70,6 +70,7 @@ CORPUS = "CORPUS"
 class EmbeddingDimension(Node):
     """A embedding dimension of a model. Represents unstructured data"""
 
+    id_attr: NodeID[int]
     name: str
     dimension: strawberry.Private[ms.EmbeddingDimension]
 
@@ -155,16 +156,16 @@ class EmbeddingDimension(Node):
         metric: DataQualityMetric,
         time_range: TimeRange,
         granularity: Granularity,
-        dataset_role: Annotated[
-            Optional[DatasetRole],
+        inferences_role: Annotated[
+            Optional[InferencesRole],
             strawberry.argument(
                 description="The dataset (primary or reference) to query",
             ),
-        ] = DatasetRole.primary,
+        ] = InferencesRole.primary,
     ) -> DataQualityTimeSeries:
-        if not isinstance(dataset_role, DatasetRole):
-            dataset_role = DatasetRole.primary
-        dataset = info.context.model[dataset_role.value]
+        if not isinstance(inferences_role, InferencesRole):
+            inferences_role = InferencesRole.primary
+        dataset = info.context.model[inferences_role.value]
         time_range, granularity = ensure_timeseries_parameters(
             dataset,
             time_range,
@@ -176,7 +177,7 @@ class EmbeddingDimension(Node):
                 metric,
                 time_range,
                 granularity,
-                dataset_role,
+                inferences_role,
             )
         )
 
@@ -314,16 +315,16 @@ class EmbeddingDimension(Node):
         model = info.context.model
         data: Dict[ID, npt.NDArray[np.float64]] = {}
         retrievals: List[Tuple[ID, Any, Any]] = []
-        for dataset in model[Dataset]:
-            dataset_id = dataset.role
-            row_id_start, row_id_stop = 0, len(dataset)
-            if dataset_id is PRIMARY:
+        for inferences in model[Inferences]:
+            inferences_id = inferences.role
+            row_id_start, row_id_stop = 0, len(inferences)
+            if inferences_id is PRIMARY:
                 row_id_start, row_id_stop = row_interval_from_sorted_time_index(
-                    time_index=cast(pd.DatetimeIndex, dataset.index),
+                    time_index=cast(pd.DatetimeIndex, inferences.index),
                     time_start=time_range.start,
                     time_stop=time_range.end,
                 )
-            vector_column = self.dimension[dataset_id]
+            vector_column = self.dimension[inferences_id]
             samples_collected = 0
             for row_id in _row_indices(
                 row_id_start,
@@ -337,7 +338,7 @@ class EmbeddingDimension(Node):
                 # of dunder method __len__.
                 if not hasattr(embedding_vector, "__len__"):
                     continue
-                event_id = create_event_id(row_id, dataset_id)
+                event_id = create_event_id(row_id, inferences_id)
                 data[event_id] = embedding_vector
                 samples_collected += 1
                 if isinstance(
@@ -347,8 +348,8 @@ class EmbeddingDimension(Node):
                     retrievals.append(
                         (
                             event_id,
-                            self.dimension.context_retrieval_ids(dataset).iloc[row_id],
-                            self.dimension.context_retrieval_scores(dataset).iloc[row_id],
+                            self.dimension.context_retrieval_ids(inferences).iloc[row_id],
+                            self.dimension.context_retrieval_scores(inferences).iloc[row_id],
                         )
                     )
 
@@ -357,13 +358,13 @@ class EmbeddingDimension(Node):
             self.dimension,
             ms.RetrievalEmbeddingDimension,
         ) and (corpus := info.context.corpus):
-            corpus_dataset = corpus[PRIMARY]
-            for row_id, document_embedding_vector in enumerate(corpus_dataset[PROMPT]):
+            corpus_inferences = corpus[PRIMARY]
+            for row_id, document_embedding_vector in enumerate(corpus_inferences[PROMPT]):
                 if not hasattr(document_embedding_vector, "__len__"):
                     continue
-                event_id = create_event_id(row_id, AncillaryDatasetRole.corpus)
+                event_id = create_event_id(row_id, AncillaryInferencesRole.corpus)
                 data[event_id] = document_embedding_vector
-            corpus_primary_key = corpus_dataset.primary_key
+            corpus_primary_key = corpus_inferences.primary_key
             for event_id, retrieval_ids, retrieval_scores in retrievals:
                 if not isinstance(retrieval_ids, Iterable):
                     continue
@@ -385,7 +386,7 @@ class EmbeddingDimension(Node):
                         )
                     except KeyError:
                         continue
-                    document_embedding_vector = corpus_dataset[PROMPT].iloc[document_row_id]
+                    document_embedding_vector = corpus_inferences[PROMPT].iloc[document_row_id]
                     if not hasattr(document_embedding_vector, "__len__"):
                         continue
                     context_retrievals.append(
@@ -393,7 +394,7 @@ class EmbeddingDimension(Node):
                             query_id=event_id,
                             document_id=create_event_id(
                                 document_row_id,
-                                AncillaryDatasetRole.corpus,
+                                AncillaryInferencesRole.corpus,
                             ),
                             relevance=document_score,
                         )
@@ -413,11 +414,13 @@ class EmbeddingDimension(Node):
             ),
         ).generate(data, n_components=n_components)
 
-        points: Dict[Union[DatasetRole, AncillaryDatasetRole], List[UMAPPoint]] = defaultdict(list)
+        points: Dict[Union[InferencesRole, AncillaryInferencesRole], List[UMAPPoint]] = defaultdict(
+            list
+        )
         for event_id, vector in vectors.items():
-            row_id, dataset_role = unpack_event_id(event_id)
-            if isinstance(dataset_role, DatasetRole):
-                dataset = model[dataset_role.value]
+            row_id, inferences_role = unpack_event_id(event_id)
+            if isinstance(inferences_role, InferencesRole):
+                dataset = model[inferences_role.value]
                 embedding_metadata = EmbeddingMetadata(
                     prediction_id=dataset[PREDICTION_ID][row_id],
                     link_to_data=dataset[self.dimension.link_to_data][row_id],
@@ -433,9 +436,12 @@ class EmbeddingDimension(Node):
                 )
             else:
                 continue
-            points[dataset_role].append(
+            points[inferences_role].append(
                 UMAPPoint(
-                    id=GlobalID(f"{type(self).__name__}:{str(dataset_role)}", row_id),
+                    id=GlobalID(
+                        type_name=f"{type(self).__name__}:{str(inferences_role)}",
+                        node_id=str(row_id),
+                    ),
                     event_id=event_id,
                     coordinates=to_gql_coordinates(vector),
                     event_metadata=EventMetadata(
@@ -449,12 +455,12 @@ class EmbeddingDimension(Node):
             )
 
         return UMAPPoints(
-            data=points[DatasetRole.primary],
-            reference_data=points[DatasetRole.reference],
+            data=points[InferencesRole.primary],
+            reference_data=points[InferencesRole.reference],
             clusters=to_gql_clusters(
                 clustered_events=clustered_events,
             ),
-            corpus_data=points[AncillaryDatasetRole.corpus],
+            corpus_data=points[AncillaryInferencesRole.corpus],
             context_retrievals=context_retrievals,
         )
 

phoenix/server/api/types/Event.py

@@ -17,10 +17,10 @@ from phoenix.core.model_schema import (
 )
 
 from ..interceptor import GqlValueMediator
-from .DatasetRole import STR_TO_DATASET_ROLE, AncillaryDatasetRole, DatasetRole
 from .Dimension import Dimension
 from .DimensionWithValue import DimensionWithValue
 from .EventMetadata import EventMetadata
+from .InferencesRole import STR_TO_INFEREENCES_ROLE, AncillaryInferencesRole, InferencesRole
 from .PromptResponse import PromptResponse
 
 
@@ -41,35 +41,35 @@ class Event:
 
 def create_event_id(
     row_id: int,
-    dataset_role: Union[DatasetRole, AncillaryDatasetRole, ms.DatasetRole],
+    inferences_role: Union[InferencesRole, AncillaryInferencesRole, ms.InferencesRole],
 ) -> ID:
-    dataset_role_str = (
-        dataset_role.value
-        if isinstance(dataset_role, (DatasetRole, AncillaryDatasetRole))
-        else dataset_role
+    inferences_role_str = (
+        inferences_role.value
+        if isinstance(inferences_role, (InferencesRole, AncillaryInferencesRole))
+        else inferences_role
     )
-    return ID(f"{row_id}:{dataset_role_str}")
+    return ID(f"{row_id}:{inferences_role_str}")
 
 
 def unpack_event_id(
     event_id: ID,
-) -> Tuple[int, Union[DatasetRole, AncillaryDatasetRole]]:
-    row_id_str, dataset_role_str = str(event_id).split(":")
+) -> Tuple[int, Union[InferencesRole, AncillaryInferencesRole]]:
+    row_id_str, inferences_role_str = str(event_id).split(":")
     row_id = int(row_id_str)
-    dataset_role = STR_TO_DATASET_ROLE[dataset_role_str]
-    return row_id, dataset_role
+    inferences_role = STR_TO_INFEREENCES_ROLE[inferences_role_str]
+    return row_id, inferences_role
 
 
-def parse_event_ids_by_dataset_role(
+def parse_event_ids_by_inferences_role(
     event_ids: List[ID],
-) -> Dict[Union[DatasetRole, AncillaryDatasetRole], List[int]]:
+) -> Dict[Union[InferencesRole, AncillaryInferencesRole], List[int]]:
     """
     Parses event IDs and returns the corresponding row indexes.
     """
-    row_indexes: Dict[Union[DatasetRole, AncillaryDatasetRole], List[int]] = defaultdict(list)
+    row_indexes: Dict[Union[InferencesRole, AncillaryInferencesRole], List[int]] = defaultdict(list)
     for event_id in event_ids:
-        row_id, dataset_role = unpack_event_id(event_id)
-        row_indexes[dataset_role].append(row_id)
+        row_id, inferences_role = unpack_event_id(event_id)
+        row_indexes[inferences_role].append(row_id)
     return row_indexes
 
 

phoenix/server/api/types/ExampleRevisionInterface.py

@@ -0,0 +1,14 @@
+import strawberry
+from strawberry.scalars import JSON
+
+
+@strawberry.interface
+class ExampleRevision:
+    """
+    Represents an example revision for generative tasks.
+    For example, you might have text -> text, text -> labels, etc.
+    """
+
+    input: JSON
+    output: JSON
+    metadata: JSON

phoenix/server/api/types/Experiment.py

@@ -0,0 +1,147 @@
+from datetime import datetime
+from typing import List, Optional
+
+import strawberry
+from sqlalchemy import select
+from sqlalchemy.orm import joinedload
+from strawberry import UNSET, Private
+from strawberry.relay import Connection, Node, NodeID
+from strawberry.scalars import JSON
+from strawberry.types import Info
+
+from phoenix.db import models
+from phoenix.server.api.context import Context
+from phoenix.server.api.types.ExperimentAnnotationSummary import ExperimentAnnotationSummary
+from phoenix.server.api.types.ExperimentRun import ExperimentRun, to_gql_experiment_run
+from phoenix.server.api.types.pagination import (
+    ConnectionArgs,
+    CursorString,
+    connection_from_list,
+)
+from phoenix.server.api.types.Project import Project
+
+
+@strawberry.type
+class Experiment(Node):
+    cached_sequence_number: Private[Optional[int]] = None
+    id_attr: NodeID[int]
+    name: str
+    project_name: Optional[str]
+    description: Optional[str]
+    metadata: JSON
+    created_at: datetime
+    updated_at: datetime
+
+    @strawberry.field(
+        description="Sequence number (1-based) of experiments belonging to the same dataset"
+    )  # type: ignore
+    async def sequence_number(
+        self,
+        info: Info[Context, None],
+    ) -> int:
+        if self.cached_sequence_number is None:
+            seq_num = await info.context.data_loaders.experiment_sequence_number.load(self.id_attr)
+            if seq_num is None:
+                raise ValueError(f"invalid experiment: id={self.id_attr}")
+            self.cached_sequence_number = seq_num
+        return self.cached_sequence_number
+
+    @strawberry.field
+    async def runs(
+        self,
+        info: Info[Context, None],
+        first: Optional[int] = 50,
+        last: Optional[int] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
+    ) -> Connection[ExperimentRun]:
+        args = ConnectionArgs(
+            first=first,
+            after=after if isinstance(after, CursorString) else None,
+            last=last,
+            before=before if isinstance(before, CursorString) else None,
+        )
+        experiment_id = self.id_attr
+        async with info.context.db() as session:
+            runs = (
+                await session.scalars(
+                    select(models.ExperimentRun)
+                    .where(models.ExperimentRun.experiment_id == experiment_id)
+                    .order_by(models.ExperimentRun.id.desc())
+                    .options(
+                        joinedload(models.ExperimentRun.trace).load_only(models.Trace.trace_id)
+                    )
+                )
+            ).all()
+        return connection_from_list([to_gql_experiment_run(run) for run in runs], args)
+
+    @strawberry.field
+    async def run_count(self, info: Info[Context, None]) -> int:
+        experiment_id = self.id_attr
+        return await info.context.data_loaders.experiment_run_counts.load(experiment_id)
+
+    @strawberry.field
+    async def annotation_summaries(
+        self, info: Info[Context, None]
+    ) -> List[ExperimentAnnotationSummary]:
+        experiment_id = self.id_attr
+        return [
+            ExperimentAnnotationSummary(
+                annotation_name=summary.annotation_name,
+                min_score=summary.min_score,
+                max_score=summary.max_score,
+                mean_score=summary.mean_score,
+                count=summary.count,
+                error_count=summary.error_count,
+            )
+            for summary in await info.context.data_loaders.experiment_annotation_summaries.load(
+                experiment_id
+            )
+        ]
+
+    @strawberry.field
+    async def error_rate(self, info: Info[Context, None]) -> Optional[float]:
+        return await info.context.data_loaders.experiment_error_rates.load(self.id_attr)
+
+    @strawberry.field
+    async def average_run_latency_ms(self, info: Info[Context, None]) -> float:
+        latency_seconds = await info.context.data_loaders.average_experiment_run_latency.load(
+            self.id_attr
+        )
+        return latency_seconds * 1000
+
+    @strawberry.field
+    async def project(self, info: Info[Context, None]) -> Optional[Project]:
+        if self.project_name is None:
+            return None
+
+        db_project = await info.context.data_loaders.project_by_name.load(self.project_name)
+
+        if db_project is None:
+            return None
+
+        return Project(
+            id_attr=db_project.id,
+            name=db_project.name,
+            gradient_start_color=db_project.gradient_start_color,
+            gradient_end_color=db_project.gradient_end_color,
+        )
+
+
+def to_gql_experiment(
+    experiment: models.Experiment,
+    sequence_number: Optional[int] = None,
+) -> Experiment:
+    """
+    Converts an ORM experiment to a GraphQL Experiment.
+    """
+    return Experiment(
+        cached_sequence_number=sequence_number,
+        id_attr=experiment.id,
+        name=experiment.name,
+        project_name=experiment.project_name,
+        description=experiment.description,
+        metadata=experiment.metadata_,
+        created_at=experiment.created_at,
+        updated_at=experiment.updated_at,
+    )

phoenix/server/api/types/ExperimentAnnotationSummary.py

@@ -0,0 +1,13 @@
+from typing import Optional
+
+import strawberry
+
+
+@strawberry.type
+class ExperimentAnnotationSummary:
+    annotation_name: str
+    min_score: Optional[float]
+    max_score: Optional[float]
+    mean_score: Optional[float]
+    count: int
+    error_count: int

phoenix/server/api/types/ExperimentComparison.py

@@ -0,0 +1,19 @@
+from typing import List
+
+import strawberry
+from strawberry.relay import GlobalID
+
+from phoenix.server.api.types.DatasetExample import DatasetExample
+from phoenix.server.api.types.ExperimentRun import ExperimentRun
+
+
+@strawberry.type
+class RunComparisonItem:
+    experiment_id: GlobalID
+    runs: List[ExperimentRun]
+
+
+@strawberry.type
+class ExperimentComparison:
+    example: DatasetExample
+    run_comparison_items: List[RunComparisonItem]