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

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

phoenix/server/api/types/DatasetSplit.py

@@ -0,0 +1,98 @@
+from datetime import datetime
+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):
+    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

phoenix/server/api/types/DocumentAnnotation.py

@@ -0,0 +1,212 @@
+from datetime import datetime
+from math import isfinite
+from typing import TYPE_CHECKING, Annotated, 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
+
+from .Annotation import Annotation
+from .AnnotationSource import AnnotationSource
+from .AnnotatorKind import AnnotatorKind
+
+if TYPE_CHECKING:
+    from .Span import Span
+    from .User import User
+
+
+@strawberry.type
+class DocumentAnnotation(Node, Annotation):
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.DocumentAnnotation]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("DocumentAnnotation ID mismatch")
+
+    @strawberry.field(description="Name of the annotation, e.g. 'helpfulness' or 'relevance'.")  # type: ignore
+    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.document_annotation_fields.load(
+                (self.id, models.DocumentAnnotation.name),
+            )
+        return val
+
+    @strawberry.field(description="The kind of annotator that produced the annotation.")  # type: ignore
+    async def annotator_kind(
+        self,
+        info: Info[Context, None],
+    ) -> AnnotatorKind:
+        if self.db_record:
+            val = self.db_record.annotator_kind
+        else:
+            val = await info.context.data_loaders.document_annotation_fields.load(
+                (self.id, models.DocumentAnnotation.annotator_kind),
+            )
+        return AnnotatorKind(val)
+
+    @strawberry.field(
+        description="Value of the annotation in the form of a string, e.g. "
+        "'helpful' or 'not helpful'. Note that the label is not necessarily binary."
+    )  # type: ignore
+    async def label(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[str]:
+        if self.db_record:
+            val = self.db_record.label
+        else:
+            val = await info.context.data_loaders.document_annotation_fields.load(
+                (self.id, models.DocumentAnnotation.label),
+            )
+        return val
+
+    @strawberry.field(
+        description="Value of the annotation in the form of a numeric score.",
+    )  # type: ignore
+    async def score(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[float]:
+        if self.db_record:
+            val = self.db_record.score
+        else:
+            val = await info.context.data_loaders.document_annotation_fields.load(
+                (self.id, models.DocumentAnnotation.score),
+            )
+        return val if val is not None and isfinite(val) else None
+
+    @strawberry.field(
+        description="The annotator's explanation for the annotation result (i.e. "
+        "score or label, or both) given to the subject."
+    )  # type: ignore
+    async def explanation(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[str]:
+        if self.db_record:
+            val = self.db_record.explanation
+        else:
+            val = await info.context.data_loaders.document_annotation_fields.load(
+                (self.id, models.DocumentAnnotation.explanation),
+            )
+        return val
+
+    @strawberry.field(description="The metadata associated with the annotation.")  # type: ignore
+    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.document_annotation_fields.load(
+                (self.id, models.DocumentAnnotation.metadata_),
+            )
+        return val
+
+    @strawberry.field(description="The position of the annotation in the document.")  # type: ignore
+    async def document_position(
+        self,
+        info: Info[Context, None],
+    ) -> int:
+        if self.db_record:
+            val = self.db_record.document_position
+        else:
+            val = await info.context.data_loaders.document_annotation_fields.load(
+                (self.id, models.DocumentAnnotation.document_position),
+            )
+        return val
+
+    @strawberry.field(description="The identifier of the annotation.")  # type: ignore
+    async def identifier(
+        self,
+        info: Info[Context, None],
+    ) -> str:
+        if self.db_record:
+            val = self.db_record.identifier
+        else:
+            val = await info.context.data_loaders.document_annotation_fields.load(
+                (self.id, models.DocumentAnnotation.identifier),
+            )
+        return val
+
+    @strawberry.field(description="The source of the annotation.")  # type: ignore
+    async def source(
+        self,
+        info: Info[Context, None],
+    ) -> AnnotationSource:
+        if self.db_record:
+            val = self.db_record.source
+        else:
+            val = await info.context.data_loaders.document_annotation_fields.load(
+                (self.id, models.DocumentAnnotation.source),
+            )
+        return AnnotationSource(val)
+
+    @strawberry.field(description="The date and time when the annotation was created.")  # type: ignore
+    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.document_annotation_fields.load(
+                (self.id, models.DocumentAnnotation.created_at),
+            )
+        return val
+
+    @strawberry.field(description="The date and time when the annotation was last updated.")  # type: ignore
+    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.document_annotation_fields.load(
+                (self.id, models.DocumentAnnotation.updated_at),
+            )
+        return val
+
+    @strawberry.field(description="The span associated with the annotation.")  # type: ignore
+    async def span(
+        self,
+        info: Info[Context, None],
+    ) -> Annotated["Span", strawberry.lazy(".Span")]:
+        if self.db_record:
+            span_rowid = self.db_record.span_rowid
+        else:
+            span_rowid = await info.context.data_loaders.document_annotation_fields.load(
+                (self.id, models.DocumentAnnotation.span_rowid),
+            )
+        from .Span import Span
+
+        return Span(id=span_rowid)
+
+    @strawberry.field(description="The user that produced the annotation.")  # type: ignore
+    async def user(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[Annotated["User", strawberry.lazy(".User")]]:
+        if self.db_record:
+            user_id = self.db_record.user_id
+        else:
+            user_id = await info.context.data_loaders.document_annotation_fields.load(
+                (self.id, models.DocumentAnnotation.user_id),
+            )
+        if user_id is None:
+            return None
+        from .User import User
+
+        return User(id=user_id)