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

phoenix/server/api/types/DocumentAnnotation.py

@@ -1,92 +1,212 @@
 from datetime import datetime
+from math import isfinite
 from typing import TYPE_CHECKING, Annotated, Optional
 
 import strawberry
-from strawberry import Private
 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 phoenix.server.api.interceptor import GqlValueMediator
 
 from .Annotation import Annotation
 from .AnnotationSource import AnnotationSource
 from .AnnotatorKind import AnnotatorKind
-from .User import User, to_gql_user
 
 if TYPE_CHECKING:
     from .Span import Span
+    from .User import User
 
 
 @strawberry.type
 class DocumentAnnotation(Node, Annotation):
-    id_attr: NodeID[int]
-    user_id: Private[Optional[int]]
-    name: str = strawberry.field(
-        description="Name of the annotation, e.g. 'helpfulness' or 'relevance'."
-    )
-    annotator_kind: AnnotatorKind
-    label: Optional[str] = strawberry.field(
+    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."
-    )
-    score: Optional[float] = strawberry.field(
+    )  # 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.",
-        default=GqlValueMediator(),
-    )
-    explanation: Optional[str] = strawberry.field(
+    )  # 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."
-    )
-    metadata: JSON
-    document_position: int
-    span_rowid: Private[int]
-    identifier: str
-    source: AnnotationSource
-    created_at: datetime = strawberry.field(
-        description="The date and time when the annotation was created."
-    )
-    updated_at: datetime = strawberry.field(
-        description="The date and time when the annotation was last updated."
-    )
-
-    @strawberry.field
-    async def span(self) -> Annotated["Span", strawberry.lazy(".Span")]:
-        from phoenix.server.api.types.Span import Span
-
-        return Span(span_rowid=self.span_rowid)
-
-    @strawberry.field
+    )  # 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[User]:
-        if self.user_id is 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
-        user = await info.context.data_loaders.users.load(self.user_id)
-        if user is None:
-            return None
-        return to_gql_user(user)
-
-
-def to_gql_document_annotation(
-    annotation: models.DocumentAnnotation,
-) -> DocumentAnnotation:
-    return DocumentAnnotation(
-        id_attr=annotation.id,
-        user_id=annotation.user_id,
-        name=annotation.name,
-        annotator_kind=AnnotatorKind(annotation.annotator_kind),
-        label=annotation.label,
-        score=annotation.score,
-        explanation=annotation.explanation,
-        metadata=annotation.metadata_,
-        span_rowid=annotation.span_rowid,
-        source=AnnotationSource(annotation.source),
-        identifier=annotation.identifier,
-        document_position=annotation.document_position,
-        created_at=annotation.created_at,
-        updated_at=annotation.updated_at,
-    )
+        from .User import User
+
+        return User(id=user_id)

phoenix/server/api/types/Experiment.py

@@ -1,9 +1,8 @@
 from datetime import datetime
-from typing import ClassVar, Optional
+from typing import TYPE_CHECKING, Annotated, Optional
 
 import strawberry
 from sqlalchemy import func, select
-from sqlalchemy.orm import joinedload
 from strawberry import UNSET, Private
 from strawberry.relay import Connection, GlobalID, Node, NodeID
 from strawberry.scalars import JSON
@@ -18,10 +17,10 @@ from phoenix.server.api.input_types.ExperimentRunSort import (
     get_experiment_run_cursor,
 )
 from phoenix.server.api.types.CostBreakdown import CostBreakdown
-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.ExperimentAnnotationSummary import ExperimentAnnotationSummary
-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.pagination import (
     ConnectionArgs,
     Cursor,
@@ -29,26 +28,128 @@ from phoenix.server.api.types.pagination import (
     connection_from_cursors_and_nodes,
     connection_from_list,
 )
-from phoenix.server.api.types.Project import Project
 from phoenix.server.api.types.SpanCostDetailSummaryEntry import SpanCostDetailSummaryEntry
 from phoenix.server.api.types.SpanCostSummary import SpanCostSummary
 
 _DEFAULT_EXPERIMENT_RUNS_PAGE_SIZE = 50
 
+if TYPE_CHECKING:
+    from .Project import Project
+
 
 @strawberry.type
 class Experiment(Node):
-    _table: ClassVar[type[models.Base]] = models.Experiment
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.Experiment]] = None
     cached_sequence_number: Private[Optional[int]] = None
-    id_attr: NodeID[int]
-    name: str
-    project_name: Optional[str]
-    description: Optional[str]
-    repetitions: int
-    dataset_version_id: GlobalID
-    metadata: JSON
-    created_at: datetime
-    updated_at: datetime
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("Experiment 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.experiment_fields.load(
+                (self.id, models.Experiment.name),
+            )
+        return val
+
+    @strawberry.field
+    async def project_name(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[str]:
+        if self.db_record:
+            val = self.db_record.project_name
+        else:
+            val = await info.context.data_loaders.experiment_fields.load(
+                (self.id, models.Experiment.project_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.experiment_fields.load(
+                (self.id, models.Experiment.description),
+            )
+        return val
+
+    @strawberry.field
+    async def repetitions(
+        self,
+        info: Info[Context, None],
+    ) -> int:
+        if self.db_record:
+            val = self.db_record.repetitions
+        else:
+            val = await info.context.data_loaders.experiment_fields.load(
+                (self.id, models.Experiment.repetitions),
+            )
+        return val
+
+    @strawberry.field
+    async def dataset_version_id(
+        self,
+        info: Info[Context, None],
+    ) -> GlobalID:
+        if self.db_record:
+            version_id = self.db_record.dataset_version_id
+        else:
+            version_id = await info.context.data_loaders.experiment_fields.load(
+                (self.id, models.Experiment.dataset_version_id),
+            )
+        return GlobalID(DatasetVersion.__name__, str(version_id))
+
+    @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.experiment_fields.load(
+                (self.id, models.Experiment.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.experiment_fields.load(
+                (self.id, models.Experiment.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.experiment_fields.load(
+                (self.id, models.Experiment.updated_at),
+            )
+        return val
 
     @strawberry.field(
         description="Sequence number (1-based) of experiments belonging to the same dataset"
@@ -58,9 +159,9 @@ class Experiment(Node):
         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)
+            seq_num = await info.context.data_loaders.experiment_sequence_number.load(self.id)
             if seq_num is None:
-                raise ValueError(f"invalid experiment: id={self.id_attr}")
+                raise ValueError(f"invalid experiment: id={self.id}")
             self.cached_sequence_number = seq_num
         return self.cached_sequence_number
 
@@ -74,12 +175,10 @@ class Experiment(Node):
     ) -> Connection[ExperimentRun]:
         if first is not None and first <= 0:
             raise BadRequest("first must be a positive integer if set")
-        experiment_rowid = self.id_attr
         page_size = first or _DEFAULT_EXPERIMENT_RUNS_PAGE_SIZE
         experiment_runs_query = (
             select(models.ExperimentRun)
-            .where(models.ExperimentRun.experiment_id == experiment_rowid)
-            .options(joinedload(models.ExperimentRun.trace).load_only(models.Trace.trace_id))
+            .where(models.ExperimentRun.experiment_id == self.id)
             .limit(page_size + 1)
         )
 
@@ -94,7 +193,7 @@ class Experiment(Node):
         experiment_runs_query = add_order_by_and_page_start_to_query(
             query=experiment_runs_query,
             sort=sort,
-            experiment_rowid=experiment_rowid,
+            experiment_rowid=self.id,
             after_experiment_run_rowid=after_experiment_run_rowid,
             after_sort_column_value=after_sort_column_value,
         )
@@ -111,7 +210,7 @@ class Experiment(Node):
         for result in results:
             run = result[0]
             annotation_score = result[1] if len(result) > 1 else None
-            gql_run = to_gql_experiment_run(run)
+            gql_run = ExperimentRun(id=run.id, db_record=run)
             cursor = get_experiment_run_cursor(
                 run=run, annotation_score=annotation_score, sort=sort
             )
@@ -125,14 +224,13 @@ class Experiment(Node):
 
     @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)
+        return await info.context.data_loaders.experiment_run_counts.load(self.id)
 
     @strawberry.field
     async def annotation_summaries(
         self, info: Info[Context, None]
     ) -> list[ExperimentAnnotationSummary]:
-        experiment_id = self.id_attr
+        experiment_id = self.id
         return [
             ExperimentAnnotationSummary(
                 annotation_name=summary.annotation_name,
@@ -149,40 +247,42 @@ class Experiment(Node):
 
     @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)
+        return await info.context.data_loaders.experiment_error_rates.load(self.id)
 
     @strawberry.field
     async def average_run_latency_ms(self, info: Info[Context, None]) -> Optional[float]:
-        latency_ms = await info.context.data_loaders.average_experiment_run_latency.load(
-            self.id_attr
-        )
+        latency_ms = await info.context.data_loaders.average_experiment_run_latency.load(self.id)
         return latency_ms
 
     @strawberry.field
-    async def project(self, info: Info[Context, None]) -> Optional[Project]:
-        if self.project_name is None:
+    async def project(
+        self, info: Info[Context, None]
+    ) -> Optional[Annotated["Project", strawberry.lazy(".Project")]]:
+        if self.db_record:
+            project_name = self.db_record.project_name
+        else:
+            project_name = await info.context.data_loaders.experiment_fields.load(
+                (self.id, models.Experiment.project_name),
+            )
+
+        if project_name is None:
             return None
 
-        db_project = await info.context.data_loaders.project_by_name.load(self.project_name)
+        db_project = await info.context.data_loaders.project_by_name.load(project_name)
 
         if db_project is None:
             return None
+        from .Project import Project
 
-        return Project(
-            project_rowid=db_project.id,
-            db_project=db_project,
-        )
+        return Project(id=db_project.id, db_record=db_project)
 
     @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)
+        return info.context.last_updated_at.get(models.Experiment, self.id)
 
     @strawberry.field
     async def cost_summary(self, info: Info[Context, None]) -> SpanCostSummary:
-        experiment_id = self.id_attr
-        summary = await info.context.data_loaders.span_cost_summary_by_experiment.load(
-            experiment_id
-        )
+        summary = await info.context.data_loaders.span_cost_summary_by_experiment.load(self.id)
         return SpanCostSummary(
             prompt=CostBreakdown(
                 tokens=summary.prompt.tokens,
@@ -202,8 +302,6 @@ class Experiment(Node):
     async def cost_detail_summary_entries(
         self, info: Info[Context, None]
     ) -> list[SpanCostDetailSummaryEntry]:
-        experiment_id = self.id_attr
-
         stmt = (
             select(
                 models.SpanCostDetail.token_type,
@@ -216,7 +314,7 @@ class Experiment(Node):
             .join(models.Span, models.SpanCost.span_rowid == models.Span.id)
             .join(models.Trace, models.Span.trace_rowid == models.Trace.id)
             .join(models.ExperimentRun, models.ExperimentRun.trace_id == models.Trace.trace_id)
-            .where(models.ExperimentRun.experiment_id == experiment_id)
+            .where(models.ExperimentRun.experiment_id == self.id)
             .group_by(models.SpanCostDetail.token_type, models.SpanCostDetail.is_prompt)
         )
 
@@ -237,9 +335,9 @@ class Experiment(Node):
         info: Info[Context, None],
     ) -> Connection[DatasetSplit]:
         """Returns the dataset splits associated with this experiment."""
-        splits = await info.context.data_loaders.experiment_dataset_splits.load(self.id_attr)
+        splits = await info.context.data_loaders.experiment_dataset_splits.load(self.id)
         return connection_from_list(
-            [to_gql_dataset_split(split) for split in splits], ConnectionArgs()
+            [DatasetSplit(id=split.id, db_record=split) for split in splits], ConnectionArgs()
         )
 
 
@@ -251,14 +349,7 @@ def to_gql_experiment(
     Converts an ORM experiment to a GraphQL Experiment.
     """
     return Experiment(
+        id=experiment.id,
+        db_record=experiment,
         cached_sequence_number=sequence_number,
-        id_attr=experiment.id,
-        name=experiment.name,
-        project_name=experiment.project_name,
-        description=experiment.description,
-        repetitions=experiment.repetitions,
-        dataset_version_id=GlobalID(DatasetVersion.__name__, str(experiment.dataset_version_id)),
-        metadata=experiment.metadata_,
-        created_at=experiment.created_at,
-        updated_at=experiment.updated_at,
     )

phoenix/server/api/types/ExperimentRepeatedRunGroup.py

@@ -26,7 +26,16 @@ DatasetExampleRowId: TypeAlias = int
 class ExperimentRepeatedRunGroup(Node):
     experiment_rowid: strawberry.Private[ExperimentRowId]
     dataset_example_rowid: strawberry.Private[DatasetExampleRowId]
-    runs: list[ExperimentRun]
+    cached_runs: strawberry.Private[Optional[list[ExperimentRun]]] = None
+
+    @strawberry.field
+    async def runs(self, info: Info[Context, None]) -> list[ExperimentRun]:
+        if self.cached_runs is not None:
+            return self.cached_runs
+        runs = await info.context.data_loaders.experiment_runs_by_experiment_and_example.load(
+            (self.experiment_rowid, self.dataset_example_rowid)
+        )
+        return [ExperimentRun(id=run.id, db_record=run) for run in runs]
 
     @classmethod
     def resolve_id(