arize-phoenix 3.24.0__py3-none-any.whl → 4.0.0__py3-none-any.whl

phoenix/server/api/types/Project.py

@@ -1,284 +1,353 @@
+import operator
 from datetime import datetime
-from itertools import chain
-from typing import List, Optional
+from typing import Any, List, Optional
 
 import strawberry
+from aioitertools.itertools import islice
+from sqlalchemy import and_, desc, distinct, select
+from sqlalchemy.orm import contains_eager
+from sqlalchemy.sql.expression import tuple_
 from strawberry import ID, UNSET
+from strawberry.types import Info
 
-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.SpanSort import SpanSort, SpanSortConfig
 from phoenix.server.api.input_types.TimeRange import TimeRange
 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,
+    connections,
 )
+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
 from phoenix.server.api.types.ValidationResult import ValidationResult
 from phoenix.trace.dsl import SpanFilter
-from phoenix.trace.schemas import SpanID, TraceID
+
+SPANS_LIMIT = 1000
 
 
 @strawberry.type
 class Project(Node):
     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:
-        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.record_counts.load(
+            ("span", self.id_attr, time_range, filter_condition),
+        )
 
     @strawberry.field
-    def trace_count(
+    async def trace_count(
         self,
+        info: Info[Context, None],
         time_range: Optional[TimeRange] = 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.record_counts.load(
+            ("trace", self.id_attr, time_range, None),
+        )
 
     @strawberry.field
-    def token_count_total(self) -> int:
-        return self.project.token_count_total
+    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
-    def latency_ms_p50(self) -> Optional[float]:
-        return self.project.root_span_latency_ms_quantiles(0.50)
+    async def token_count_prompt(
+        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(
+            ("prompt", self.id_attr, time_range, filter_condition),
+        )
 
     @strawberry.field
-    def latency_ms_p99(self) -> Optional[float]:
-        return self.project.root_span_latency_ms_quantiles(0.99)
+    async def token_count_completion(
+        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(
+            ("completion", self.id_attr, time_range, filter_condition),
+        )
 
     @strawberry.field
-    def trace(self, trace_id: ID) -> Optional[Trace]:
-        if self.project.has_trace(TraceID(trace_id)):
-            return Trace(trace_id=trace_id, project=self.project)
-        return None
+    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 spans(
+    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
+    async def trace(self, trace_id: ID, info: Info[Context, None]) -> Optional[Trace]:
+        stmt = (
+            select(models.Trace.id)
+            .where(models.Trace.trace_id == str(trace_id))
+            .where(models.Trace.project_rowid == self.id_attr)
+        )
+        async with info.context.db() as session:
+            if (id_attr := await session.scalar(stmt)) is None:
+                return None
+        return Trace(id_attr=id_attr)
+
+    @strawberry.field
+    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,
+        stmt = (
+            select(models.Span)
+            .join(models.Trace)
+            .where(models.Trace.project_rowid == self.id_attr)
+            .options(contains_eager(models.Span.trace))
         )
-        start_time = time_range.start if time_range else None
-        stop_time = time_range.end if time_range else None
-        if not (project := self.project).span_count(
-            start_time=start_time,
-            stop_time=stop_time,
-        ):
-            return connection_from_list(data=[], args=args)
-        predicate = (
-            SpanFilter(
-                condition=filter_condition,
-                evals=project,
-            )
-            if filter_condition
-            else None
-        )
-        if not trace_ids:
-            spans = project.get_spans(
-                start_time=start_time,
-                stop_time=stop_time,
-                root_spans_only=root_spans_only,
-            )
-        else:
-            spans = chain.from_iterable(
-                project.get_trace(trace_id) for trace_id in map(TraceID, trace_ids)
+        if time_range:
+            stmt = stmt.where(
+                and_(
+                    time_range.start <= models.Span.start_time,
+                    models.Span.start_time < time_range.end,
+                )
             )
-        if predicate:
-            spans = filter(predicate, spans)
+        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:
-            spans = sort(spans, evals=project)
-        data = [to_gql_span(span, project) for span in spans]
-        return connection_from_list(data=data, args=args)
+            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)
+        data = []
+        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]
+                sort_column_value = span_record[1] if len(span_record) > 1 else None
+                cursor = Cursor(
+                    rowid=span.id,
+                    sort_column=(
+                        CursorSortColumn(
+                            type=sort_config.column_data_type,
+                            value=sort_column_value,
+                        )
+                        if sort_config
+                        else None
+                    ),
+                )
+                data.append((cursor, to_gql_span(span)))
+            has_next_page = True
+            try:
+                next(span_records)
+            except StopIteration:
+                has_next_page = False
+
+        return connections(
+            data,
+            has_previous_page=False,
+            has_next_page=has_next_page,
+        )
 
     @strawberry.field(
         description="Names of all available evaluations for traces. "
         "(The list contains no duplicates.)"
     )  # type: ignore
-    def trace_evaluation_names(self) -> List[str]:
-        return self.project.get_trace_evaluation_names()
+    async def trace_evaluation_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)
+            .where(models.TraceAnnotation.annotator_kind == "LLM")
+        )
+        async with info.context.db() as session:
+            return list(await session.scalars(stmt))
 
     @strawberry.field(
         description="Names of all available evaluations for spans. "
         "(The list contains no duplicates.)"
     )  # type: ignore
-    def span_evaluation_names(self) -> List[str]:
-        return self.project.get_span_evaluation_names()
+    async def span_evaluation_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)
+            .where(models.SpanAnnotation.annotator_kind == "LLM")
+        )
+        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),
+        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 trace_evaluation_summary(
+    async def trace_evaluation_summary(
         self,
+        info: Info[Context, None],
         evaluation_name: str,
         time_range: Optional[TimeRange] = UNSET,
     ) -> Optional[EvaluationSummary]:
-        project = self.project
-        eval_trace_ids = project.get_trace_evaluation_trace_ids(evaluation_name)
-        if not eval_trace_ids:
-            return None
-        trace_ids = project.get_trace_ids(
-            start_time=time_range.start if time_range else None,
-            stop_time=time_range.end if time_range else None,
-            trace_ids=eval_trace_ids,
+        return await info.context.data_loaders.evaluation_summaries.load(
+            ("trace", self.id_attr, time_range, None, evaluation_name),
         )
-        evaluations = tuple(
-            evaluation
-            for trace_id in trace_ids
-            if (
-                evaluation := project.get_trace_evaluation(
-                    trace_id,
-                    evaluation_name,
-                )
-            )
-            is not None
-        )
-        if not evaluations:
-            return None
-        labels = project.get_trace_evaluation_labels(evaluation_name)
-        return EvaluationSummary(evaluations, labels)
 
     @strawberry.field
-    def span_evaluation_summary(
+    async def span_evaluation_summary(
         self,
+        info: Info[Context, None],
         evaluation_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
+        return await info.context.data_loaders.evaluation_summaries.load(
+            ("span", self.id_attr, time_range, filter_condition, evaluation_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.streaming_last_updated_at(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_evaluation_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: