arize-phoenix 6.2.0__py3-none-any.whl → 7.0.0__py3-none-any.whl

phoenix/server/api/types/Project.py

@@ -1,23 +1,26 @@
 import operator
 from datetime import datetime
-from typing import (
-    Any,
-    ClassVar,
-    Optional,
-)
+from typing import Any, ClassVar, Optional
 
 import strawberry
 from aioitertools.itertools import islice
-from sqlalchemy import and_, desc, distinct, select
+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.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
@@ -28,9 +31,10 @@ from phoenix.server.api.types.pagination import (
     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
+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
 
@@ -127,7 +131,13 @@ class Project(Node):
         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),
+            (
+                "trace",
+                self.id_attr,
+                time_range,
+                None,
+                probability,
+            ),
         )
 
     @strawberry.field
@@ -139,20 +149,26 @@ class Project(Node):
         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),
+            (
+                "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)
+            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 (id_attr := await session.scalar(stmt)) is None:
+            if (trace := await session.scalar(stmt)) is None:
                 return None
-        return Trace(id_attr=id_attr, trace_id=trace_id, project_rowid=self.id_attr)
+        return to_gql_trace(trace)
 
     @strawberry.field
     async def spans(
@@ -241,6 +257,124 @@ class Project(Node):
             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:
+            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
+            )
+        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 annotations for traces. "
         "(The list contains no duplicates.)"
@@ -363,3 +497,7 @@ def to_gql_project(project: models.Project) -> Project:
         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(".")

phoenix/server/api/types/ProjectSession.py

@@ -0,0 +1,139 @@
+from datetime import datetime
+from typing import TYPE_CHECKING, Annotated, ClassVar, Optional, Type
+
+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 phoenix.db import models
+from phoenix.server.api.context import Context
+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.SpanIOValue import SpanIOValue
+from phoenix.server.api.types.TokenUsage import TokenUsage
+
+if TYPE_CHECKING:
+    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
+
+    @strawberry.field
+    async def project_id(self) -> GlobalID:
+        from phoenix.server.api.types.Project import Project
+
+        return GlobalID(type_name=Project.__name__, node_id=str(self.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)
+
+    @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)
+
+    @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)
+        if record is None:
+            return None
+        return SpanIOValue(
+            mime_type=MimeType(record.mime_type.value),
+            value=record.value,
+        )
+
+    @strawberry.field
+    async def last_output(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[SpanIOValue]:
+        record = await info.context.data_loaders.session_last_outputs.load(self.id_attr)
+        if record is None:
+            return None
+        return SpanIOValue(
+            mime_type=MimeType(record.mime_type.value),
+            value=record.value,
+        )
+
+    @strawberry.field
+    async def token_usage(
+        self,
+        info: Info[Context, None],
+    ) -> TokenUsage:
+        usage = await info.context.data_loaders.session_token_usages.load(self.id_attr)
+        return TokenUsage(
+            prompt=usage.prompt,
+            completion=usage.completion,
+        )
+
+    @strawberry.field
+    async def traces(
+        self,
+        info: Info[Context, None],
+        first: Optional[int] = 50,
+        last: Optional[int] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
+    ) -> Connection[Annotated["Trace", lazy(".Trace")]]:
+        from phoenix.server.api.types.Trace import to_gql_trace
+
+        args = ConnectionArgs(
+            first=first,
+            after=after if isinstance(after, CursorString) else None,
+            last=last,
+            before=before if isinstance(before, CursorString) else None,
+        )
+        stmt = (
+            select(models.Trace)
+            .filter_by(project_session_rowid=self.id_attr)
+            .order_by(models.Trace.start_time)
+            .limit(first)
+        )
+        async with info.context.db() as session:
+            traces = await session.stream_scalars(stmt)
+            data = [to_gql_trace(trace) async for trace in traces]
+        return connection_from_list(data=data, args=args)
+
+    @strawberry.field
+    async def trace_latency_ms_quantile(
+        self,
+        info: Info[Context, None],
+        probability: float,
+    ) -> Optional[float]:
+        return await info.context.data_loaders.session_trace_latency_ms_quantile.load(
+            (self.id_attr, probability)
+        )
+
+
+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,
+    )
+
+
+INPUT_VALUE = SpanAttributes.INPUT_VALUE.split(".")
+INPUT_MIME_TYPE = SpanAttributes.INPUT_MIME_TYPE.split(".")
+OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE.split(".")
+OUTPUT_MIME_TYPE = SpanAttributes.OUTPUT_MIME_TYPE.split(".")

phoenix/server/api/types/Span.py

@@ -24,17 +24,16 @@ from phoenix.server.api.input_types.SpanAnnotationSort import (
     SpanAnnotationColumn,
     SpanAnnotationSort,
 )
+from phoenix.server.api.types.DocumentRetrievalMetrics import DocumentRetrievalMetrics
+from phoenix.server.api.types.Evaluation import DocumentEvaluation
+from phoenix.server.api.types.ExampleRevisionInterface import ExampleRevision
 from phoenix.server.api.types.GenerativeProvider import GenerativeProvider
+from phoenix.server.api.types.MimeType import MimeType
 from phoenix.server.api.types.SortDir import SortDir
-from phoenix.server.api.types.SpanAnnotation import to_gql_span_annotation
+from phoenix.server.api.types.SpanAnnotation import SpanAnnotation, to_gql_span_annotation
+from phoenix.server.api.types.SpanIOValue import SpanIOValue
 from phoenix.trace.attributes import get_attribute_value
 
-from .DocumentRetrievalMetrics import DocumentRetrievalMetrics
-from .Evaluation import DocumentEvaluation
-from .ExampleRevisionInterface import ExampleRevision
-from .MimeType import MimeType
-from .SpanAnnotation import SpanAnnotation
-
 if TYPE_CHECKING:
     from phoenix.server.api.types.Project import Project
 
@@ -71,18 +70,6 @@ class SpanContext:
     span_id: ID
 
 
-@strawberry.type
-class SpanIOValue:
-    mime_type: MimeType
-    value: str
-
-    @strawberry.field(
-        description="Truncate value up to `chars` characters, appending '...' if truncated.",
-    )  # type: ignore
-    def truncated_value(self, chars: int = 100) -> str:
-        return f"{self.value[: max(0, chars - 3)]}..." if len(self.value) > chars else self.value
-
-
 @strawberry.enum
 class SpanStatusCode(Enum):
     OK = "OK"

phoenix/server/api/types/SpanIOValue.py

@@ -0,0 +1,15 @@
+import strawberry
+
+from phoenix.server.api.types.MimeType import MimeType
+
+
+@strawberry.type
+class SpanIOValue:
+    mime_type: MimeType
+    value: str
+
+    @strawberry.field(
+        description="Truncate value up to `chars` characters, appending '...' if truncated.",
+    )  # type: ignore
+    def truncated_value(self, chars: int = 100) -> str:
+        return f"{self.value[: max(0, chars - 3)]}..." if len(self.value) > chars else self.value

phoenix/server/api/types/TokenUsage.py

@@ -0,0 +1,11 @@
+import strawberry
+
+
+@strawberry.type
+class TokenUsage:
+    prompt: int = 0
+    completion: int = 0
+
+    @strawberry.field
+    async def total(self) -> int:
+        return self.prompt + self.completion

phoenix/server/api/types/Trace.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
-from typing import Optional
+from datetime import datetime
+from typing import TYPE_CHECKING, Annotated, Optional, Union
 
 import strawberry
+from openinference.semconv.trace import SpanAttributes
 from sqlalchemy import desc, select
 from sqlalchemy.orm import contains_eager
-from strawberry import UNSET, Private
+from strawberry import UNSET, Private, lazy
 from strawberry.relay import Connection, GlobalID, Node, NodeID
 from strawberry.types import Info
 
@@ -21,12 +23,18 @@ from phoenix.server.api.types.SortDir import SortDir
 from phoenix.server.api.types.Span import Span, to_gql_span
 from phoenix.server.api.types.TraceAnnotation import TraceAnnotation, to_gql_trace_annotation
 
+if TYPE_CHECKING:
+    from phoenix.server.api.types.ProjectSession import ProjectSession
+
 
 @strawberry.type
 class Trace(Node):
     id_attr: NodeID[int]
     project_rowid: Private[int]
+    project_session_rowid: Private[Optional[int]]
     trace_id: str
+    start_time: datetime
+    end_time: datetime
 
     @strawberry.field
     async def project_id(self) -> GlobalID:
@@ -34,6 +42,40 @@ class Trace(Node):
 
         return GlobalID(type_name=Project.__name__, node_id=str(self.project_rowid))
 
+    @strawberry.field
+    async def project_session_id(self) -> Optional[GlobalID]:
+        if self.project_session_rowid is None:
+            return None
+        from phoenix.server.api.types.ProjectSession import ProjectSession
+
+        return GlobalID(type_name=ProjectSession.__name__, node_id=str(self.project_session_rowid))
+
+    @strawberry.field
+    async def session(
+        self,
+        info: Info[Context, None],
+    ) -> Union[Annotated["ProjectSession", lazy(".ProjectSession")], None]:
+        if self.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=self.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)
+
+    @strawberry.field
+    async def root_span(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[Span]:
+        span = await info.context.data_loaders.trace_root_spans.load(self.id_attr)
+        if span is None:
+            return None
+        return to_gql_span(span)
+
     @strawberry.field
     async def spans(
         self,
@@ -82,3 +124,18 @@ class Trace(Node):
                 stmt = stmt.order_by(models.TraceAnnotation.created_at.desc())
             annotations = await session.scalars(stmt)
         return [to_gql_trace_annotation(annotation) for annotation in annotations]
+
+
+def to_gql_trace(trace: models.Trace) -> Trace:
+    return Trace(
+        id_attr=trace.id,
+        project_rowid=trace.project_rowid,
+        project_session_rowid=trace.project_session_rowid,
+        trace_id=trace.trace_id,
+        start_time=trace.start_time,
+        end_time=trace.end_time,
+    )
+
+
+INPUT_VALUE = SpanAttributes.INPUT_VALUE.split(".")
+OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE.split(".")

phoenix/server/app.py

@@ -87,12 +87,18 @@ from phoenix.server.api.dataloaders import (
     MinStartOrMaxEndTimeDataLoader,
     ProjectByNameDataLoader,
     RecordCountDataLoader,
+    SessionIODataLoader,
+    SessionNumTracesDataLoader,
+    SessionNumTracesWithErrorDataLoader,
+    SessionTokenUsagesDataLoader,
+    SessionTraceLatencyMsQuantileDataLoader,
     SpanAnnotationsDataLoader,
     SpanDatasetExamplesDataLoader,
     SpanDescendantsDataLoader,
     SpanProjectsDataLoader,
     TokenCountDataLoader,
-    TraceRowIdsDataLoader,
+    TraceByTraceIdsDataLoader,
+    TraceRootSpansDataLoader,
     UserRolesDataLoader,
     UsersDataLoader,
 )
@@ -609,6 +615,12 @@ def create_graphql_router(
                     db,
                     cache_map=cache_for_dataloaders.record_count if cache_for_dataloaders else None,
                 ),
+                session_first_inputs=SessionIODataLoader(db, "first_input"),
+                session_last_outputs=SessionIODataLoader(db, "last_output"),
+                session_num_traces=SessionNumTracesDataLoader(db),
+                session_num_traces_with_error=SessionNumTracesWithErrorDataLoader(db),
+                session_token_usages=SessionTokenUsagesDataLoader(db),
+                session_trace_latency_ms_quantile=SessionTraceLatencyMsQuantileDataLoader(db),
                 span_annotations=SpanAnnotationsDataLoader(db),
                 span_dataset_examples=SpanDatasetExamplesDataLoader(db),
                 span_descendants=SpanDescendantsDataLoader(db),
@@ -617,7 +629,8 @@ def create_graphql_router(
                     db,
                     cache_map=cache_for_dataloaders.token_count if cache_for_dataloaders else None,
                 ),
-                trace_row_ids=TraceRowIdsDataLoader(db),
+                trace_by_trace_ids=TraceByTraceIdsDataLoader(db),
+                trace_root_spans=TraceRootSpansDataLoader(db),
                 project_by_name=ProjectByNameDataLoader(db),
                 users=UsersDataLoader(db),
                 user_roles=UserRolesDataLoader(db),