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

phoenix/db/models.py

@@ -156,14 +156,37 @@ class Project(Base):
     )
 
 
+class ProjectSession(Base):
+    __tablename__ = "project_sessions"
+    id: Mapped[int] = mapped_column(primary_key=True)
+    session_id: Mapped[str] = mapped_column(String, nullable=False, unique=True)
+    project_id: Mapped[int] = mapped_column(
+        ForeignKey("projects.id", ondelete="CASCADE"),
+        nullable=False,
+        index=True,
+    )
+    start_time: Mapped[datetime] = mapped_column(UtcTimeStamp, index=True, nullable=False)
+    end_time: Mapped[datetime] = mapped_column(UtcTimeStamp, index=True, nullable=False)
+    traces: Mapped[list["Trace"]] = relationship(
+        "Trace",
+        back_populates="project_session",
+        uselist=True,
+    )
+
+
 class Trace(Base):
     __tablename__ = "traces"
     id: Mapped[int] = mapped_column(primary_key=True)
     project_rowid: Mapped[int] = mapped_column(
         ForeignKey("projects.id", ondelete="CASCADE"),
+        nullable=False,
         index=True,
     )
     trace_id: Mapped[str]
+    project_session_rowid: Mapped[Optional[int]] = mapped_column(
+        ForeignKey("project_sessions.id", ondelete="CASCADE"),
+        index=True,
+    )
     start_time: Mapped[datetime] = mapped_column(UtcTimeStamp, index=True)
     end_time: Mapped[datetime] = mapped_column(UtcTimeStamp)
 
@@ -188,6 +211,10 @@ class Trace(Base):
         cascade="all, delete-orphan",
         uselist=True,
     )
+    project_session: Mapped[ProjectSession] = relationship(
+        "ProjectSession",
+        back_populates="traces",
+    )
     experiment_runs: Mapped[list["ExperimentRun"]] = relationship(
         primaryjoin="foreign(ExperimentRun.trace_id) == Trace.trace_id",
         back_populates="trace",

phoenix/server/api/context.py

@@ -31,12 +31,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,
 )
@@ -68,12 +74,19 @@ class DataLoaders:
     latency_ms_quantile: LatencyMsQuantileDataLoader
     min_start_or_max_end_times: MinStartOrMaxEndTimeDataLoader
     record_counts: RecordCountDataLoader
+    session_first_inputs: SessionIODataLoader
+    session_last_outputs: SessionIODataLoader
+    session_num_traces: SessionNumTracesDataLoader
+    session_num_traces_with_error: SessionNumTracesWithErrorDataLoader
+    session_token_usages: SessionTokenUsagesDataLoader
+    session_trace_latency_ms_quantile: SessionTraceLatencyMsQuantileDataLoader
     span_annotations: SpanAnnotationsDataLoader
     span_dataset_examples: SpanDatasetExamplesDataLoader
     span_descendants: SpanDescendantsDataLoader
     span_projects: SpanProjectsDataLoader
     token_counts: TokenCountDataLoader
-    trace_row_ids: TraceRowIdsDataLoader
+    trace_by_trace_ids: TraceByTraceIdsDataLoader
+    trace_root_spans: TraceRootSpansDataLoader
     project_by_name: ProjectByNameDataLoader
     users: UsersDataLoader
     user_roles: UserRolesDataLoader

phoenix/server/api/dataloaders/__init__.py

@@ -19,12 +19,18 @@ from .latency_ms_quantile import LatencyMsQuantileCache, LatencyMsQuantileDataLo
 from .min_start_or_max_end_times import MinStartOrMaxEndTimeCache, MinStartOrMaxEndTimeDataLoader
 from .project_by_name import ProjectByNameDataLoader
 from .record_counts import RecordCountCache, RecordCountDataLoader
+from .session_io import SessionIODataLoader
+from .session_num_traces import SessionNumTracesDataLoader
+from .session_num_traces_with_error import SessionNumTracesWithErrorDataLoader
+from .session_token_usages import SessionTokenUsagesDataLoader
+from .session_trace_latency_ms_quantile import SessionTraceLatencyMsQuantileDataLoader
 from .span_annotations import SpanAnnotationsDataLoader
 from .span_dataset_examples import SpanDatasetExamplesDataLoader
 from .span_descendants import SpanDescendantsDataLoader
 from .span_projects import SpanProjectsDataLoader
 from .token_counts import TokenCountCache, TokenCountDataLoader
-from .trace_row_ids import TraceRowIdsDataLoader
+from .trace_by_trace_ids import TraceByTraceIdsDataLoader
+from .trace_root_spans import TraceRootSpansDataLoader
 from .user_roles import UserRolesDataLoader
 from .users import UsersDataLoader
 
@@ -45,11 +51,17 @@ __all__ = [
     "LatencyMsQuantileDataLoader",
     "MinStartOrMaxEndTimeDataLoader",
     "RecordCountDataLoader",
+    "SessionIODataLoader",
+    "SessionNumTracesDataLoader",
+    "SessionNumTracesWithErrorDataLoader",
+    "SessionTokenUsagesDataLoader",
+    "SessionTraceLatencyMsQuantileDataLoader",
     "SpanDatasetExamplesDataLoader",
     "SpanDescendantsDataLoader",
     "SpanProjectsDataLoader",
     "TokenCountDataLoader",
-    "TraceRowIdsDataLoader",
+    "TraceByTraceIdsDataLoader",
+    "TraceRootSpansDataLoader",
     "ProjectByNameDataLoader",
     "SpanAnnotationsDataLoader",
     "UsersDataLoader",

phoenix/server/api/dataloaders/session_io.py

@@ -0,0 +1,75 @@
+from functools import cached_property
+from typing import Literal, Optional, cast
+
+from openinference.semconv.trace import SpanAttributes
+from sqlalchemy import Select, func, select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias, assert_never
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+from phoenix.trace.schemas import MimeType, SpanIOValue
+
+Key: TypeAlias = int
+Result: TypeAlias = Optional[SpanIOValue]
+
+Kind = Literal["first_input", "last_output"]
+
+
+class SessionIODataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory, kind: Kind) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+        self._kind = kind
+
+    @cached_property
+    def _subq(self) -> Select[tuple[Optional[int], str, str, int]]:
+        stmt = (
+            select(models.Trace.project_session_rowid.label("id_"))
+            .join_from(models.Span, models.Trace)
+            .where(models.Span.parent_id.is_(None))
+        )
+        if self._kind == "first_input":
+            stmt = stmt.add_columns(
+                models.Span.attributes[INPUT_VALUE].label("value"),
+                models.Span.attributes[INPUT_MIME_TYPE].label("mime_type"),
+                func.row_number()
+                .over(
+                    partition_by=models.Trace.project_session_rowid,
+                    order_by=[models.Trace.start_time.asc(), models.Trace.id.asc()],
+                )
+                .label("rank"),
+            )
+        elif self._kind == "last_output":
+            stmt = stmt.add_columns(
+                models.Span.attributes[OUTPUT_VALUE].label("value"),
+                models.Span.attributes[OUTPUT_MIME_TYPE].label("mime_type"),
+                func.row_number()
+                .over(
+                    partition_by=models.Trace.project_session_rowid,
+                    order_by=[models.Trace.start_time.desc(), models.Trace.id.desc()],
+                )
+                .label("rank"),
+            )
+        else:
+            assert_never(self._kind)
+        return cast(Select[tuple[Optional[int], str, str, int]], stmt)
+
+    def _stmt(self, *keys: Key) -> Select[tuple[int, str, str]]:
+        subq = self._subq.where(models.Trace.project_session_rowid.in_(keys)).subquery()
+        return select(subq.c.id_, subq.c.value, subq.c.mime_type).filter_by(rank=1)
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        async with self._db() as session:
+            result: dict[Key, SpanIOValue] = {
+                id_: SpanIOValue(value=value, mime_type=MimeType(mime_type))
+                async for id_, value, mime_type in await session.stream(self._stmt(*keys))
+                if id_ is not None
+            }
+        return [result.get(key) for key in keys]
+
+
+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/dataloaders/session_num_traces.py

@@ -0,0 +1,30 @@
+from sqlalchemy import func, select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+Key: TypeAlias = int
+Result: TypeAlias = int
+
+
+class SessionNumTracesDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        stmt = (
+            select(
+                models.Trace.project_session_rowid.label("id_"),
+                func.count(models.Trace.id).label("value"),
+            )
+            .group_by(models.Trace.project_session_rowid)
+            .where(models.Trace.project_session_rowid.in_(keys))
+        )
+        async with self._db() as session:
+            result: dict[Key, int] = {
+                id_: value async for id_, value in await session.stream(stmt) if id_ is not None
+            }
+        return [result.get(key, 0) for key in keys]

phoenix/server/api/dataloaders/session_num_traces_with_error.py

@@ -0,0 +1,32 @@
+from sqlalchemy import func, select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+Key: TypeAlias = int
+Result: TypeAlias = int
+
+
+class SessionNumTracesWithErrorDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        stmt = (
+            select(
+                models.Trace.project_session_rowid.label("id_"),
+                func.count(models.Trace.id).label("value"),
+            )
+            .join(models.Span)
+            .group_by(models.Trace.project_session_rowid)
+            .where(models.Span.cumulative_error_count > 0)
+            .where(models.Trace.project_session_rowid.in_(keys))
+        )
+        async with self._db() as session:
+            result: dict[Key, int] = {
+                id_: value async for id_, value in await session.stream(stmt) if id_ is not None
+            }
+        return [result.get(key, 0) for key in keys]

phoenix/server/api/dataloaders/session_token_usages.py

@@ -0,0 +1,41 @@
+from sqlalchemy import func, select
+from sqlalchemy.sql.functions import coalesce
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+from phoenix.trace.schemas import TokenUsage
+
+Key: TypeAlias = int
+Result: TypeAlias = TokenUsage
+
+
+class SessionTokenUsagesDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        stmt = (
+            select(
+                models.Trace.project_session_rowid.label("id_"),
+                func.sum(coalesce(models.Span.cumulative_llm_token_count_prompt, 0)).label(
+                    "prompt"
+                ),
+                func.sum(coalesce(models.Span.cumulative_llm_token_count_completion, 0)).label(
+                    "completion"
+                ),
+            )
+            .join_from(models.Span, models.Trace)
+            .where(models.Span.parent_id.is_(None))
+            .where(models.Trace.project_session_rowid.in_(keys))
+            .group_by(models.Trace.project_session_rowid)
+        )
+        async with self._db() as session:
+            result: dict[Key, TokenUsage] = {
+                id_: TokenUsage(prompt=prompt, completion=completion)
+                async for id_, prompt, completion in await session.stream(stmt)
+                if id_ is not None
+            }
+        return [result.get(key, TokenUsage()) for key in keys]

phoenix/server/api/dataloaders/session_trace_latency_ms_quantile.py

@@ -0,0 +1,55 @@
+from collections import defaultdict
+from typing import Optional
+
+import numpy as np
+from aioitertools.itertools import groupby
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+SessionId: TypeAlias = int
+Probability: TypeAlias = float
+QuantileValue: TypeAlias = float
+
+Key: TypeAlias = tuple[SessionId, Probability]
+Result: TypeAlias = Optional[QuantileValue]
+ResultPosition: TypeAlias = int
+
+DEFAULT_VALUE: Result = None
+
+
+class SessionTraceLatencyMsQuantileDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        results: list[Result] = [DEFAULT_VALUE] * len(keys)
+        argument_position_map: defaultdict[
+            SessionId, defaultdict[Probability, list[ResultPosition]]
+        ] = defaultdict(lambda: defaultdict(list))
+        session_rowids = {session_id for session_id, _ in keys}
+        for position, (session_id, probability) in enumerate(keys):
+            argument_position_map[session_id][probability].append(position)
+        stmt = (
+            select(
+                models.Trace.project_session_rowid,
+                models.Trace.latency_ms,
+            )
+            .where(models.Trace.project_session_rowid.in_(session_rowids))
+            .order_by(models.Trace.project_session_rowid)
+        )
+        async with self._db() as session:
+            data = await session.stream(stmt)
+            async for project_session_rowid, group in groupby(
+                data, lambda row: row.project_session_rowid
+            ):
+                session_latencies = [row.latency_ms for row in group]
+                for probability, positions in argument_position_map[project_session_rowid].items():
+                    quantile_value = np.quantile(session_latencies, probability)
+                    for position in positions:
+                        results[position] = quantile_value
+        return results

phoenix/server/api/dataloaders/trace_by_trace_ids.py

@@ -0,0 +1,25 @@
+from typing import List, Optional
+
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+Key: TypeAlias = str
+Result: TypeAlias = Optional[models.Trace]
+
+
+class TraceByTraceIdsDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+        stmt = select(models.Trace).where(models.Trace.trace_id.in_(keys))
+        async with self._db() as session:
+            result: dict[Key, models.Trace] = {
+                trace.trace_id: trace async for trace in await session.stream_scalars(stmt)
+            }
+        return [result.get(trace_id) for trace_id in keys]

phoenix/server/api/dataloaders/trace_root_spans.py

@@ -0,0 +1,32 @@
+from typing import List, Optional
+
+from sqlalchemy import select
+from sqlalchemy.orm import contains_eager
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+Key: TypeAlias = int
+Result: TypeAlias = Optional[models.Span]
+
+
+class TraceRootSpansDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+        stmt = (
+            select(models.Span)
+            .join(models.Trace)
+            .where(models.Span.parent_id.is_(None))
+            .where(models.Trace.id.in_(keys))
+            .options(contains_eager(models.Span.trace).load_only(models.Trace.trace_id))
+        )
+        async with self._db() as session:
+            result: dict[Key, models.Span] = {
+                span.trace_rowid: span async for span in await session.stream_scalars(stmt)
+            }
+        return [result.get(key) for key in keys]

phoenix/server/api/input_types/ProjectSessionSort.py

@@ -0,0 +1,29 @@
+from enum import Enum, auto
+
+import strawberry
+from typing_extensions import assert_never
+
+from phoenix.server.api.types.pagination import CursorSortColumnDataType
+from phoenix.server.api.types.SortDir import SortDir
+
+
+@strawberry.enum
+class ProjectSessionColumn(Enum):
+    startTime = auto()
+    endTime = auto()
+    tokenCountTotal = auto()
+    numTraces = auto()
+
+    @property
+    def data_type(self) -> CursorSortColumnDataType:
+        if self is ProjectSessionColumn.tokenCountTotal or self is ProjectSessionColumn.numTraces:
+            return CursorSortColumnDataType.INT
+        if self is ProjectSessionColumn.startTime or self is ProjectSessionColumn.endTime:
+            return CursorSortColumnDataType.DATETIME
+        assert_never(self)
+
+
+@strawberry.input(description="The sort key and direction for ProjectSession connections.")
+class ProjectSessionSort:
+    col: ProjectSessionColumn
+    dir: SortDir

phoenix/server/api/mutations/project_mutations.py

@@ -38,10 +38,20 @@ class ProjectMutationMixin:
         project_id = from_global_id_with_expected_type(
             global_id=input.id, expected_type_name="Project"
         )
-        delete_statement = delete(models.Trace).where(models.Trace.project_rowid == project_id)
+        delete_statement = (
+            delete(models.Trace)
+            .where(models.Trace.project_rowid == project_id)
+            .returning(models.Trace.project_session_rowid)
+        )
         if input.end_time:
             delete_statement = delete_statement.where(models.Trace.start_time < input.end_time)
         async with info.context.db() as session:
-            await session.execute(delete_statement)
+            deleted_trace_project_session_ids = await session.scalars(delete_statement)
+            if deleted_trace_project_session_ids:
+                await session.execute(
+                    delete(models.ProjectSession).where(
+                        models.ProjectSession.id.in_(set(deleted_trace_project_session_ids))
+                    )
+                )
         info.context.event_queue.put(SpanDeleteEvent((project_id,)))
         return Query()

phoenix/server/api/queries.py

@@ -78,10 +78,11 @@ from phoenix.server.api.types.pagination import (
     connection_from_list,
 )
 from phoenix.server.api.types.Project import Project
+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.SystemApiKey import SystemApiKey
-from phoenix.server.api.types.Trace import Trace
+from phoenix.server.api.types.Trace import to_gql_trace
 from phoenix.server.api.types.User import User, to_gql_user
 from phoenix.server.api.types.UserApiKey import UserApiKey, to_gql_api_key
 from phoenix.server.api.types.UserRole import UserRole
@@ -445,17 +446,12 @@ class Query:
                 gradient_end_color=project.gradient_end_color,
             )
         elif type_name == "Trace":
-            trace_stmt = select(
-                models.Trace.id,
-                models.Trace.project_rowid,
-            ).where(models.Trace.id == node_id)
+            trace_stmt = select(models.Trace).filter_by(id=node_id)
             async with info.context.db() as session:
-                trace = (await session.execute(trace_stmt)).first()
+                trace = await session.scalar(trace_stmt)
             if trace is None:
                 raise NotFound(f"Unknown trace: {id}")
-            return Trace(
-                id_attr=trace.id, trace_id=trace.trace_id, project_rowid=trace.project_rowid
-            )
+            return to_gql_trace(trace)
         elif type_name == Span.__name__:
             span_stmt = (
                 select(models.Span)
@@ -544,6 +540,15 @@ class Query:
                 ):
                     raise NotFound(f"Unknown user: {id}")
             return to_gql_user(user)
+        elif type_name == ProjectSession.__name__:
+            async with info.context.db() as session:
+                if not (
+                    project_session := await session.scalar(
+                        select(models.ProjectSession).filter_by(id=node_id)
+                    )
+                ):
+                    raise NotFound(f"Unknown user: {id}")
+            return to_gql_project_session(project_session)
         raise NotFound(f"Unknown node type: {type_name}")
 
     @strawberry.field

phoenix/server/api/types/ExperimentRun.py

@@ -20,7 +20,7 @@ from phoenix.server.api.types.pagination import (
     CursorString,
     connection_from_list,
 )
-from phoenix.server.api.types.Trace import Trace
+from phoenix.server.api.types.Trace import Trace, to_gql_trace
 
 if TYPE_CHECKING:
     from phoenix.server.api.types.DatasetExample import DatasetExample
@@ -61,11 +61,10 @@ class ExperimentRun(Node):
     async def trace(self, info: Info) -> Optional[Trace]:
         if not self.trace_id:
             return None
-        dataloader = info.context.data_loaders.trace_row_ids
+        dataloader = info.context.data_loaders.trace_by_trace_ids
         if (trace := await dataloader.load(self.trace_id)) is None:
             return None
-        trace_rowid, project_rowid = trace
-        return Trace(id_attr=trace_rowid, trace_id=self.trace_id, project_rowid=project_rowid)
+        return to_gql_trace(trace)
 
     @strawberry.field
     async def example(

phoenix/server/api/types/ExperimentRunAnnotation.py

@@ -8,7 +8,7 @@ from strawberry.scalars import JSON
 
 from phoenix.db import models
 from phoenix.server.api.types.AnnotatorKind import ExperimentRunAnnotatorKind
-from phoenix.server.api.types.Trace import Trace
+from phoenix.server.api.types.Trace import Trace, to_gql_trace
 
 
 @strawberry.type
@@ -29,11 +29,10 @@ class ExperimentRunAnnotation(Node):
     async def trace(self, info: Info) -> Optional[Trace]:
         if not self.trace_id:
             return None
-        dataloader = info.context.data_loaders.trace_row_ids
+        dataloader = info.context.data_loaders.trace_by_trace_ids
         if (trace := await dataloader.load(self.trace_id)) is None:
             return None
-        trace_row_id, project_row_id = trace
-        return Trace(id_attr=trace_row_id, trace_id=self.trace_id, project_rowid=project_row_id)
+        return to_gql_trace(trace)
 
 
 def to_gql_experiment_run_annotation(