arize-phoenix 3.16.0__py3-none-any.whl → 7.7.0__py3-none-any.whl

phoenix/server/api/types/Project.py

@@ -1,234 +1,483 @@
+import operator
 from datetime import datetime
-from itertools import chain
-from typing import List, Optional
+from typing import Any, ClassVar, Optional
 
 import strawberry
+from aioitertools.itertools import islice
+from openinference.semconv.trace import SpanAttributes
+from sqlalchemy import and_, desc, distinct, func, or_, select
+from sqlalchemy.orm import contains_eager
+from sqlalchemy.sql.elements import ColumnElement
+from sqlalchemy.sql.expression import tuple_
 from strawberry import ID, UNSET
+from strawberry.relay import Connection, Node, NodeID
+from strawberry.types import Info
+from typing_extensions import assert_never
 
-from phoenix.core.project import Project as CoreProject
-from phoenix.metrics.retrieval_metrics import RetrievalMetrics
-from phoenix.server.api.input_types.SpanSort import SpanSort
+from phoenix.datetime_utils import right_open_time_range
+from phoenix.db import models
+from phoenix.server.api.context import Context
+from phoenix.server.api.input_types.ProjectSessionSort import (
+    ProjectSessionColumn,
+    ProjectSessionSort,
+)
+from phoenix.server.api.input_types.SpanSort import SpanSort, SpanSortConfig
 from phoenix.server.api.input_types.TimeRange import TimeRange
+from phoenix.server.api.types.AnnotationSummary import AnnotationSummary
 from phoenix.server.api.types.DocumentEvaluationSummary import DocumentEvaluationSummary
-from phoenix.server.api.types.EvaluationSummary import EvaluationSummary
-from phoenix.server.api.types.node import Node
 from phoenix.server.api.types.pagination import (
-    Connection,
-    ConnectionArgs,
     Cursor,
-    connection_from_list,
+    CursorSortColumn,
+    CursorString,
+    connection_from_cursors_and_nodes,
 )
+from phoenix.server.api.types.ProjectSession import ProjectSession, to_gql_project_session
+from phoenix.server.api.types.SortDir import SortDir
 from phoenix.server.api.types.Span import Span, to_gql_span
+from phoenix.server.api.types.Trace import Trace, to_gql_trace
 from phoenix.server.api.types.ValidationResult import ValidationResult
 from phoenix.trace.dsl import SpanFilter
-from phoenix.trace.schemas import SpanID, TraceID
 
 
 @strawberry.type
 class Project(Node):
+    _table: ClassVar[type[models.Base]] = models.Project
+    id_attr: NodeID[int]
     name: str
-    project: strawberry.Private[CoreProject]
+    gradient_start_color: str
+    gradient_end_color: str
 
     @strawberry.field
-    def start_time(self) -> Optional[datetime]:
-        start_time, _ = self.project.right_open_time_range
+    async def start_time(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[datetime]:
+        start_time = await info.context.data_loaders.min_start_or_max_end_times.load(
+            (self.id_attr, "start"),
+        )
+        start_time, _ = right_open_time_range(start_time, None)
         return start_time
 
     @strawberry.field
-    def end_time(self) -> Optional[datetime]:
-        _, end_time = self.project.right_open_time_range
+    async def end_time(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[datetime]:
+        end_time = await info.context.data_loaders.min_start_or_max_end_times.load(
+            (self.id_attr, "end"),
+        )
+        _, end_time = right_open_time_range(None, end_time)
         return end_time
 
     @strawberry.field
-    def record_count(
+    async def record_count(
+        self,
+        info: Info[Context, None],
+        time_range: Optional[TimeRange] = UNSET,
+        filter_condition: Optional[str] = UNSET,
+    ) -> int:
+        return await info.context.data_loaders.record_counts.load(
+            ("span", self.id_attr, time_range, filter_condition),
+        )
+
+    @strawberry.field
+    async def trace_count(
+        self,
+        info: Info[Context, None],
+        time_range: Optional[TimeRange] = UNSET,
+    ) -> int:
+        return await info.context.data_loaders.record_counts.load(
+            ("trace", self.id_attr, time_range, None),
+        )
+
+    @strawberry.field
+    async def token_count_total(
         self,
+        info: Info[Context, None],
         time_range: Optional[TimeRange] = UNSET,
+        filter_condition: Optional[str] = UNSET,
+    ) -> int:
+        return await info.context.data_loaders.token_counts.load(
+            ("total", self.id_attr, time_range, filter_condition),
+        )
+
+    @strawberry.field
+    async def token_count_prompt(
+        self,
+        info: Info[Context, None],
+        time_range: Optional[TimeRange] = UNSET,
+        filter_condition: Optional[str] = UNSET,
     ) -> int:
-        if not time_range:
-            return self.project.span_count()
-        return self.project.span_count(time_range.start, time_range.end)
+        return await info.context.data_loaders.token_counts.load(
+            ("prompt", self.id_attr, time_range, filter_condition),
+        )
 
     @strawberry.field
-    def trace_count(
+    async def token_count_completion(
         self,
+        info: Info[Context, None],
         time_range: Optional[TimeRange] = UNSET,
+        filter_condition: Optional[str] = UNSET,
     ) -> int:
-        if not time_range:
-            return self.project.trace_count()
-        return self.project.trace_count(time_range.start, time_range.end)
+        return await info.context.data_loaders.token_counts.load(
+            ("completion", self.id_attr, time_range, filter_condition),
+        )
 
     @strawberry.field
-    def token_count_total(self) -> int:
-        return self.project.token_count_total
+    async def latency_ms_quantile(
+        self,
+        info: Info[Context, None],
+        probability: float,
+        time_range: Optional[TimeRange] = UNSET,
+    ) -> Optional[float]:
+        return await info.context.data_loaders.latency_ms_quantile.load(
+            (
+                "trace",
+                self.id_attr,
+                time_range,
+                None,
+                probability,
+            ),
+        )
 
     @strawberry.field
-    def latency_ms_p50(self) -> Optional[float]:
-        return self.project.root_span_latency_ms_quantiles(0.50)
+    async def span_latency_ms_quantile(
+        self,
+        info: Info[Context, None],
+        probability: float,
+        time_range: Optional[TimeRange] = UNSET,
+        filter_condition: Optional[str] = UNSET,
+    ) -> Optional[float]:
+        return await info.context.data_loaders.latency_ms_quantile.load(
+            (
+                "span",
+                self.id_attr,
+                time_range,
+                filter_condition,
+                probability,
+            ),
+        )
 
     @strawberry.field
-    def latency_ms_p99(self) -> Optional[float]:
-        return self.project.root_span_latency_ms_quantiles(0.99)
+    async def trace(self, trace_id: ID, info: Info[Context, None]) -> Optional[Trace]:
+        stmt = (
+            select(models.Trace)
+            .where(models.Trace.trace_id == str(trace_id))
+            .where(models.Trace.project_rowid == self.id_attr)
+        )
+        async with info.context.db() as session:
+            if (trace := await session.scalar(stmt)) is None:
+                return None
+        return to_gql_trace(trace)
 
     @strawberry.field
-    def spans(
+    async def spans(
         self,
+        info: Info[Context, None],
         time_range: Optional[TimeRange] = UNSET,
-        trace_ids: Optional[List[ID]] = UNSET,
         first: Optional[int] = 50,
         last: Optional[int] = UNSET,
-        after: Optional[Cursor] = UNSET,
-        before: Optional[Cursor] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
         sort: Optional[SpanSort] = UNSET,
         root_spans_only: Optional[bool] = UNSET,
         filter_condition: Optional[str] = UNSET,
     ) -> Connection[Span]:
-        args = ConnectionArgs(
-            first=first,
-            after=after if isinstance(after, Cursor) else None,
-            last=last,
-            before=before if isinstance(before, Cursor) else None,
-        )
-        if not (project := self.project).span_count():
-            return connection_from_list(data=[], args=args)
-        predicate = (
-            SpanFilter(
-                condition=filter_condition,
-                evals=project,
+        stmt = (
+            select(models.Span)
+            .join(models.Trace)
+            .where(models.Trace.project_rowid == self.id_attr)
+            .options(contains_eager(models.Span.trace).load_only(models.Trace.trace_id))
+        )
+        if time_range:
+            stmt = stmt.where(
+                and_(
+                    time_range.start <= models.Span.start_time,
+                    models.Span.start_time < time_range.end,
+                )
             )
-            if filter_condition
-            else None
-        )
-        if not trace_ids:
-            spans = project.get_spans(
-                start_time=time_range.start if time_range else None,
-                stop_time=time_range.end if time_range else None,
-                root_spans_only=root_spans_only,
+        if root_spans_only:
+            # A root span is any span whose parent span is missing in the
+            # database, even if its `parent_span_id` may not be NULL.
+            parent = select(models.Span.span_id).alias()
+            stmt = stmt.outerjoin(
+                parent,
+                models.Span.parent_id == parent.c.span_id,
+            ).where(parent.c.span_id.is_(None))
+        if filter_condition:
+            span_filter = SpanFilter(condition=filter_condition)
+            stmt = span_filter(stmt)
+        sort_config: Optional[SpanSortConfig] = None
+        cursor_rowid_column: Any = models.Span.id
+        if sort:
+            sort_config = sort.update_orm_expr(stmt)
+            stmt = sort_config.stmt
+            if sort_config.dir is SortDir.desc:
+                cursor_rowid_column = desc(cursor_rowid_column)
+        if after:
+            cursor = Cursor.from_string(after)
+            if sort_config and cursor.sort_column:
+                sort_column = cursor.sort_column
+                compare = operator.lt if sort_config.dir is SortDir.desc else operator.gt
+                stmt = stmt.where(
+                    compare(
+                        tuple_(sort_config.orm_expression, models.Span.id),
+                        (sort_column.value, cursor.rowid),
+                    )
+                )
+            else:
+                stmt = stmt.where(models.Span.id > cursor.rowid)
+        if first:
+            stmt = stmt.limit(
+                first + 1  # overfetch by one to determine whether there's a next page
             )
+        stmt = stmt.order_by(cursor_rowid_column)
+        cursors_and_nodes = []
+        async with info.context.db() as session:
+            span_records = await session.execute(stmt)
+            async for span_record in islice(span_records, first):
+                span = span_record[0]
+                cursor = Cursor(rowid=span.id)
+                if sort_config:
+                    assert len(span_record) > 1
+                    cursor.sort_column = CursorSortColumn(
+                        type=sort_config.column_data_type,
+                        value=span_record[1],
+                    )
+                cursors_and_nodes.append((cursor, to_gql_span(span)))
+            has_next_page = True
+            try:
+                next(span_records)
+            except StopIteration:
+                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
+    async def sessions(
+        self,
+        info: Info[Context, None],
+        time_range: Optional[TimeRange] = UNSET,
+        first: Optional[int] = 50,
+        after: Optional[CursorString] = UNSET,
+        sort: Optional[ProjectSessionSort] = UNSET,
+        filter_io_substring: Optional[str] = UNSET,
+    ) -> Connection[ProjectSession]:
+        table = models.ProjectSession
+        stmt = select(table).filter_by(project_id=self.id_attr)
+        if time_range:
+            if time_range.start:
+                stmt = stmt.where(time_range.start <= table.start_time)
+            if time_range.end:
+                stmt = stmt.where(table.start_time < time_range.end)
+        if filter_io_substring:
+            filter_subq = (
+                stmt.with_only_columns(distinct(table.id).label("id"))
+                .join_from(table, models.Trace)
+                .join_from(models.Trace, models.Span)
+                .where(models.Span.parent_id.is_(None))
+                .where(
+                    or_(
+                        models.TextContains(
+                            models.Span.attributes[INPUT_VALUE].as_string(),
+                            filter_io_substring,
+                        ),
+                        models.TextContains(
+                            models.Span.attributes[OUTPUT_VALUE].as_string(),
+                            filter_io_substring,
+                        ),
+                    )
+                )
+            ).subquery()
+            stmt = stmt.join(filter_subq, table.id == filter_subq.c.id)
+        if sort:
+            key: ColumnElement[Any]
+            if sort.col is ProjectSessionColumn.startTime:
+                key = table.start_time.label("key")
+            elif sort.col is ProjectSessionColumn.endTime:
+                key = table.end_time.label("key")
+            elif (
+                sort.col is ProjectSessionColumn.tokenCountTotal
+                or sort.col is ProjectSessionColumn.numTraces
+            ):
+                if sort.col is ProjectSessionColumn.tokenCountTotal:
+                    sort_subq = (
+                        select(
+                            models.Trace.project_session_rowid.label("id"),
+                            func.sum(models.Span.cumulative_llm_token_count_total).label("key"),
+                        )
+                        .join_from(models.Trace, models.Span)
+                        .where(models.Span.parent_id.is_(None))
+                        .group_by(models.Trace.project_session_rowid)
+                    ).subquery()
+                elif sort.col is ProjectSessionColumn.numTraces:
+                    sort_subq = (
+                        select(
+                            models.Trace.project_session_rowid.label("id"),
+                            func.count(models.Trace.id).label("key"),
+                        ).group_by(models.Trace.project_session_rowid)
+                    ).subquery()
+                else:
+                    assert_never(sort.col)
+                key = sort_subq.c.key
+                stmt = stmt.join(sort_subq, table.id == sort_subq.c.id)
+            else:
+                assert_never(sort.col)
+            stmt = stmt.add_columns(key)
+            if sort.dir is SortDir.asc:
+                stmt = stmt.order_by(key.asc(), table.id.asc())
+            else:
+                stmt = stmt.order_by(key.desc(), table.id.desc())
+            if after:
+                cursor = Cursor.from_string(after)
+                assert cursor.sort_column is not None
+                compare = operator.lt if sort.dir is SortDir.desc else operator.gt
+                stmt = stmt.where(
+                    compare(
+                        tuple_(key, table.id),
+                        (cursor.sort_column.value, cursor.rowid),
+                    )
+                )
         else:
-            spans = chain.from_iterable(
-                project.get_trace(trace_id) for trace_id in map(TraceID, trace_ids)
+            stmt = stmt.order_by(table.id.desc())
+            if after:
+                cursor = Cursor.from_string(after)
+                stmt = stmt.where(table.id < cursor.rowid)
+        if first:
+            stmt = stmt.limit(
+                first + 1  # over-fetch by one to determine whether there's a next page
             )
-        if predicate:
-            spans = filter(predicate, spans)
-        if sort:
-            spans = sort(spans, evals=project)
-        data = [to_gql_span(span, project) for span in spans]
-        return connection_from_list(data=data, args=args)
+        cursors_and_nodes = []
+        async with info.context.db() as session:
+            records = await session.stream(stmt)
+            async for record in islice(records, first):
+                project_session = record[0]
+                cursor = Cursor(rowid=project_session.id)
+                if sort:
+                    assert len(record) > 1
+                    cursor.sort_column = CursorSortColumn(
+                        type=sort.col.data_type,
+                        value=record[1],
+                    )
+                cursors_and_nodes.append((cursor, to_gql_project_session(project_session)))
+            has_next_page = True
+            try:
+                await records.__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="Names of all available evaluations for spans. "
+        description="Names of all available annotations for traces. "
         "(The list contains no duplicates.)"
     )  # type: ignore
-    def span_evaluation_names(self) -> List[str]:
-        return self.project.get_span_evaluation_names()
+    async def trace_annotations_names(
+        self,
+        info: Info[Context, None],
+    ) -> list[str]:
+        stmt = (
+            select(distinct(models.TraceAnnotation.name))
+            .join(models.Trace)
+            .where(models.Trace.project_rowid == self.id_attr)
+        )
+        async with info.context.db() as session:
+            return list(await session.scalars(stmt))
+
+    @strawberry.field(
+        description="Names of all available annotations for spans. "
+        "(The list contains no duplicates.)"
+    )  # type: ignore
+    async def span_annotation_names(
+        self,
+        info: Info[Context, None],
+    ) -> list[str]:
+        stmt = (
+            select(distinct(models.SpanAnnotation.name))
+            .join(models.Span)
+            .join(models.Trace, models.Span.trace_rowid == models.Trace.id)
+            .where(models.Trace.project_rowid == self.id_attr)
+        )
+        async with info.context.db() as session:
+            return list(await session.scalars(stmt))
 
     @strawberry.field(
         description="Names of available document evaluations.",
     )  # type: ignore
-    def document_evaluation_names(
+    async def document_evaluation_names(
         self,
+        info: Info[Context, None],
         span_id: Optional[ID] = UNSET,
-    ) -> List[str]:
-        return self.project.get_document_evaluation_names(
-            None if span_id is UNSET else SpanID(span_id),
+    ) -> list[str]:
+        stmt = (
+            select(distinct(models.DocumentAnnotation.name))
+            .join(models.Span)
+            .join(models.Trace, models.Span.trace_rowid == models.Trace.id)
+            .where(models.Trace.project_rowid == self.id_attr)
+            .where(models.DocumentAnnotation.annotator_kind == "LLM")
         )
+        if span_id:
+            stmt = stmt.where(models.Span.span_id == str(span_id))
+        async with info.context.db() as session:
+            return list(await session.scalars(stmt))
 
     @strawberry.field
-    def span_evaluation_summary(
+    async def trace_annotation_summary(
         self,
-        evaluation_name: str,
+        info: Info[Context, None],
+        annotation_name: str,
+        time_range: Optional[TimeRange] = UNSET,
+    ) -> Optional[AnnotationSummary]:
+        return await info.context.data_loaders.annotation_summaries.load(
+            ("trace", self.id_attr, time_range, None, annotation_name),
+        )
+
+    @strawberry.field
+    async def span_annotation_summary(
+        self,
+        info: Info[Context, None],
+        annotation_name: str,
         time_range: Optional[TimeRange] = UNSET,
         filter_condition: Optional[str] = UNSET,
-    ) -> Optional[EvaluationSummary]:
-        project = self.project
-        predicate = (
-            SpanFilter(
-                condition=filter_condition,
-                evals=project,
-            )
-            if filter_condition
-            else None
-        )
-        span_ids = project.get_span_evaluation_span_ids(evaluation_name)
-        if not span_ids:
-            return None
-        spans = project.get_spans(
-            start_time=time_range.start if time_range else None,
-            stop_time=time_range.end if time_range else None,
-            span_ids=span_ids,
-        )
-        if predicate:
-            spans = filter(predicate, spans)
-        evaluations = tuple(
-            evaluation
-            for span in spans
-            if (
-                evaluation := project.get_span_evaluation(
-                    span.context.span_id,
-                    evaluation_name,
-                )
-            )
-            is not None
+    ) -> Optional[AnnotationSummary]:
+        return await info.context.data_loaders.annotation_summaries.load(
+            ("span", self.id_attr, time_range, filter_condition, annotation_name),
         )
-        if not evaluations:
-            return None
-        labels = project.get_span_evaluation_labels(evaluation_name)
-        return EvaluationSummary(evaluations, labels)
 
     @strawberry.field
-    def document_evaluation_summary(
+    async def document_evaluation_summary(
         self,
+        info: Info[Context, None],
         evaluation_name: str,
         time_range: Optional[TimeRange] = UNSET,
         filter_condition: Optional[str] = UNSET,
     ) -> Optional[DocumentEvaluationSummary]:
-        project = self.project
-        predicate = (
-            SpanFilter(condition=filter_condition, evals=project) if filter_condition else None
-        )
-        span_ids = project.get_document_evaluation_span_ids(evaluation_name)
-        if not span_ids:
-            return None
-        spans = project.get_spans(
-            start_time=time_range.start if time_range else None,
-            stop_time=time_range.end if time_range else None,
-            span_ids=span_ids,
-        )
-        if predicate:
-            spans = filter(predicate, spans)
-        metrics_collection = []
-        for span in spans:
-            span_id = span.context.span_id
-            num_documents = project.get_num_documents(span_id)
-            if not num_documents:
-                continue
-            evaluation_scores = project.get_document_evaluation_scores(
-                span_id=span_id,
-                evaluation_name=evaluation_name,
-                num_documents=num_documents,
-            )
-            metrics_collection.append(RetrievalMetrics(evaluation_scores))
-        if not metrics_collection:
-            return None
-        return DocumentEvaluationSummary(
-            evaluation_name=evaluation_name,
-            metrics_collection=metrics_collection,
+        return await info.context.data_loaders.document_evaluation_summaries.load(
+            (self.id_attr, time_range, filter_condition, evaluation_name),
         )
 
     @strawberry.field
     def streaming_last_updated_at(
         self,
+        info: Info[Context, None],
     ) -> Optional[datetime]:
-        return self.project.last_updated_at
+        return info.context.last_updated_at.get(self._table, self.id_attr)
 
     @strawberry.field
-    def validate_span_filter_condition(self, condition: str) -> ValidationResult:
-        valid_eval_names = self.project.get_span_evaluation_names()
+    async def validate_span_filter_condition(self, condition: str) -> ValidationResult:
+        # This query is too expensive to run on every validation
+        # valid_eval_names = await self.span_annotation_names()
         try:
             SpanFilter(
                 condition=condition,
-                evals=self.project,
-                valid_eval_names=valid_eval_names,
+                # valid_eval_names=valid_eval_names,
             )
             return ValidationResult(is_valid=True, error_message=None)
         except SyntaxError as e:
@@ -236,3 +485,19 @@ class Project(Node):
                 is_valid=False,
                 error_message=e.msg,
             )
+
+
+def to_gql_project(project: models.Project) -> Project:
+    """
+    Converts an ORM project to a GraphQL Project.
+    """
+    return Project(
+        id_attr=project.id,
+        name=project.name,
+        gradient_start_color=project.gradient_start_color,
+        gradient_end_color=project.gradient_end_color,
+    )
+
+
+INPUT_VALUE = SpanAttributes.INPUT_VALUE.split(".")
+OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE.split(".")