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

phoenix/server/api/types/ProjectSession.py

@@ -1,58 +1,116 @@
+from collections import defaultdict
+from dataclasses import asdict, dataclass
 from datetime import datetime
-from typing import TYPE_CHECKING, Annotated, ClassVar, Optional, Type
+from typing import TYPE_CHECKING, Annotated, Optional
 
+import pandas as pd
 import strawberry
 from openinference.semconv.trace import SpanAttributes
 from sqlalchemy import select
-from strawberry import UNSET, Info, Private, lazy
-from strawberry.relay import Connection, GlobalID, Node, NodeID
+from strawberry import UNSET, Info, lazy
+from strawberry.relay import Connection, Node, NodeID
 
 from phoenix.db import models
 from phoenix.server.api.context import Context
+from phoenix.server.api.input_types.AnnotationFilter import AnnotationFilter, satisfies_filter
+from phoenix.server.api.types.AnnotationSummary import AnnotationSummary
+from phoenix.server.api.types.CostBreakdown import CostBreakdown
 from phoenix.server.api.types.MimeType import MimeType
 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.SpanIOValue import SpanIOValue
 from phoenix.server.api.types.TokenUsage import TokenUsage
 
 if TYPE_CHECKING:
+    from phoenix.server.api.types.Project import Project
+    from phoenix.server.api.types.ProjectSessionAnnotation import ProjectSessionAnnotation
     from phoenix.server.api.types.Trace import Trace
 
 
 @strawberry.type
 class ProjectSession(Node):
-    _table: ClassVar[Type[models.ProjectSession]] = models.ProjectSession
-    id_attr: NodeID[int]
-    project_rowid: Private[int]
-    session_id: str
-    start_time: datetime
-    end_time: datetime
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.ProjectSession]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("ProjectSession ID mismatch")
+
+    @strawberry.field
+    async def session_id(
+        self,
+        info: Info[Context, None],
+    ) -> str:
+        if self.db_record:
+            val = self.db_record.session_id
+        else:
+            val = await info.context.data_loaders.project_session_fields.load(
+                (self.id, models.ProjectSession.session_id),
+            )
+        return val
+
+    @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.project_session_fields.load(
+                (self.id, models.ProjectSession.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.project_session_fields.load(
+                (self.id, models.ProjectSession.end_time),
+            )
+        return val
 
     @strawberry.field
-    async def project_id(self) -> GlobalID:
+    async def project(
+        self,
+        info: Info[Context, None],
+    ) -> Annotated["Project", lazy(".Project")]:
         from phoenix.server.api.types.Project import Project
 
-        return GlobalID(type_name=Project.__name__, node_id=str(self.project_rowid))
+        if self.db_record:
+            project_rowid = self.db_record.project_id
+        else:
+            project_rowid = await info.context.data_loaders.project_session_fields.load(
+                (self.id, models.ProjectSession.project_id),
+            )
+        return Project(id=project_rowid)
 
     @strawberry.field
     async def num_traces(
         self,
         info: Info[Context, None],
     ) -> int:
-        return await info.context.data_loaders.session_num_traces.load(self.id_attr)
+        return await info.context.data_loaders.session_num_traces.load(self.id)
 
     @strawberry.field
     async def num_traces_with_error(
         self,
         info: Info[Context, None],
     ) -> int:
-        return await info.context.data_loaders.session_num_traces_with_error.load(self.id_attr)
+        return await info.context.data_loaders.session_num_traces_with_error.load(self.id)
 
     @strawberry.field
     async def first_input(
         self,
         info: Info[Context, None],
     ) -> Optional[SpanIOValue]:
-        record = await info.context.data_loaders.session_first_inputs.load(self.id_attr)
+        record = await info.context.data_loaders.session_first_inputs.load(self.id)
         if record is None:
             return None
         return SpanIOValue(
@@ -65,7 +123,7 @@ class ProjectSession(Node):
         self,
         info: Info[Context, None],
     ) -> Optional[SpanIOValue]:
-        record = await info.context.data_loaders.session_last_outputs.load(self.id_attr)
+        record = await info.context.data_loaders.session_last_outputs.load(self.id)
         if record is None:
             return None
         return SpanIOValue(
@@ -78,7 +136,7 @@ class ProjectSession(Node):
         self,
         info: Info[Context, None],
     ) -> TokenUsage:
-        usage = await info.context.data_loaders.session_token_usages.load(self.id_attr)
+        usage = await info.context.data_loaders.session_token_usages.load(self.id)
         return TokenUsage(
             prompt=usage.prompt,
             completion=usage.completion,
@@ -103,13 +161,12 @@ class ProjectSession(Node):
         )
         stmt = (
             select(models.Trace)
-            .filter_by(project_session_rowid=self.id_attr)
+            .filter_by(project_session_rowid=self.id)
             .order_by(models.Trace.start_time)
-            .limit(first)
         )
         async with info.context.db() as session:
             traces = await session.stream_scalars(stmt)
-            data = [Trace(trace_rowid=trace.id, db_trace=trace) async for trace in traces]
+            data = [Trace(id=trace.id, db_record=trace) async for trace in traces]
         return connection_from_list(data=data, args=args)
 
     @strawberry.field
@@ -119,18 +176,121 @@ class ProjectSession(Node):
         probability: float,
     ) -> Optional[float]:
         return await info.context.data_loaders.session_trace_latency_ms_quantile.load(
-            (self.id_attr, probability)
+            (self.id, probability)
+        )
+
+    @strawberry.field
+    async def cost_summary(
+        self,
+        info: Info[Context, None],
+    ) -> SpanCostSummary:
+        loader = info.context.data_loaders.span_cost_summary_by_project_session
+        summary = await loader.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]:
+        loader = info.context.data_loaders.span_cost_detail_summary_entries_by_project_session
+        summary = await loader.load(self.id)
+        return [
+            SpanCostDetailSummaryEntry(
+                token_type=entry.token_type,
+                is_prompt=entry.is_prompt,
+                value=CostBreakdown(
+                    tokens=entry.value.tokens,
+                    cost=entry.value.cost,
+                ),
+            )
+            for entry in summary
+        ]
+
+    @strawberry.field
+    async def session_annotations(
+        self,
+        info: Info[Context, None],
+    ) -> list[Annotated["ProjectSessionAnnotation", lazy(".ProjectSessionAnnotation")]]:
+        """Get all annotations for this session."""
+        from .ProjectSessionAnnotation import ProjectSessionAnnotation
+
+        stmt = select(models.ProjectSessionAnnotation).filter_by(project_session_id=self.id)
+        async with info.context.db() as session:
+            annotations = await session.stream_scalars(stmt)
+            return [
+                ProjectSessionAnnotation(id=annotation.id, db_record=annotation)
+                async for annotation in annotations
+            ]
+
+    @strawberry.field(
+        description="Summarizes each annotation (by name) associated with the session"
+    )  # type: ignore
+    async def session_annotation_summaries(
+        self,
+        info: Info[Context, None],
+        filter: Optional[AnnotationFilter] = None,
+    ) -> list[AnnotationSummary]:
+        """
+        Retrieves and summarizes annotations associated with this span.
+
+        This method aggregates annotation data by name and label, calculating metrics
+        such as count of occurrences and sum of scores. The results are organized
+        into a structured format that can be easily converted to a DataFrame.
+
+        Args:
+            info: GraphQL context information
+            filter: Optional filter to apply to annotations before processing
+
+        Returns:
+            A list of AnnotationSummary objects, each containing:
+            - name: The name of the annotation
+            - data: A list of dictionaries with label statistics
+        """
+        # Load all annotations for this span from the data loader
+        annotations = await info.context.data_loaders.session_annotations_by_session.load(self.id)
+
+        # Apply filter if provided to narrow down the annotations
+        if filter:
+            annotations = [
+                annotation for annotation in annotations if satisfies_filter(annotation, filter)
+            ]
+
+        @dataclass
+        class Metrics:
+            record_count: int = 0
+            label_count: int = 0
+            score_sum: float = 0
+            score_count: int = 0
+
+        summaries: defaultdict[str, defaultdict[Optional[str], Metrics]] = defaultdict(
+            lambda: defaultdict(Metrics)
+        )
+        for annotation in annotations:
+            metrics = summaries[annotation.name][annotation.label]
+            metrics.record_count += 1
+            metrics.label_count += int(annotation.label is not None)
+            metrics.score_sum += annotation.score or 0
+            metrics.score_count += int(annotation.score is not None)
 
-def to_gql_project_session(project_session: models.ProjectSession) -> ProjectSession:
-    return ProjectSession(
-        id_attr=project_session.id,
-        session_id=project_session.session_id,
-        start_time=project_session.start_time,
-        project_rowid=project_session.project_id,
-        end_time=project_session.end_time,
-    )
+        result: list[AnnotationSummary] = []
+        for name, label_metrics in summaries.items():
+            rows = [{"label": label, **asdict(metrics)} for label, metrics in label_metrics.items()]
+            result.append(AnnotationSummary(name=name, df=pd.DataFrame(rows), simple_avg=True))
+        return result
 
 
 INPUT_VALUE = SpanAttributes.INPUT_VALUE.split(".")

phoenix/server/api/types/ProjectSessionAnnotation.py

@@ -0,0 +1,187 @@
+from math import isfinite
+from typing import TYPE_CHECKING, Annotated, Optional
+
+import strawberry
+from strawberry.relay import GlobalID, 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.types.AnnotatorKind import AnnotatorKind
+
+from .Annotation import Annotation
+from .AnnotationSource import AnnotationSource
+
+if TYPE_CHECKING:
+    from .ProjectSession import ProjectSession
+    from .User import User
+
+
+@strawberry.type
+class ProjectSessionAnnotation(Node, Annotation):
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.ProjectSessionAnnotation]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("ProjectSessionAnnotation 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.project_session_annotation_fields.load(
+                (self.id, models.ProjectSessionAnnotation.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.project_session_annotation_fields.load(
+                (self.id, models.ProjectSessionAnnotation.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."  # 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.project_session_annotation_fields.load(
+                (self.id, models.ProjectSessionAnnotation.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.project_session_annotation_fields.load(
+                (self.id, models.ProjectSessionAnnotation.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.project_session_annotation_fields.load(
+                (self.id, models.ProjectSessionAnnotation.explanation),
+            )
+        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.project_session_annotation_fields.load(
+                (self.id, models.ProjectSessionAnnotation.metadata_),
+            )
+        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.project_session_annotation_fields.load(
+                (self.id, models.ProjectSessionAnnotation.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.project_session_annotation_fields.load(
+                (self.id, models.ProjectSessionAnnotation.source),
+            )
+        return AnnotationSource(val)
+
+    @strawberry.field(description="The project session associated with the annotation.")  # type: ignore
+    async def project_session_id(
+        self,
+        info: Info[Context, None],
+    ) -> GlobalID:
+        from phoenix.server.api.types.ProjectSession import ProjectSession
+
+        if self.db_record:
+            project_session_id = self.db_record.project_session_id
+        else:
+            project_session_id = (
+                await info.context.data_loaders.project_session_annotation_fields.load(
+                    (self.id, models.ProjectSessionAnnotation.project_session_id),
+                )
+            )
+        return GlobalID(type_name=ProjectSession.__name__, node_id=str(project_session_id))
+
+    @strawberry.field(description="The project session associated with the annotation.")  # type: ignore
+    async def project_session(
+        self,
+        info: Info[Context, None],
+    ) -> Annotated["ProjectSession", strawberry.lazy(".ProjectSession")]:
+        if self.db_record:
+            project_session_id = self.db_record.project_session_id
+        else:
+            project_session_id = (
+                await info.context.data_loaders.project_session_annotation_fields.load(
+                    (self.id, models.ProjectSessionAnnotation.project_session_id),
+                )
+            )
+        from .ProjectSession import ProjectSession
+
+        return ProjectSession(id=project_session_id)
+
+    @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.project_session_annotation_fields.load(
+                (self.id, models.ProjectSessionAnnotation.user_id),
+            )
+        if user_id is None:
+            return None
+        from .User import User
+
+        return User(id=user_id)

phoenix/server/api/types/ProjectTraceRetentionPolicy.py

@@ -106,5 +106,5 @@ class ProjectTraceRetentionPolicy(Node):
         project_rowids = await info.context.data_loaders.projects_by_trace_retention_policy_id.load(
             self.id
         )
-        data = [Project(project_rowid=project_rowid) for project_rowid in project_rowids]
+        data = [Project(id=project_rowid) for project_rowid in project_rowids]
         return connection_from_list(data=data, args=args)