arize-phoenix 11.23.1__py3-none-any.whl → 12.28.1__py3-none-any.whl

phoenix/server/api/types/Dataset.py

@@ -1,9 +1,9 @@
 from collections.abc import AsyncIterable
 from datetime import datetime
-from typing import ClassVar, Optional, cast
+from typing import Optional, cast
 
 import strawberry
-from sqlalchemy import and_, func, or_, select
+from sqlalchemy import Text, and_, func, or_, select
 from sqlalchemy.sql.functions import count
 from strawberry import UNSET
 from strawberry.relay import Connection, GlobalID, Node, NodeID
@@ -15,9 +15,13 @@ from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest
 from phoenix.server.api.input_types.DatasetVersionSort import DatasetVersionSort
 from phoenix.server.api.types.DatasetExample import DatasetExample
+from phoenix.server.api.types.DatasetExperimentAnnotationSummary import (
+    DatasetExperimentAnnotationSummary,
+)
+from phoenix.server.api.types.DatasetLabel import DatasetLabel
+from phoenix.server.api.types.DatasetSplit import DatasetSplit
 from phoenix.server.api.types.DatasetVersion import DatasetVersion
 from phoenix.server.api.types.Experiment import Experiment, to_gql_experiment
-from phoenix.server.api.types.ExperimentAnnotationSummary import ExperimentAnnotationSummary
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.pagination import (
     ConnectionArgs,
@@ -29,13 +33,77 @@ from phoenix.server.api.types.SortDir import SortDir
 
 @strawberry.type
 class Dataset(Node):
-    _table: ClassVar[type[models.Base]] = models.Experiment
-    id_attr: NodeID[int]
-    name: str
-    description: Optional[str]
-    metadata: JSON
-    created_at: datetime
-    updated_at: datetime
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.Dataset]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("Dataset ID mismatch")
+
+    @strawberry.field
+    async def name(
+        self,
+        info: Info[Context, None],
+    ) -> str:
+        if self.db_record:
+            val = self.db_record.name
+        else:
+            val = await info.context.data_loaders.dataset_fields.load(
+                (self.id, models.Dataset.name),
+            )
+        return val
+
+    @strawberry.field
+    async def description(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[str]:
+        if self.db_record:
+            val = self.db_record.description
+        else:
+            val = await info.context.data_loaders.dataset_fields.load(
+                (self.id, models.Dataset.description),
+            )
+        return val
+
+    @strawberry.field
+    async def metadata(
+        self,
+        info: Info[Context, None],
+    ) -> JSON:
+        if self.db_record:
+            val = self.db_record.metadata_
+        else:
+            val = await info.context.data_loaders.dataset_fields.load(
+                (self.id, models.Dataset.metadata_),
+            )
+        return val
+
+    @strawberry.field
+    async def created_at(
+        self,
+        info: Info[Context, None],
+    ) -> datetime:
+        if self.db_record:
+            val = self.db_record.created_at
+        else:
+            val = await info.context.data_loaders.dataset_fields.load(
+                (self.id, models.Dataset.created_at),
+            )
+        return val
+
+    @strawberry.field
+    async def updated_at(
+        self,
+        info: Info[Context, None],
+    ) -> datetime:
+        if self.db_record:
+            val = self.db_record.updated_at
+        else:
+            val = await info.context.data_loaders.dataset_fields.load(
+                (self.id, models.Dataset.updated_at),
+            )
+        return val
 
     @strawberry.field
     async def versions(
@@ -54,7 +122,7 @@ class Dataset(Node):
             before=before if isinstance(before, CursorString) else None,
         )
         async with info.context.db() as session:
-            stmt = select(models.DatasetVersion).filter_by(dataset_id=self.id_attr)
+            stmt = select(models.DatasetVersion).filter_by(dataset_id=self.id)
             if sort:
                 # For now assume the the column names match 1:1 with the enum values
                 sort_col = getattr(models.DatasetVersion, sort.col.value)
@@ -65,15 +133,7 @@ class Dataset(Node):
             else:
                 stmt = stmt.order_by(models.DatasetVersion.created_at.desc())
             versions = await session.scalars(stmt)
-        data = [
-            DatasetVersion(
-                id_attr=version.id,
-                description=version.description,
-                metadata=version.metadata_,
-                created_at=version.created_at,
-            )
-            for version in versions
-        ]
+        data = [DatasetVersion(id=version.id, db_record=version) for version in versions]
         return connection_from_list(data=data, args=args)
 
     @strawberry.field(
@@ -84,8 +144,9 @@ class Dataset(Node):
         self,
         info: Info[Context, None],
         dataset_version_id: Optional[GlobalID] = UNSET,
+        split_ids: Optional[list[GlobalID]] = UNSET,
     ) -> int:
-        dataset_id = self.id_attr
+        dataset_id = self.id
         version_id = (
             from_global_id_with_expected_type(
                 global_id=dataset_version_id,
@@ -94,6 +155,20 @@ class Dataset(Node):
             if dataset_version_id
             else None
         )
+
+        # Parse split IDs if provided
+        split_rowids: Optional[list[int]] = None
+        if split_ids:
+            split_rowids = []
+            for split_id in split_ids:
+                try:
+                    split_rowid = from_global_id_with_expected_type(
+                        global_id=split_id, expected_type_name=models.DatasetSplit.__name__
+                    )
+                    split_rowids.append(split_rowid)
+                except Exception:
+                    raise BadRequest(f"Invalid split ID: {split_id}")
+
         revision_ids = (
             select(func.max(models.DatasetExampleRevision.id))
             .join(models.DatasetExample)
@@ -110,11 +185,36 @@ class Dataset(Node):
             revision_ids = revision_ids.where(
                 models.DatasetExampleRevision.dataset_version_id <= version_id_subquery
             )
-        stmt = (
-            select(count(models.DatasetExampleRevision.id))
-            .where(models.DatasetExampleRevision.id.in_(revision_ids))
-            .where(models.DatasetExampleRevision.revision_kind != "DELETE")
-        )
+
+        # Build the count query
+        if split_rowids:
+            # When filtering by splits, count distinct examples that belong to those splits
+            stmt = (
+                select(count(models.DatasetExample.id.distinct()))
+                .join(
+                    models.DatasetExampleRevision,
+                    onclause=(
+                        models.DatasetExample.id == models.DatasetExampleRevision.dataset_example_id
+                    ),
+                )
+                .join(
+                    models.DatasetSplitDatasetExample,
+                    onclause=(
+                        models.DatasetExample.id
+                        == models.DatasetSplitDatasetExample.dataset_example_id
+                    ),
+                )
+                .where(models.DatasetExampleRevision.id.in_(revision_ids))
+                .where(models.DatasetExampleRevision.revision_kind != "DELETE")
+                .where(models.DatasetSplitDatasetExample.dataset_split_id.in_(split_rowids))
+            )
+        else:
+            stmt = (
+                select(count(models.DatasetExampleRevision.id))
+                .where(models.DatasetExampleRevision.id.in_(revision_ids))
+                .where(models.DatasetExampleRevision.revision_kind != "DELETE")
+            )
+
         async with info.context.db() as session:
             return (await session.scalar(stmt)) or 0
 
@@ -123,10 +223,12 @@ class Dataset(Node):
         self,
         info: Info[Context, None],
         dataset_version_id: Optional[GlobalID] = UNSET,
+        split_ids: Optional[list[GlobalID]] = UNSET,
         first: Optional[int] = 50,
         last: Optional[int] = UNSET,
         after: Optional[CursorString] = UNSET,
         before: Optional[CursorString] = UNSET,
+        filter: Optional[str] = UNSET,
     ) -> Connection[DatasetExample]:
         args = ConnectionArgs(
             first=first,
@@ -134,7 +236,7 @@ class Dataset(Node):
             last=last,
             before=before if isinstance(before, CursorString) else None,
         )
-        dataset_id = self.id_attr
+        dataset_id = self.id
         version_id = (
             from_global_id_with_expected_type(
                 global_id=dataset_version_id, expected_type_name=DatasetVersion.__name__
@@ -142,6 +244,20 @@ class Dataset(Node):
             if dataset_version_id
             else None
         )
+
+        # Parse split IDs if provided
+        split_rowids: Optional[list[int]] = None
+        if split_ids:
+            split_rowids = []
+            for split_id in split_ids:
+                try:
+                    split_rowid = from_global_id_with_expected_type(
+                        global_id=split_id, expected_type_name=models.DatasetSplit.__name__
+                    )
+                    split_rowids.append(split_rowid)
+                except Exception:
+                    raise BadRequest(f"Invalid split ID: {split_id}")
+
         revision_ids = (
             select(func.max(models.DatasetExampleRevision.id))
             .join(models.DatasetExample)
@@ -171,19 +287,51 @@ class Dataset(Node):
                     models.DatasetExampleRevision.revision_kind != "DELETE",
                 )
             )
-            .order_by(models.DatasetExampleRevision.dataset_example_id.desc())
+            .order_by(models.DatasetExample.id.desc())
         )
+
+        # Filter by split IDs if provided
+        if split_rowids:
+            query = (
+                query.join(
+                    models.DatasetSplitDatasetExample,
+                    onclause=(
+                        models.DatasetExample.id
+                        == models.DatasetSplitDatasetExample.dataset_example_id
+                    ),
+                )
+                .where(models.DatasetSplitDatasetExample.dataset_split_id.in_(split_rowids))
+                .distinct()
+            )
+        # Apply filter if provided - search through JSON fields (input, output, metadata)
+        if filter is not UNSET and filter:
+            # Create a filter that searches for the filter string in JSON fields
+            # Using PostgreSQL's JSON operators for case-insensitive text search
+            filter_condition = or_(
+                func.cast(models.DatasetExampleRevision.input, Text).ilike(f"%{filter}%"),
+                func.cast(models.DatasetExampleRevision.output, Text).ilike(f"%{filter}%"),
+                func.cast(models.DatasetExampleRevision.metadata_, Text).ilike(f"%{filter}%"),
+            )
+            query = query.where(filter_condition)
+
         async with info.context.db() as session:
             dataset_examples = [
                 DatasetExample(
-                    id_attr=example.id,
+                    id=example.id,
+                    db_record=example,
                     version_id=version_id,
-                    created_at=example.created_at,
                 )
                 async for example in await session.stream_scalars(query)
             ]
         return connection_from_list(data=dataset_examples, args=args)
 
+    @strawberry.field
+    async def splits(self, info: Info[Context, None]) -> list[DatasetSplit]:
+        return [
+            DatasetSplit(id=split.id, db_record=split)
+            for split in await info.context.data_loaders.dataset_dataset_splits.load(self.id)
+        ]
+
     @strawberry.field(
         description="Number of experiments for a specific version if version is specified, "
         "or for all versions if version is not specified."
@@ -193,9 +341,7 @@ class Dataset(Node):
         info: Info[Context, None],
         dataset_version_id: Optional[GlobalID] = UNSET,
     ) -> int:
-        stmt = select(count(models.Experiment.id)).where(
-            models.Experiment.dataset_id == self.id_attr
-        )
+        stmt = select(count(models.Experiment.id)).where(models.Experiment.dataset_id == self.id)
         version_id = (
             from_global_id_with_expected_type(
                 global_id=dataset_version_id,
@@ -228,7 +374,7 @@ class Dataset(Node):
             last=last,
             before=before if isinstance(before, CursorString) else None,
         )
-        dataset_id = self.id_attr
+        dataset_id = self.id
         row_number = func.row_number().over(order_by=models.Experiment.id).label("row_number")
         query = (
             select(models.Experiment, row_number)
@@ -270,17 +416,15 @@ class Dataset(Node):
     @strawberry.field
     async def experiment_annotation_summaries(
         self, info: Info[Context, None]
-    ) -> list[ExperimentAnnotationSummary]:
-        dataset_id = self.id_attr
+    ) -> list[DatasetExperimentAnnotationSummary]:
+        dataset_id = self.id
         query = (
             select(
-                models.ExperimentRunAnnotation.name,
-                func.min(models.ExperimentRunAnnotation.score),
-                func.max(models.ExperimentRunAnnotation.score),
-                func.avg(models.ExperimentRunAnnotation.score),
-                func.count(),
-                func.count(models.ExperimentRunAnnotation.error),
+                models.ExperimentRunAnnotation.name.label("annotation_name"),
+                func.min(models.ExperimentRunAnnotation.score).label("min_score"),
+                func.max(models.ExperimentRunAnnotation.score).label("max_score"),
             )
+            .select_from(models.ExperimentRunAnnotation)
             .join(
                 models.ExperimentRun,
                 models.ExperimentRunAnnotation.experiment_run_id == models.ExperimentRun.id,
@@ -295,38 +439,21 @@ class Dataset(Node):
         )
         async with info.context.db() as session:
             return [
-                ExperimentAnnotationSummary(
-                    annotation_name=annotation_name,
-                    min_score=min_score,
-                    max_score=max_score,
-                    mean_score=mean_score,
-                    count=count_,
-                    error_count=error_count,
+                DatasetExperimentAnnotationSummary(
+                    annotation_name=scores_tuple.annotation_name,
+                    min_score=scores_tuple.min_score,
+                    max_score=scores_tuple.max_score,
                 )
-                async for (
-                    annotation_name,
-                    min_score,
-                    max_score,
-                    mean_score,
-                    count_,
-                    error_count,
-                ) in await session.stream(query)
+                async for scores_tuple in await session.stream(query)
             ]
 
     @strawberry.field
-    def last_updated_at(self, info: Info[Context, None]) -> Optional[datetime]:
-        return info.context.last_updated_at.get(self._table, self.id_attr)
-
+    async def labels(self, info: Info[Context, None]) -> list[DatasetLabel]:
+        return [
+            DatasetLabel(id=label.id, db_record=label)
+            for label in await info.context.data_loaders.dataset_labels.load(self.id)
+        ]
 
-def to_gql_dataset(dataset: models.Dataset) -> Dataset:
-    """
-    Converts an ORM dataset to a GraphQL dataset.
-    """
-    return Dataset(
-        id_attr=dataset.id,
-        name=dataset.name,
-        description=dataset.description,
-        metadata=dataset.metadata_,
-        created_at=dataset.created_at,
-        updated_at=dataset.updated_at,
-    )
+    @strawberry.field
+    def last_updated_at(self, info: Info[Context, None]) -> Optional[datetime]:
+        return info.context.last_updated_at.get(models.Dataset, self.id)

phoenix/server/api/types/DatasetExample.py

@@ -1,40 +1,59 @@
 from datetime import datetime
-from typing import Optional
+from typing import TYPE_CHECKING, Annotated, Optional
 
 import strawberry
 from sqlalchemy import select
-from sqlalchemy.orm import joinedload
 from strawberry import UNSET
 from strawberry.relay.types import Connection, GlobalID, Node, NodeID
 from strawberry.types import Info
 
 from phoenix.db import models
 from phoenix.server.api.context import Context
+from phoenix.server.api.exceptions import BadRequest
 from phoenix.server.api.types.DatasetExampleRevision import DatasetExampleRevision
+from phoenix.server.api.types.DatasetSplit import DatasetSplit
 from phoenix.server.api.types.DatasetVersion import DatasetVersion
-from phoenix.server.api.types.ExperimentRun import ExperimentRun, to_gql_experiment_run
+from phoenix.server.api.types.ExperimentRepeatedRunGroup import (
+    ExperimentRepeatedRunGroup,
+)
+from phoenix.server.api.types.ExperimentRun import ExperimentRun
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.pagination import (
     ConnectionArgs,
     CursorString,
     connection_from_list,
 )
-from phoenix.server.api.types.Span import Span
+
+if TYPE_CHECKING:
+    from .Span import Span
 
 
 @strawberry.type
 class DatasetExample(Node):
-    id_attr: NodeID[int]
-    created_at: datetime
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.DatasetExample]] = None
     version_id: strawberry.Private[Optional[int]] = None
 
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("DatasetExample ID mismatch")
+
+    @strawberry.field
+    async def created_at(self, info: Info[Context, None]) -> datetime:
+        if self.db_record:
+            val = self.db_record.created_at
+        else:
+            val = await info.context.data_loaders.dataset_example_fields.load(
+                (self.id, models.DatasetExample.created_at),
+            )
+        return val
+
     @strawberry.field
     async def revision(
         self,
         info: Info[Context, None],
         dataset_version_id: Optional[GlobalID] = UNSET,
     ) -> DatasetExampleRevision:
-        example_id = self.id_attr
         version_id: Optional[int] = None
         if dataset_version_id:
             version_id = from_global_id_with_expected_type(
@@ -42,18 +61,18 @@ class DatasetExample(Node):
             )
         elif self.version_id is not None:
             version_id = self.version_id
-        return await info.context.data_loaders.dataset_example_revisions.load(
-            (example_id, version_id)
-        )
+        return await info.context.data_loaders.dataset_example_revisions.load((self.id, version_id))
 
     @strawberry.field
     async def span(
         self,
         info: Info[Context, None],
-    ) -> Optional[Span]:
+    ) -> Optional[Annotated["Span", strawberry.lazy(".Span")]]:
+        from .Span import Span
+
         return (
-            Span(span_rowid=span.id, db_span=span)
-            if (span := await info.context.data_loaders.dataset_example_spans.load(self.id_attr))
+            Span(id=span.id, db_record=span)
+            if (span := await info.context.data_loaders.dataset_example_spans.load(self.id))
             else None
         )
 
@@ -65,6 +84,7 @@ class DatasetExample(Node):
         last: Optional[int] = UNSET,
         after: Optional[CursorString] = UNSET,
         before: Optional[CursorString] = UNSET,
+        experiment_ids: Optional[list[GlobalID]] = UNSET,
     ) -> Connection[ExperimentRun]:
         args = ConnectionArgs(
             first=first,
@@ -72,14 +92,64 @@ class DatasetExample(Node):
             last=last,
             before=before if isinstance(before, CursorString) else None,
         )
-        example_id = self.id_attr
         query = (
             select(models.ExperimentRun)
-            .options(joinedload(models.ExperimentRun.trace).load_only(models.Trace.trace_id))
             .join(models.Experiment, models.Experiment.id == models.ExperimentRun.experiment_id)
-            .where(models.ExperimentRun.dataset_example_id == example_id)
-            .order_by(models.Experiment.id.desc())
+            .where(models.ExperimentRun.dataset_example_id == self.id)
+            .order_by(
+                models.ExperimentRun.experiment_id.asc(),
+                models.ExperimentRun.repetition_number.asc(),
+            )
         )
+        if experiment_ids:
+            experiment_db_ids = [
+                from_global_id_with_expected_type(
+                    global_id=experiment_id,
+                    expected_type_name=models.Experiment.__name__,
+                )
+                for experiment_id in experiment_ids or []
+            ]
+            query = query.where(models.ExperimentRun.experiment_id.in_(experiment_db_ids))
         async with info.context.db() as session:
             runs = (await session.scalars(query)).all()
-        return connection_from_list([to_gql_experiment_run(run) for run in runs], args)
+        return connection_from_list([ExperimentRun(id=run.id, db_record=run) for run in runs], args)
+
+    @strawberry.field
+    async def experiment_repeated_run_groups(
+        self,
+        info: Info[Context, None],
+        experiment_ids: list[GlobalID],
+    ) -> list[ExperimentRepeatedRunGroup]:
+        experiment_rowids = []
+        for experiment_id in experiment_ids:
+            try:
+                experiment_rowid = from_global_id_with_expected_type(
+                    global_id=experiment_id,
+                    expected_type_name=models.Experiment.__name__,
+                )
+            except Exception:
+                raise BadRequest(f"Invalid experiment ID: {experiment_id}")
+            experiment_rowids.append(experiment_rowid)
+        repeated_run_groups = (
+            await info.context.data_loaders.experiment_repeated_run_groups.load_many(
+                [(experiment_rowid, self.id) for experiment_rowid in experiment_rowids]
+            )
+        )
+        return [
+            ExperimentRepeatedRunGroup(
+                experiment_rowid=group.experiment_rowid,
+                dataset_example_rowid=group.dataset_example_rowid,
+                cached_runs=[ExperimentRun(id=run.id, db_record=run) for run in group.runs],
+            )
+            for group in repeated_run_groups
+        ]
+
+    @strawberry.field
+    async def dataset_splits(
+        self,
+        info: Info[Context, None],
+    ) -> list[DatasetSplit]:
+        return [
+            DatasetSplit(id=split.id, db_record=split)
+            for split in await info.context.data_loaders.dataset_example_splits.load(self.id)
+        ]

phoenix/server/api/types/DatasetExperimentAnnotationSummary.py

@@ -0,0 +1,10 @@
+from typing import Optional
+
+import strawberry
+
+
+@strawberry.type
+class DatasetExperimentAnnotationSummary:
+    annotation_name: str
+    min_score: Optional[float]
+    max_score: Optional[float]

phoenix/server/api/types/DatasetLabel.py

@@ -0,0 +1,57 @@
+from typing import Optional
+
+import strawberry
+from strawberry.relay import Node, NodeID
+from strawberry.types import Info
+
+from phoenix.db import models
+from phoenix.server.api.context import Context
+
+
+@strawberry.type
+class DatasetLabel(Node):
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.DatasetLabel]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("DatasetLabel ID mismatch")
+
+    @strawberry.field
+    async def name(
+        self,
+        info: Info[Context, None],
+    ) -> str:
+        if self.db_record:
+            val = self.db_record.name
+        else:
+            val = await info.context.data_loaders.dataset_label_fields.load(
+                (self.id, models.DatasetLabel.name),
+            )
+        return val
+
+    @strawberry.field
+    async def description(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[str]:
+        if self.db_record:
+            val = self.db_record.description
+        else:
+            val = await info.context.data_loaders.dataset_label_fields.load(
+                (self.id, models.DatasetLabel.description),
+            )
+        return val
+
+    @strawberry.field
+    async def color(
+        self,
+        info: Info[Context, None],
+    ) -> str:
+        if self.db_record:
+            val = self.db_record.color
+        else:
+            val = await info.context.data_loaders.dataset_label_fields.load(
+                (self.id, models.DatasetLabel.color),
+            )
+        return val