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

phoenix/db/migrations/data_migration_scripts/populate_project_sessions.py

@@ -0,0 +1,199 @@
+# /// script
+# dependencies = [
+#   "arize-phoenix[pg]",
+# ]
+# ///
+"""
+Populate the `project_sessions` table with data from the traces and spans tables.
+
+Environment variables.
+
+- `PHOENIX_SQL_DATABASE_URL` must be set to the database connection string.
+- (optional) Postgresql schema can be set via `PHOENIX_SQL_DATABASE_SCHEMA`.
+"""
+
+import os
+from datetime import datetime
+from time import perf_counter
+from typing import Any, Optional, Union
+
+import sqlean
+from openinference.semconv.trace import SpanAttributes
+from sqlalchemy import (
+    JSON,
+    Engine,
+    NullPool,
+    create_engine,
+    event,
+    func,
+    insert,
+    make_url,
+    select,
+    update,
+)
+from sqlalchemy.dialects import postgresql
+from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, sessionmaker
+
+from phoenix.config import ENV_PHOENIX_SQL_DATABASE_SCHEMA, get_env_database_connection_str
+from phoenix.db.engines import set_postgresql_search_path
+
+
+class JSONB(JSON):
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    __visit_name__ = "JSONB"
+
+
+@compiles(JSONB, "sqlite")
+def _(*args: Any, **kwargs: Any) -> str:
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    return "JSONB"
+
+
+JSON_ = (
+    JSON()
+    .with_variant(
+        postgresql.JSONB(),  # type: ignore
+        "postgresql",
+    )
+    .with_variant(
+        JSONB(),
+        "sqlite",
+    )
+)
+
+
+class Base(DeclarativeBase): ...
+
+
+class ProjectSession(Base):
+    __tablename__ = "project_sessions"
+    id: Mapped[int] = mapped_column(primary_key=True)
+    session_id: Mapped[str]
+    project_id: Mapped[int]
+    start_time: Mapped[datetime]
+    end_time: Mapped[datetime]
+
+
+class Trace(Base):
+    __tablename__ = "traces"
+    id: Mapped[int] = mapped_column(primary_key=True)
+    project_session_rowid: Mapped[Union[int, None]]
+    project_rowid: Mapped[int]
+    start_time: Mapped[datetime]
+    end_time: Mapped[datetime]
+
+
+class Span(Base):
+    __tablename__ = "spans"
+    id: Mapped[int] = mapped_column(primary_key=True)
+    trace_rowid: Mapped[int]
+    parent_id: Mapped[Optional[str]]
+    attributes: Mapped[dict[str, Any]] = mapped_column(JSON_, nullable=False)
+
+
+SESSION_ID = SpanAttributes.SESSION_ID.split(".")
+USER_ID = SpanAttributes.USER_ID.split(".")
+
+
+def populate_project_sessions(
+    engine: Engine,
+) -> None:
+    sessions_from_span = (
+        select(
+            Span.attributes[SESSION_ID].as_string().label("session_id"),
+            Trace.project_rowid.label("project_id"),
+            Trace.start_time.label("start_time"),
+            func.row_number()
+            .over(
+                partition_by=Span.attributes[SESSION_ID],
+                order_by=[Trace.start_time, Trace.id, Span.id],
+            )
+            .label("rank"),
+            func.max(Trace.end_time)
+            .over(partition_by=Span.attributes[SESSION_ID])
+            .label("end_time"),
+        )
+        .join_from(Span, Trace, Span.trace_rowid == Trace.id)
+        .where(Span.parent_id.is_(None))
+        .where(Span.attributes[SESSION_ID].as_string() != "")
+        .subquery()
+    )
+    sessions_for_trace_id = (
+        select(
+            Span.trace_rowid,
+            ProjectSession.id.label("project_session_rowid"),
+        )
+        .join_from(
+            Span,
+            ProjectSession,
+            Span.attributes[SESSION_ID].as_string() == ProjectSession.session_id,
+        )
+        .where(Span.parent_id.is_(None))
+        .where(Span.attributes[SESSION_ID].as_string() != "")
+        .subquery()
+    )
+    start_time = perf_counter()
+    with sessionmaker(engine).begin() as session:
+        session.execute(
+            insert(ProjectSession).from_select(
+                [
+                    "session_id",
+                    "project_id",
+                    "start_time",
+                    "end_time",
+                ],
+                select(
+                    sessions_from_span.c.session_id,
+                    sessions_from_span.c.project_id,
+                    sessions_from_span.c.start_time,
+                    sessions_from_span.c.end_time,
+                ).where(sessions_from_span.c.rank == 1),
+            )
+        )
+        session.execute(
+            (
+                update(Trace)
+                .values(project_session_rowid=sessions_for_trace_id.c.project_session_rowid)
+                .where(Trace.id == sessions_for_trace_id.c.trace_rowid)
+            )
+        )
+    elapsed_time = perf_counter() - start_time
+    print(f"✅ Populated project_sessions in {elapsed_time:.3f} seconds.")
+
+
+if __name__ == "__main__":
+    sql_database_url = make_url(get_env_database_connection_str())
+    print(f"Using database URL: {sql_database_url}")
+    ans = input("Is that correct? [y]/n: ")
+    if ans.lower().startswith("n"):
+        url = input("Please enter the correct database URL: ")
+        sql_database_url = make_url(url)
+    backend = sql_database_url.get_backend_name()
+    if backend == "sqlite":
+        file = sql_database_url.database
+        engine = create_engine(
+            url=sql_database_url.set(drivername="sqlite"),
+            creator=lambda: sqlean.connect(f"file:///{file}", uri=True),
+            poolclass=NullPool,
+            echo=True,
+        )
+    elif backend == "postgresql":
+        schema = os.getenv(ENV_PHOENIX_SQL_DATABASE_SCHEMA)
+        if schema:
+            print(f"Using schema: {schema}")
+        else:
+            print("No PostgreSQL schema set. (This is the default.)")
+        ans = input("Is that correct? [y]/n: ")
+        if ans.lower().startswith("n"):
+            schema = input("Please enter the correct schema: ")
+        engine = create_engine(
+            url=sql_database_url.set(drivername="postgresql+psycopg"),
+            poolclass=NullPool,
+            echo=True,
+        )
+        if schema:
+            event.listen(engine, "connect", set_postgresql_search_path(schema))
+    else:
+        raise ValueError(f"Unknown database backend: {backend}")
+    populate_project_sessions(engine)

phoenix/db/migrations/versions/4ded9e43755f_create_project_sessions_table.py

@@ -0,0 +1,66 @@
+"""create project_session table
+
+Revision ID: 4ded9e43755f
+Revises: cd164e83824f
+Create Date: 2024-10-08 22:53:24.539786
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "4ded9e43755f"
+down_revision: Union[str, None] = "cd164e83824f"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "project_sessions",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column("session_id", sa.String, unique=True, nullable=False),
+        sa.Column(
+            "project_id",
+            sa.Integer,
+            sa.ForeignKey("projects.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column(
+            "start_time",
+            sa.TIMESTAMP(timezone=True),
+            index=True,
+            nullable=False,
+        ),
+        sa.Column(
+            "end_time",
+            sa.TIMESTAMP(timezone=True),
+            index=True,
+            nullable=False,
+        ),
+    )
+    with op.batch_alter_table("traces") as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "project_session_rowid",
+                sa.Integer,
+                sa.ForeignKey("project_sessions.id", ondelete="CASCADE"),
+                nullable=True,
+            ),
+        )
+    op.create_index(
+        "ix_traces_project_session_rowid",
+        "traces",
+        ["project_session_rowid"],
+    )
+
+
+def downgrade() -> None:
+    op.drop_index("ix_traces_project_session_rowid")
+    with op.batch_alter_table("traces") as batch_op:
+        batch_op.drop_column("project_session_rowid")
+    op.drop_table("project_sessions")

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/metrics/wrappers.py

@@ -27,6 +27,8 @@ from sklearn import metrics as sk
 from sklearn.utils.multiclass import check_classification_targets
 from wrapt import PartialCallableObjectProxy
 
+from phoenix.config import SKLEARN_VERSION
+
 
 class Eval(PartialCallableObjectProxy, ABC):  # type: ignore
     def __call__(
@@ -232,5 +234,9 @@ class SkEval(Enum):
     r2_score = RegressionEval(sk.r2_score)
     recall_score = ClassificationEval(sk.recall_score)
     roc_auc_score = ScoredClassificationEval(sk.roc_auc_score)
-    root_mean_squared_error = RegressionEval(sk.mean_squared_error, squared=False)
+    root_mean_squared_error = (
+        RegressionEval(sk.mean_squared_error, squared=False)
+        if SKLEARN_VERSION < (1, 6)
+        else RegressionEval(sk.root_mean_squared_error)
+    )
     zero_one_loss = ClassificationEval(sk.zero_one_loss)

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]