arize-phoenix 12.8.0__py3-none-any.whl → 12.9.0__py3-none-any.whl

phoenix/server/api/types/Dataset.py

@@ -1,6 +1,6 @@
 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 Text, and_, func, or_, select
@@ -18,8 +18,8 @@ from phoenix.server.api.types.DatasetExample import DatasetExample
 from phoenix.server.api.types.DatasetExperimentAnnotationSummary import (
     DatasetExperimentAnnotationSummary,
 )
-from phoenix.server.api.types.DatasetLabel import DatasetLabel, to_gql_dataset_label
-from phoenix.server.api.types.DatasetSplit import DatasetSplit, to_gql_dataset_split
+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.node import from_global_id_with_expected_type
@@ -33,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(
@@ -58,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)
@@ -69,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(
@@ -90,7 +146,7 @@ class Dataset(Node):
         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,
@@ -180,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__
@@ -261,9 +317,9 @@ class Dataset(Node):
         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)
             ]
@@ -272,8 +328,8 @@ class Dataset(Node):
     @strawberry.field
     async def splits(self, info: Info[Context, None]) -> list[DatasetSplit]:
         return [
-            to_gql_dataset_split(split)
-            for split in await info.context.data_loaders.dataset_dataset_splits.load(self.id_attr)
+            DatasetSplit(id=split.id, db_record=split)
+            for split in await info.context.data_loaders.dataset_dataset_splits.load(self.id)
         ]
 
     @strawberry.field(
@@ -285,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,
@@ -320,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)
@@ -363,7 +417,7 @@ class Dataset(Node):
     async def experiment_annotation_summaries(
         self, info: Info[Context, None]
     ) -> list[DatasetExperimentAnnotationSummary]:
-        dataset_id = self.id_attr
+        dataset_id = self.id
         query = (
             select(
                 models.ExperimentRunAnnotation.name.label("annotation_name"),
@@ -396,24 +450,10 @@ class Dataset(Node):
     @strawberry.field
     async def labels(self, info: Info[Context, None]) -> list[DatasetLabel]:
         return [
-            to_gql_dataset_label(label)
-            for label in await info.context.data_loaders.dataset_labels.load(self.id_attr)
+            DatasetLabel(id=label.id, db_record=label)
+            for label in await info.context.data_loaders.dataset_labels.load(self.id)
         ]
 
     @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)
-
-
-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,
-    )
+        return info.context.last_updated_at.get(models.Dataset, self.id)

phoenix/server/api/types/DatasetExample.py

@@ -1,9 +1,8 @@
 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
@@ -12,34 +11,49 @@ 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, to_gql_dataset_split
+from phoenix.server.api.types.DatasetSplit import DatasetSplit
 from phoenix.server.api.types.DatasetVersion import DatasetVersion
 from phoenix.server.api.types.ExperimentRepeatedRunGroup import (
     ExperimentRepeatedRunGroup,
 )
-from phoenix.server.api.types.ExperimentRun import ExperimentRun, to_gql_experiment_run
+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(
@@ -47,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
         )
 
@@ -78,12 +92,10 @@ 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)
+            .where(models.ExperimentRun.dataset_example_id == self.id)
             .order_by(
                 models.ExperimentRun.experiment_id.asc(),
                 models.ExperimentRun.repetition_number.asc(),
@@ -100,7 +112,7 @@ class DatasetExample(Node):
             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(
@@ -108,7 +120,6 @@ class DatasetExample(Node):
         info: Info[Context, None],
         experiment_ids: list[GlobalID],
     ) -> list[ExperimentRepeatedRunGroup]:
-        example_rowid = self.id_attr
         experiment_rowids = []
         for experiment_id in experiment_ids:
             try:
@@ -121,14 +132,14 @@ class DatasetExample(Node):
             experiment_rowids.append(experiment_rowid)
         repeated_run_groups = (
             await info.context.data_loaders.experiment_repeated_run_groups.load_many(
-                [(experiment_rowid, example_rowid) for experiment_rowid in experiment_rowids]
+                [(experiment_rowid, self.id) for experiment_rowid in experiment_rowids]
             )
         )
         return [
             ExperimentRepeatedRunGroup(
                 experiment_rowid=group.experiment_rowid,
                 dataset_example_rowid=group.dataset_example_rowid,
-                runs=[to_gql_experiment_run(run) for run in group.runs],
+                cached_runs=[ExperimentRun(id=run.id, db_record=run) for run in group.runs],
             )
             for group in repeated_run_groups
         ]
@@ -139,13 +150,6 @@ class DatasetExample(Node):
         info: Info[Context, None],
     ) -> list[DatasetSplit]:
         return [
-            to_gql_dataset_split(split)
-            for split in await info.context.data_loaders.dataset_example_splits.load(self.id_attr)
+            DatasetSplit(id=split.id, db_record=split)
+            for split in await info.context.data_loaders.dataset_example_splits.load(self.id)
         ]
-
-
-def to_gql_dataset_example(example: models.DatasetExample) -> DatasetExample:
-    return DatasetExample(
-        id_attr=example.id,
-        created_at=example.created_at,
-    )

phoenix/server/api/types/DatasetLabel.py

@@ -2,22 +2,56 @@ 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_attr: NodeID[int]
-    name: str
-    description: Optional[str]
-    color: str
-
-
-def to_gql_dataset_label(dataset_label: models.DatasetLabel) -> DatasetLabel:
-    return DatasetLabel(
-        id_attr=dataset_label.id,
-        name=dataset_label.name,
-        description=dataset_label.description,
-        color=dataset_label.color,
-    )
+    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

phoenix/server/api/types/DatasetSplit.py

@@ -1,32 +1,98 @@
 from datetime import datetime
-from typing import ClassVar, Optional
+from typing import Optional
 
 import strawberry
 from strawberry.relay import Node, NodeID
 from strawberry.scalars import JSON
+from strawberry.types import Info
 
 from phoenix.db import models
+from phoenix.server.api.context import Context
 
 
 @strawberry.type
 class DatasetSplit(Node):
-    _table: ClassVar[type[models.Base]] = models.DatasetSplit
-    id_attr: NodeID[int]
-    name: str
-    description: Optional[str]
-    metadata: JSON
-    color: str
-    created_at: datetime
-    updated_at: datetime
-
-
-def to_gql_dataset_split(dataset_split: models.DatasetSplit) -> DatasetSplit:
-    return DatasetSplit(
-        id_attr=dataset_split.id,
-        name=dataset_split.name,
-        description=dataset_split.description,
-        color=dataset_split.color or "#ffffff",
-        metadata=dataset_split.metadata_,
-        created_at=dataset_split.created_at,
-        updated_at=dataset_split.updated_at,
-    )
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.DatasetSplit]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("DatasetSplit 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_split_fields.load(
+                (self.id, models.DatasetSplit.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_split_fields.load(
+                (self.id, models.DatasetSplit.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_split_fields.load(
+                (self.id, models.DatasetSplit.metadata_),
+            )
+        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_split_fields.load(
+                (self.id, models.DatasetSplit.color),
+            )
+        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_split_fields.load(
+                (self.id, models.DatasetSplit.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_split_fields.load(
+                (self.id, models.DatasetSplit.updated_at),
+            )
+        return val

phoenix/server/api/types/DatasetVersion.py

@@ -4,11 +4,56 @@ from typing import Optional
 import strawberry
 from strawberry.relay import Node, NodeID
 from strawberry.scalars import JSON
+from strawberry.types import Info
+
+from phoenix.db import models
+from phoenix.server.api.context import Context
 
 
 @strawberry.type
 class DatasetVersion(Node):
-    id_attr: NodeID[int]
-    description: Optional[str]
-    metadata: JSON
-    created_at: datetime
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.DatasetVersion]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("DatasetVersion ID mismatch")
+
+    @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_version_fields.load(
+                (self.id, models.DatasetVersion.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_version_fields.load(
+                (self.id, models.DatasetVersion.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_version_fields.load(
+                (self.id, models.DatasetVersion.created_at),
+            )
+        return val