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

phoenix/server/api/types/ExperimentRun.py

@@ -2,8 +2,8 @@ from datetime import datetime
 from typing import TYPE_CHECKING, Annotated, Optional
 
 import strawberry
-from sqlalchemy import select
-from sqlalchemy.orm import load_only
+from sqlalchemy import func, select
+from sqlalchemy.sql.functions import coalesce
 from strawberry import UNSET
 from strawberry.relay import Connection, GlobalID, Node, NodeID
 from strawberry.scalars import JSON
@@ -11,30 +11,112 @@ from strawberry.types import Info
 
 from phoenix.db import models
 from phoenix.server.api.context import Context
-from phoenix.server.api.types.ExperimentRunAnnotation import (
-    ExperimentRunAnnotation,
-    to_gql_experiment_run_annotation,
-)
+from phoenix.server.api.types.CostBreakdown import CostBreakdown
+from phoenix.server.api.types.ExperimentRunAnnotation import ExperimentRunAnnotation
 from phoenix.server.api.types.pagination import (
     ConnectionArgs,
     CursorString,
     connection_from_list,
 )
+from phoenix.server.api.types.SpanCostDetailSummaryEntry import SpanCostDetailSummaryEntry
+from phoenix.server.api.types.SpanCostSummary import SpanCostSummary
 from phoenix.server.api.types.Trace import Trace
 
 if TYPE_CHECKING:
-    from phoenix.server.api.types.DatasetExample import DatasetExample
+    from .DatasetExample import DatasetExample
+    from .Trace import Trace
 
 
 @strawberry.type
 class ExperimentRun(Node):
-    id_attr: NodeID[int]
-    experiment_id: GlobalID
-    trace_id: Optional[str]
-    output: Optional[JSON]
-    start_time: datetime
-    end_time: datetime
-    error: Optional[str]
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.ExperimentRun]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("ExperimentRun ID mismatch")
+
+    @strawberry.field
+    async def experiment_id(self, info: Info[Context, None]) -> GlobalID:
+        from .Experiment import Experiment
+
+        if self.db_record:
+            experiment_id = self.db_record.experiment_id
+        else:
+            experiment_id = await info.context.data_loaders.experiment_run_fields.load(
+                (self.id, models.ExperimentRun.experiment_id),
+            )
+        return GlobalID(Experiment.__name__, str(experiment_id))
+
+    @strawberry.field
+    async def repetition_number(self, info: Info[Context, None]) -> int:
+        if self.db_record:
+            val = self.db_record.repetition_number
+        else:
+            val = await info.context.data_loaders.experiment_run_fields.load(
+                (self.id, models.ExperimentRun.repetition_number),
+            )
+        return val
+
+    @strawberry.field
+    async def trace_id(self, info: Info[Context, None]) -> Optional[str]:
+        if self.db_record:
+            val = self.db_record.trace_id
+        else:
+            val = await info.context.data_loaders.experiment_run_fields.load(
+                (self.id, models.ExperimentRun.trace_id),
+            )
+        return val
+
+    @strawberry.field
+    async def output(self, info: Info[Context, None]) -> Optional[JSON]:
+        if self.db_record:
+            output_dict = self.db_record.output
+        else:
+            output_dict = await info.context.data_loaders.experiment_run_fields.load(
+                (self.id, models.ExperimentRun.output),
+            )
+        return output_dict.get("task_output") if output_dict else None
+
+    @strawberry.field
+    async def start_time(self, info: Info[Context, None]) -> datetime:
+        if self.db_record:
+            val = self.db_record.start_time
+        else:
+            val = await info.context.data_loaders.experiment_run_fields.load(
+                (self.id, models.ExperimentRun.start_time),
+            )
+        return val
+
+    @strawberry.field
+    async def end_time(self, info: Info[Context, None]) -> datetime:
+        if self.db_record:
+            val = self.db_record.end_time
+        else:
+            val = await info.context.data_loaders.experiment_run_fields.load(
+                (self.id, models.ExperimentRun.end_time),
+            )
+        return val
+
+    @strawberry.field
+    async def error(self, info: Info[Context, None]) -> Optional[str]:
+        if self.db_record:
+            val = self.db_record.error
+        else:
+            val = await info.context.data_loaders.experiment_run_fields.load(
+                (self.id, models.ExperimentRun.error),
+            )
+        return val
+
+    @strawberry.field
+    async def latency_ms(self, info: Info[Context, None]) -> float:
+        if self.db_record:
+            val = self.db_record.latency_ms
+        else:
+            val = await info.context.data_loaders.experiment_run_fields.load(
+                (self.id, models.ExperimentRun.latency_ms),
+            )
+        return val
 
     @strawberry.field
     async def annotations(
@@ -51,69 +133,91 @@ class ExperimentRun(Node):
             last=last,
             before=before if isinstance(before, CursorString) else None,
         )
-        run_id = self.id_attr
-        annotations = await info.context.data_loaders.experiment_run_annotations.load(run_id)
+        annotations = await info.context.data_loaders.experiment_run_annotations.load(self.id)
         return connection_from_list(
-            [to_gql_experiment_run_annotation(annotation) for annotation in annotations], args
+            [
+                ExperimentRunAnnotation(id=annotation.id, db_record=annotation)
+                for annotation in annotations
+            ],
+            args,
         )
 
     @strawberry.field
-    async def trace(self, info: Info) -> Optional[Trace]:
-        if not self.trace_id:
+    async def trace(
+        self, info: Info[Context, None]
+    ) -> Optional[Annotated["Trace", strawberry.lazy(".Trace")]]:
+        if self.db_record:
+            trace_id = self.db_record.trace_id
+        else:
+            trace_id = await info.context.data_loaders.experiment_run_fields.load(
+                (self.id, models.ExperimentRun.trace_id),
+            )
+        if not trace_id:
             return None
-        dataloader = info.context.data_loaders.trace_by_trace_ids
-        if (trace := await dataloader.load(self.trace_id)) is None:
+        loader = info.context.data_loaders.trace_by_trace_ids
+        if (trace := await loader.load(trace_id)) is None:
             return None
-        return Trace(trace_rowid=trace.id, db_trace=trace)
+        from .Trace import Trace
+
+        return Trace(id=trace.id, db_record=trace)
 
     @strawberry.field
     async def example(
-        self, info: Info
+        self, info: Info[Context, None]
     ) -> Annotated[
-        "DatasetExample", strawberry.lazy("phoenix.server.api.types.DatasetExample")
+        "DatasetExample", strawberry.lazy(".DatasetExample")
     ]:  # use lazy types to avoid circular import: https://strawberry.rocks/docs/types/lazy
-        from phoenix.server.api.types.DatasetExample import DatasetExample
+        from .DatasetExample import DatasetExample
 
-        async with info.context.db() as session:
-            assert (
-                result := await session.execute(
-                    select(models.DatasetExample, models.Experiment.dataset_version_id)
-                    .select_from(models.ExperimentRun)
-                    .join(
-                        models.DatasetExample,
-                        models.DatasetExample.id == models.ExperimentRun.dataset_example_id,
-                    )
-                    .join(
-                        models.Experiment,
-                        models.Experiment.id == models.ExperimentRun.experiment_id,
-                    )
-                    .where(models.ExperimentRun.id == self.id_attr)
-                    .options(load_only(models.DatasetExample.id, models.DatasetExample.created_at))
-                )
-            ) is not None
-            example, version_id = result.first()
-        return DatasetExample(
-            id_attr=example.id,
-            created_at=example.created_at,
-            version_id=version_id,
+        loader = info.context.data_loaders.dataset_examples_and_versions_by_experiment_run
+        (example, version_id) = await loader.load(self.id)
+        return DatasetExample(id=example.id, db_record=example, version_id=version_id)
+
+    @strawberry.field
+    async def cost_summary(self, info: Info[Context, None]) -> SpanCostSummary:
+        summary = await info.context.data_loaders.span_cost_summary_by_experiment_run.load(self.id)
+        return SpanCostSummary(
+            prompt=CostBreakdown(
+                tokens=summary.prompt.tokens,
+                cost=summary.prompt.cost,
+            ),
+            completion=CostBreakdown(
+                tokens=summary.completion.tokens,
+                cost=summary.completion.cost,
+            ),
+            total=CostBreakdown(
+                tokens=summary.total.tokens,
+                cost=summary.total.cost,
+            ),
         )
 
+    @strawberry.field
+    async def cost_detail_summary_entries(
+        self, info: Info[Context, None]
+    ) -> list[SpanCostDetailSummaryEntry]:
+        stmt = (
+            select(
+                models.SpanCostDetail.token_type,
+                models.SpanCostDetail.is_prompt,
+                coalesce(func.sum(models.SpanCostDetail.cost), 0).label("cost"),
+                coalesce(func.sum(models.SpanCostDetail.tokens), 0).label("tokens"),
+            )
+            .select_from(models.SpanCostDetail)
+            .join(models.SpanCost, models.SpanCostDetail.span_cost_id == models.SpanCost.id)
+            .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.id == self.id)
+            .group_by(models.SpanCostDetail.token_type, models.SpanCostDetail.is_prompt)
+        )
 
-def to_gql_experiment_run(run: models.ExperimentRun) -> ExperimentRun:
-    """
-    Converts an ORM experiment run to a GraphQL ExperimentRun.
-    """
-
-    from phoenix.server.api.types.Experiment import Experiment
-
-    return ExperimentRun(
-        id_attr=run.id,
-        experiment_id=GlobalID(Experiment.__name__, str(run.experiment_id)),
-        trace_id=trace_id
-        if (trace := run.trace) and (trace_id := trace.trace_id) is not None
-        else None,
-        output=run.output.get("task_output"),
-        start_time=run.start_time,
-        end_time=run.end_time,
-        error=run.error,
-    )
+        async with info.context.db() as session:
+            data = await session.stream(stmt)
+            return [
+                SpanCostDetailSummaryEntry(
+                    token_type=token_type,
+                    is_prompt=is_prompt,
+                    value=CostBreakdown(tokens=tokens, cost=cost),
+                )
+                async for token_type, is_prompt, cost, tokens in data
+            ]

phoenix/server/api/types/ExperimentRunAnnotation.py

@@ -1,56 +1,175 @@
 from datetime import datetime
+from math import isfinite
 from typing import Optional
 
 import strawberry
 from strawberry import Info
-from strawberry.relay import Node, NodeID
+from strawberry.relay import GlobalID, Node, NodeID
 from strawberry.scalars import JSON
 
 from phoenix.db import models
+from phoenix.server.api.context import Context
 from phoenix.server.api.types.AnnotatorKind import ExperimentRunAnnotatorKind
 from phoenix.server.api.types.Trace import Trace
 
 
 @strawberry.type
 class ExperimentRunAnnotation(Node):
-    id_attr: NodeID[int]
-    name: str
-    annotator_kind: ExperimentRunAnnotatorKind
-    label: Optional[str]
-    score: Optional[float]
-    explanation: Optional[str]
-    error: Optional[str]
-    metadata: JSON
-    start_time: datetime
-    end_time: datetime
-    trace_id: Optional[str]
-
-    @strawberry.field
-    async def trace(self, info: Info) -> Optional[Trace]:
-        if not self.trace_id:
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.ExperimentRunAnnotation]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("ExperimentRunAnnotation 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.experiment_run_annotation_fields.load(
+                (self.id, models.ExperimentRunAnnotation.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],
+    ) -> ExperimentRunAnnotatorKind:
+        if self.db_record:
+            val = self.db_record.annotator_kind
+        else:
+            val = await info.context.data_loaders.experiment_run_annotation_fields.load(
+                (self.id, models.ExperimentRunAnnotation.annotator_kind),
+            )
+        return ExperimentRunAnnotatorKind(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."  # noqa: E501
+    )  # 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.experiment_run_annotation_fields.load(
+                (self.id, models.ExperimentRunAnnotation.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.experiment_run_annotation_fields.load(
+                (self.id, models.ExperimentRunAnnotation.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."  # noqa: E501
+    )  # 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.experiment_run_annotation_fields.load(
+                (self.id, models.ExperimentRunAnnotation.explanation),
+            )
+        return val
+
+    @strawberry.field(description="Error message if the annotation failed to produce a result.")  # type: ignore
+    async def error(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[str]:
+        if self.db_record:
+            val = self.db_record.error
+        else:
+            val = await info.context.data_loaders.experiment_run_annotation_fields.load(
+                (self.id, models.ExperimentRunAnnotation.error),
+            )
+        return val
+
+    @strawberry.field(description="Metadata about 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.experiment_run_annotation_fields.load(
+                (self.id, models.ExperimentRunAnnotation.metadata_),
+            )
+        return val
+
+    @strawberry.field(description="The date and time when the annotation was created.")  # type: ignore
+    async def start_time(
+        self,
+        info: Info[Context, None],
+    ) -> datetime:
+        if self.db_record:
+            val = self.db_record.start_time
+        else:
+            val = await info.context.data_loaders.experiment_run_annotation_fields.load(
+                (self.id, models.ExperimentRunAnnotation.start_time),
+            )
+        return val
+
+    @strawberry.field(description="The date and time when the annotation was last updated.")  # type: ignore
+    async def end_time(
+        self,
+        info: Info[Context, None],
+    ) -> datetime:
+        if self.db_record:
+            val = self.db_record.end_time
+        else:
+            val = await info.context.data_loaders.experiment_run_annotation_fields.load(
+                (self.id, models.ExperimentRunAnnotation.end_time),
+            )
+        return val
+
+    @strawberry.field(description="The identifier of the trace associated with the annotation.")  # type: ignore
+    async def trace_id(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[GlobalID]:
+        if self.db_record:
+            val = self.db_record.trace_id
+        else:
+            val = await info.context.data_loaders.experiment_run_annotation_fields.load(
+                (self.id, models.ExperimentRunAnnotation.trace_id),
+            )
+        return None if val is None else GlobalID(type_name=Trace.__name__, node_id=val)
+
+    @strawberry.field(description="The trace associated with the annotation.")  # type: ignore
+    async def trace(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[Trace]:
+        if self.db_record:
+            trace_id = self.db_record.trace_id
+        else:
+            trace_id = await info.context.data_loaders.experiment_run_annotation_fields.load(
+                (self.id, models.ExperimentRunAnnotation.trace_id),
+            )
+        if not trace_id:
             return None
         dataloader = info.context.data_loaders.trace_by_trace_ids
-        if (trace := await dataloader.load(self.trace_id)) is None:
+        if (trace := await dataloader.load(trace_id)) is None:
             return None
-        return Trace(trace_rowid=trace.id, db_trace=trace)
-
-
-def to_gql_experiment_run_annotation(
-    annotation: models.ExperimentRunAnnotation,
-) -> ExperimentRunAnnotation:
-    """
-    Converts an ORM experiment run annotation to a GraphQL ExperimentRunAnnotation.
-    """
-    return ExperimentRunAnnotation(
-        id_attr=annotation.id,
-        name=annotation.name,
-        annotator_kind=ExperimentRunAnnotatorKind(annotation.annotator_kind),
-        label=annotation.label,
-        score=annotation.score,
-        explanation=annotation.explanation,
-        error=annotation.error,
-        metadata=annotation.metadata_,
-        start_time=annotation.start_time,
-        end_time=annotation.end_time,
-        trace_id=annotation.trace_id,
-    )
+        return Trace(id=trace.id, db_record=trace)