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

phoenix/server/api/types/Trace.py

@@ -1,27 +1,36 @@
 from __future__ import annotations
 
+from collections import defaultdict
+from dataclasses import asdict, dataclass
 from datetime import datetime
 from typing import TYPE_CHECKING, Annotated, Optional, Union
 
+import pandas as pd
 import strawberry
+from aioitertools.itertools import islice
 from openinference.semconv.trace import SpanAttributes
-from sqlalchemy import desc, select
-from strawberry import ID, UNSET, Private, lazy
+from sqlalchemy import desc, or_, select
+from strawberry import ID, UNSET, lazy
 from strawberry.relay import Connection, GlobalID, Node, NodeID
 from strawberry.types import Info
 from typing_extensions import TypeAlias
 
 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.input_types.TraceAnnotationSort import TraceAnnotationSort
+from phoenix.server.api.types.AnnotationSummary import AnnotationSummary
+from phoenix.server.api.types.CostBreakdown import CostBreakdown
 from phoenix.server.api.types.pagination import (
-    ConnectionArgs,
+    Cursor,
     CursorString,
-    connection_from_list,
+    connection_from_cursors_and_nodes,
 )
 from phoenix.server.api.types.SortDir import SortDir
 from phoenix.server.api.types.Span import Span
-from phoenix.server.api.types.TraceAnnotation import TraceAnnotation, to_gql_trace_annotation
+from phoenix.server.api.types.SpanCostDetailSummaryEntry import SpanCostDetailSummaryEntry
+from phoenix.server.api.types.SpanCostSummary import SpanCostSummary
+from phoenix.server.api.types.TraceAnnotation import TraceAnnotation
 
 if TYPE_CHECKING:
     from phoenix.server.api.types.Project import Project
@@ -33,11 +42,11 @@ TraceRowId: TypeAlias = int
 
 @strawberry.type
 class Trace(Node):
-    trace_rowid: NodeID[TraceRowId]
-    db_trace: Private[models.Trace] = UNSET
+    id: NodeID[TraceRowId]
+    db_record: strawberry.Private[Optional[models.Trace]] = None
 
     def __post_init__(self) -> None:
-        if self.db_trace and self.trace_rowid != self.db_trace.id:
+        if self.db_record and self.id != self.db_record.id:
             raise ValueError("Trace ID mismatch")
 
     @strawberry.field
@@ -45,11 +54,11 @@ class Trace(Node):
         self,
         info: Info[Context, None],
     ) -> ID:
-        if self.db_trace:
-            trace_id = self.db_trace.trace_id
+        if self.db_record:
+            trace_id = self.db_record.trace_id
         else:
             trace_id = await info.context.data_loaders.trace_fields.load(
-                (self.trace_rowid, models.Trace.trace_id),
+                (self.id, models.Trace.trace_id),
             )
         return ID(trace_id)
 
@@ -58,11 +67,11 @@ class Trace(Node):
         self,
         info: Info[Context, None],
     ) -> datetime:
-        if self.db_trace:
-            start_time = self.db_trace.start_time
+        if self.db_record:
+            start_time = self.db_record.start_time
         else:
             start_time = await info.context.data_loaders.trace_fields.load(
-                (self.trace_rowid, models.Trace.start_time),
+                (self.id, models.Trace.start_time),
             )
         return start_time
 
@@ -71,11 +80,11 @@ class Trace(Node):
         self,
         info: Info[Context, None],
     ) -> datetime:
-        if self.db_trace:
-            end_time = self.db_trace.end_time
+        if self.db_record:
+            end_time = self.db_record.end_time
         else:
             end_time = await info.context.data_loaders.trace_fields.load(
-                (self.trace_rowid, models.Trace.end_time),
+                (self.id, models.Trace.end_time),
             )
         return end_time
 
@@ -84,11 +93,11 @@ class Trace(Node):
         self,
         info: Info[Context, None],
     ) -> Optional[float]:
-        if self.db_trace:
-            latency_ms = self.db_trace.latency_ms
+        if self.db_record:
+            latency_ms = self.db_record.latency_ms
         else:
             latency_ms = await info.context.data_loaders.trace_fields.load(
-                (self.trace_rowid, models.Trace.latency_ms),
+                (self.id, models.Trace.latency_ms),
             )
         return latency_ms
 
@@ -97,26 +106,26 @@ class Trace(Node):
         self,
         info: Info[Context, None],
     ) -> Annotated["Project", strawberry.lazy(".Project")]:
-        if self.db_trace:
-            project_rowid = self.db_trace.project_rowid
+        if self.db_record:
+            project_rowid = self.db_record.project_rowid
         else:
             project_rowid = await info.context.data_loaders.trace_fields.load(
-                (self.trace_rowid, models.Trace.project_rowid),
+                (self.id, models.Trace.project_rowid),
             )
         from phoenix.server.api.types.Project import Project
 
-        return Project(project_rowid=project_rowid)
+        return Project(id=project_rowid)
 
     @strawberry.field
     async def project_id(
         self,
         info: Info[Context, None],
     ) -> GlobalID:
-        if self.db_trace:
-            project_rowid = self.db_trace.project_rowid
+        if self.db_record:
+            project_rowid = self.db_record.project_rowid
         else:
             project_rowid = await info.context.data_loaders.trace_fields.load(
-                (self.trace_rowid, models.Trace.project_rowid),
+                (self.id, models.Trace.project_rowid),
             )
         from phoenix.server.api.types.Project import Project
 
@@ -127,11 +136,11 @@ class Trace(Node):
         self,
         info: Info[Context, None],
     ) -> Optional[GlobalID]:
-        if self.db_trace:
-            project_session_rowid = self.db_trace.project_session_rowid
+        if self.db_record:
+            project_session_rowid = self.db_record.project_session_rowid
         else:
             project_session_rowid = await info.context.data_loaders.trace_fields.load(
-                (self.trace_rowid, models.Trace.project_session_rowid),
+                (self.id, models.Trace.project_session_rowid),
             )
         if project_session_rowid is None:
             return None
@@ -144,39 +153,40 @@ class Trace(Node):
         self,
         info: Info[Context, None],
     ) -> Union[Annotated["ProjectSession", lazy(".ProjectSession")], None]:
-        if self.db_trace:
-            project_session_rowid = self.db_trace.project_session_rowid
+        if self.db_record:
+            project_session_rowid = self.db_record.project_session_rowid
         else:
             project_session_rowid = await info.context.data_loaders.trace_fields.load(
-                (self.trace_rowid, models.Trace.project_session_rowid),
+                (self.id, models.Trace.project_session_rowid),
             )
         if project_session_rowid is None:
             return None
-        from phoenix.server.api.types.ProjectSession import to_gql_project_session
 
         stmt = select(models.ProjectSession).filter_by(id=project_session_rowid)
         async with info.context.db() as session:
             project_session = await session.scalar(stmt)
         if project_session is None:
             return None
-        return to_gql_project_session(project_session)
+        from .ProjectSession import ProjectSession
+
+        return ProjectSession(id=project_session.id, db_record=project_session)
 
     @strawberry.field
     async def root_span(
         self,
         info: Info[Context, None],
     ) -> Optional[Span]:
-        span_rowid = await info.context.data_loaders.trace_root_spans.load(self.trace_rowid)
+        span_rowid = await info.context.data_loaders.trace_root_spans.load(self.id)
         if span_rowid is None:
             return None
-        return Span(span_rowid=span_rowid)
+        return Span(id=span_rowid)
 
     @strawberry.field
     async def num_spans(
         self,
         info: Info[Context, None],
     ) -> int:
-        return await info.context.data_loaders.num_spans_per_trace.load(self.trace_rowid)
+        return await info.context.data_loaders.num_spans_per_trace.load(self.id)
 
     @strawberry.field
     async def spans(
@@ -186,26 +196,94 @@ class Trace(Node):
         last: Optional[int] = UNSET,
         after: Optional[CursorString] = UNSET,
         before: Optional[CursorString] = UNSET,
+        root_spans_only: Optional[bool] = UNSET,
+        orphan_span_as_root_span: Optional[bool] = True,
     ) -> Connection[Span]:
-        args = ConnectionArgs(
-            first=first,
-            after=after if isinstance(after, CursorString) else None,
-            last=last,
-            before=before if isinstance(before, CursorString) else None,
-        )
-        stmt = (
+        # Validate pagination arguments
+        if isinstance(first, int) and first <= 0:
+            raise ValueError('Argument "first" must be a positive int')
+
+        # Build base query for spans in this trace
+        base_query = (
             select(models.Span.id)
             .join(models.Trace)
-            .where(models.Trace.id == self.trace_rowid)
+            .where(models.Trace.id == self.id)
             # Sort descending because the root span tends to show up later
             # in the ingestion process.
             .order_by(desc(models.Span.id))
-            .limit(first)
         )
+        # Handle cursor pagination (forward pagination only)
+        if after is not UNSET and after is not None:
+            # Type narrowing: after is guaranteed to be str at this point
+            assert after is not None  # Type narrowing for mypy
+            try:
+                cursor = Cursor.from_string(after)
+            except Exception as e:
+                raise ValueError(f"Invalid cursor format: {after}") from e
+            # For descending order, "after" means we want spans with smaller IDs
+            # (going forward in descending order)
+            base_query = base_query.where(models.Span.id < cursor.rowid)
+        # Note: backward pagination (last/before) is not yet implemented
+        # as it requires more complex handling with reversed ordering
+        if before is not UNSET or (last is not UNSET and last is not None):
+            raise ValueError("Backward pagination (last/before) is not yet supported")
+
+        # Build final query based on filtering requirements
+        if root_spans_only:
+            if orphan_span_as_root_span:
+                # A root span is either a span with no parent_id or an orphan span
+                # (a span whose parent_id references a span that doesn't exist in the current trace)
+                # We need parent_id to check for orphan spans, so add it to the query
+                # and create a CTE
+                candidate_spans = base_query.add_columns(models.Span.parent_id).cte(
+                    "candidate_spans"
+                )
+                # Subquery to get all span_ids that exist in this trace
+                parent_spans_in_trace = (
+                    select(models.Span.span_id)
+                    .where(models.Span.trace_rowid == self.id)
+                    .alias("parent_spans")
+                )
+                # Filter candidates to only root spans (NULL parent_id or orphan spans)
+                stmt = (
+                    select(candidate_spans.c.id)
+                    .where(
+                        or_(
+                            candidate_spans.c.parent_id.is_(None),
+                            ~select(1)
+                            .where(candidate_spans.c.parent_id == parent_spans_in_trace.c.span_id)
+                            .exists(),
+                        )
+                    )
+                    .order_by(desc(candidate_spans.c.id))
+                )
+            else:
+                # Only include explicit root spans (spans with parent_id = NULL)
+                stmt = base_query.where(models.Span.parent_id.is_(None))
+        else:
+            # Return all spans (no root span filtering)
+            stmt = base_query
+
+        # Over-fetch by one to determine whether there's a next page
+        limit = first if isinstance(first, int) else 50
+        stmt = stmt.limit(limit + 1)
+
+        cursors_and_nodes = []
         async with info.context.db() as session:
             span_rowids = await session.stream_scalars(stmt)
-            data = [Span(span_rowid=span_rowid) async for span_rowid in span_rowids]
-        return connection_from_list(data=data, args=args)
+            async for span_rowid in islice(span_rowids, limit):
+                cursor = Cursor(rowid=span_rowid)
+                cursors_and_nodes.append((cursor, Span(id=span_rowid)))
+            has_next_page = True
+            try:
+                await span_rowids.__anext__()
+            except StopAsyncIteration:
+                has_next_page = False
+        return connection_from_cursors_and_nodes(
+            cursors_and_nodes,
+            has_previous_page=False,
+            has_next_page=has_next_page,
+        )
 
     @strawberry.field(description="Annotations associated with the trace.")  # type: ignore
     async def trace_annotations(
@@ -214,7 +292,7 @@ class Trace(Node):
         sort: Optional[TraceAnnotationSort] = None,
     ) -> list[TraceAnnotation]:
         async with info.context.db() as session:
-            stmt = select(models.TraceAnnotation).filter_by(trace_rowid=self.trace_rowid)
+            stmt = select(models.TraceAnnotation).filter_by(trace_rowid=self.id)
             if sort:
                 sort_col = getattr(models.TraceAnnotation, sort.col.value)
                 if sort.dir is SortDir.desc:
@@ -224,7 +302,101 @@ class Trace(Node):
             else:
                 stmt = stmt.order_by(models.TraceAnnotation.created_at.desc())
             annotations = await session.scalars(stmt)
-        return [to_gql_trace_annotation(annotation) for annotation in annotations]
+        return [
+            TraceAnnotation(id=annotation.id, db_record=annotation) for annotation in annotations
+        ]
+
+    @strawberry.field(description="Summarizes each annotation (by name) associated with the trace")  # type: ignore
+    async def trace_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.trace_annotations_by_trace.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)
+
+        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
+
+    @strawberry.field
+    async def cost_summary(
+        self,
+        info: Info[Context, None],
+    ) -> SpanCostSummary:
+        loader = info.context.data_loaders.span_cost_summary_by_trace
+        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_trace
+        entries = 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 entries
+        ]
 
 
 INPUT_VALUE = SpanAttributes.INPUT_VALUE.split(".")

phoenix/server/api/types/TraceAnnotation.py

@@ -1,8 +1,8 @@
-from typing import Optional
+from math import isfinite
+from typing import TYPE_CHECKING, Annotated, Optional
 
 import strawberry
-from strawberry import Private
-from strawberry.relay import GlobalID, Node, NodeID
+from strawberry.relay import Node, NodeID
 from strawberry.scalars import JSON
 from strawberry.types import Info
 
@@ -11,58 +11,157 @@ from phoenix.server.api.context import Context
 from phoenix.server.api.types.AnnotatorKind import AnnotatorKind
 
 from .AnnotationSource import AnnotationSource
-from .User import User, to_gql_user
+
+if TYPE_CHECKING:
+    from .Trace import Trace
+    from .User import User
 
 
 @strawberry.type
 class TraceAnnotation(Node):
-    id_attr: NodeID[int]
-    user_id: Private[Optional[int]]
-    name: str
-    annotator_kind: AnnotatorKind
-    label: Optional[str]
-    score: Optional[float]
-    explanation: Optional[str]
-    metadata: JSON
-    trace_rowid: Private[Optional[int]]
-    identifier: str
-    source: AnnotationSource
-
-    @strawberry.field
-    async def trace_id(self) -> GlobalID:
-        from phoenix.server.api.types.Trace import Trace
-
-        return GlobalID(type_name=Trace.__name__, node_id=str(self.trace_rowid))
-
-    @strawberry.field
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.TraceAnnotation]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("TraceAnnotation 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.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.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.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.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.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.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.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.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.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.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.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.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.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.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.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.source),
+            )
+        return AnnotationSource(val)
+
+    @strawberry.field(description="The trace associated with the annotation.")  # type: ignore
+    async def trace(
+        self,
+        info: Info[Context, None],
+    ) -> Annotated["Trace", strawberry.lazy(".Trace")]:
+        if self.db_record:
+            trace_rowid = self.db_record.trace_rowid
+        else:
+            trace_rowid = await info.context.data_loaders.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.trace_rowid),
+            )
+        from .Trace import Trace
+
+        return Trace(id=trace_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.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.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_trace_annotation(
-    annotation: models.TraceAnnotation,
-) -> TraceAnnotation:
-    """
-    Converts an ORM trace annotation to a GraphQL TraceAnnotation.
-    """
-    return TraceAnnotation(
-        id_attr=annotation.id,
-        user_id=annotation.user_id,
-        trace_rowid=annotation.trace_rowid,
-        name=annotation.name,
-        annotator_kind=AnnotatorKind(annotation.annotator_kind),
-        label=annotation.label,
-        score=annotation.score,
-        explanation=annotation.explanation,
-        metadata=annotation.metadata_,
-        identifier=annotation.identifier,
-        source=AnnotationSource(annotation.source),
-    )
+        from .User import User
+
+        return User(id=user_id)