arize-phoenix 3.16.1__py3-none-any.whl → 7.7.1__py3-none-any.whl

phoenix/pointcloud/umap_parameters.py

@@ -1,52 +1,53 @@
-from dataclasses import dataclass
-from typing import Any, Mapping, Optional
-
-DEFAULT_MIN_DIST = 0.0
-DEFAULT_N_NEIGHBORS = 30
-DEFAULT_N_SAMPLES = 500
-
-MIN_NEIGHBORS = 5
-MAX_NEIGHBORS = 100
-MIN_SAMPLES = 1
-MAX_SAMPLES = 1000
-MIN_MIN_DIST = 0.0
-MAX_MIN_DIST = 0.99
-
-
-@dataclass
-class UMAPParameters:
-    min_dist: float = DEFAULT_MIN_DIST
-    n_neighbors: int = DEFAULT_N_NEIGHBORS
-    n_samples: int = DEFAULT_N_SAMPLES
-
-    def __post_init__(self) -> None:
-        if not isinstance(self.min_dist, float) or not (
-            MIN_MIN_DIST <= self.min_dist <= MAX_MIN_DIST
-        ):
-            raise ValueError(
-                f"minDist must be float type, and between {MIN_MIN_DIST} and {MAX_MIN_DIST}"
-            )
-
-        if not isinstance(self.n_neighbors, int) or not (
-            MIN_NEIGHBORS <= self.n_neighbors <= MAX_NEIGHBORS
-        ):
-            raise ValueError(
-                f"nNeighbors must be int type, and between {MIN_NEIGHBORS} and {MAX_NEIGHBORS}"
-            )
-
-        if not isinstance(self.n_samples, int) or not (
-            MIN_SAMPLES <= self.n_samples <= MAX_SAMPLES
-        ):
-            raise ValueError(
-                f"nSamples must be int type, and between {MIN_SAMPLES} and {MAX_SAMPLES}"
-            )
-
-
-def get_umap_parameters(default_umap_parameters: Optional[Mapping[str, Any]]) -> UMAPParameters:
-    if not default_umap_parameters:
-        return UMAPParameters()
-    return UMAPParameters(
-        min_dist=float(default_umap_parameters.get("min_dist", DEFAULT_MIN_DIST)),
-        n_neighbors=int(default_umap_parameters.get("n_neighbors", DEFAULT_N_NEIGHBORS)),
-        n_samples=int(default_umap_parameters.get("n_samples", DEFAULT_N_SAMPLES)),
-    )
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any, Optional
+
+DEFAULT_MIN_DIST = 0.0
+DEFAULT_N_NEIGHBORS = 30
+DEFAULT_N_SAMPLES = 500
+
+MIN_NEIGHBORS = 5
+MAX_NEIGHBORS = 100
+MIN_SAMPLES = 1
+MAX_SAMPLES = 1000
+MIN_MIN_DIST = 0.0
+MAX_MIN_DIST = 0.99
+
+
+@dataclass
+class UMAPParameters:
+    min_dist: float = DEFAULT_MIN_DIST
+    n_neighbors: int = DEFAULT_N_NEIGHBORS
+    n_samples: int = DEFAULT_N_SAMPLES
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.min_dist, float) or not (
+            MIN_MIN_DIST <= self.min_dist <= MAX_MIN_DIST
+        ):
+            raise ValueError(
+                f"minDist must be float type, and between {MIN_MIN_DIST} and {MAX_MIN_DIST}"
+            )
+
+        if not isinstance(self.n_neighbors, int) or not (
+            MIN_NEIGHBORS <= self.n_neighbors <= MAX_NEIGHBORS
+        ):
+            raise ValueError(
+                f"nNeighbors must be int type, and between {MIN_NEIGHBORS} and {MAX_NEIGHBORS}"
+            )
+
+        if not isinstance(self.n_samples, int) or not (
+            MIN_SAMPLES <= self.n_samples <= MAX_SAMPLES
+        ):
+            raise ValueError(
+                f"nSamples must be int type, and between {MIN_SAMPLES} and {MAX_SAMPLES}"
+            )
+
+
+def get_umap_parameters(default_umap_parameters: Optional[Mapping[str, Any]]) -> UMAPParameters:
+    if not default_umap_parameters:
+        return UMAPParameters()
+    return UMAPParameters(
+        min_dist=float(default_umap_parameters.get("min_dist", DEFAULT_MIN_DIST)),
+        n_neighbors=int(default_umap_parameters.get("n_neighbors", DEFAULT_N_NEIGHBORS)),
+        n_samples=int(default_umap_parameters.get("n_samples", DEFAULT_N_SAMPLES)),
+    )

phoenix/server/api/README.md

@@ -0,0 +1,28 @@
+# Permission Matrix for GraphQL API
+
+## Mutations
+
+| Action                       | Admin | Member |
+|:-----------------------------|:-----:|:------:|
+| Create User                  |  Yes  |   No   |
+| Delete User                  |  Yes  |   No   |
+| Change Own Password          |  Yes  |  Yes   |
+| Change Other's Password      |  Yes  |   No   |
+| Change Own Username          |  Yes  |  Yes   |
+| Change Other's Username      |  Yes  |   No   |
+| Change Own Email             |  No   |   No   |
+| Change Other's Email         |  No   |   No   |
+| Create System API Keys       |  Yes  |   No   |
+| Delete System API Keys       |  Yes  |   No   |
+| Create Own User API Keys     |  Yes  |  Yes   |
+| Delete Own User API Keys     |  Yes  |  Yes   |
+| Delete Other's User API Keys |  Yes  |   No   |
+
+## Queries
+
+| Action                               | Admin | Member |
+|:-------------------------------------|:-----:|:------:|
+| List All System API Keys             |  Yes  |   No   |
+| List All User API Keys               |  Yes  |   No   |
+| List All Users                       |  Yes  |   No   |
+| Fetch Other User's Info, e.g. emails |  Yes  |   No   |

phoenix/server/api/auth.py

@@ -0,0 +1,44 @@
+from abc import ABC
+from typing import Any
+
+from strawberry import Info
+from strawberry.permission import BasePermission
+
+from phoenix.server.api.exceptions import Unauthorized
+from phoenix.server.bearer_auth import PhoenixUser
+
+
+class Authorization(BasePermission, ABC):
+    def on_unauthorized(self) -> None:
+        raise Unauthorized(self.message)
+
+
+class IsNotReadOnly(Authorization):
+    message = "Application is read-only"
+
+    def has_permission(self, source: Any, info: Info, **kwargs: Any) -> bool:
+        return not info.context.read_only
+
+
+class IsLocked(Authorization):
+    """
+    Disables mutations and subscriptions that create or update data but allows
+    queries and delete mutations.
+    """
+
+    message = "Operations that write or modify data are locked"
+
+    def has_permission(self, source: Any, info: Info, **kwargs: Any) -> bool:
+        return not info.context.locked
+
+
+MSG_ADMIN_ONLY = "Only admin can perform this action"
+
+
+class IsAdmin(Authorization):
+    message = MSG_ADMIN_ONLY
+
+    def has_permission(self, source: Any, info: Info, **kwargs: Any) -> bool:
+        if not info.context.auth_enabled:
+            return False
+        return isinstance((user := info.context.user), PhoenixUser) and user.is_admin

phoenix/server/api/context.py

@@ -1,20 +1,163 @@
+from asyncio import get_running_loop
 from dataclasses import dataclass
+from functools import cached_property, partial
 from pathlib import Path
-from typing import Optional, Union
+from typing import Any, Optional, cast
 
-from starlette.requests import Request
-from starlette.responses import Response
-from starlette.websockets import WebSocket
+from starlette.requests import Request as StarletteRequest
+from starlette.responses import Response as StarletteResponse
+from strawberry.fastapi import BaseContext
 
+from phoenix.auth import (
+    compute_password_hash,
+)
 from phoenix.core.model_schema import Model
-from phoenix.core.traces import Traces
+from phoenix.db import models
+from phoenix.server.api.dataloaders import (
+    AnnotationSummaryDataLoader,
+    AverageExperimentRunLatencyDataLoader,
+    CacheForDataLoaders,
+    DatasetExampleRevisionsDataLoader,
+    DatasetExampleSpansDataLoader,
+    DocumentEvaluationsDataLoader,
+    DocumentEvaluationSummaryDataLoader,
+    DocumentRetrievalMetricsDataLoader,
+    ExperimentAnnotationSummaryDataLoader,
+    ExperimentErrorRatesDataLoader,
+    ExperimentRunAnnotations,
+    ExperimentRunCountsDataLoader,
+    ExperimentSequenceNumberDataLoader,
+    LatencyMsQuantileDataLoader,
+    MinStartOrMaxEndTimeDataLoader,
+    ProjectByNameDataLoader,
+    RecordCountDataLoader,
+    SessionIODataLoader,
+    SessionNumTracesDataLoader,
+    SessionNumTracesWithErrorDataLoader,
+    SessionTokenUsagesDataLoader,
+    SessionTraceLatencyMsQuantileDataLoader,
+    SpanAnnotationsDataLoader,
+    SpanDatasetExamplesDataLoader,
+    SpanDescendantsDataLoader,
+    SpanProjectsDataLoader,
+    TokenCountDataLoader,
+    TraceByTraceIdsDataLoader,
+    TraceRootSpansDataLoader,
+    UserRolesDataLoader,
+    UsersDataLoader,
+)
+from phoenix.server.bearer_auth import PhoenixUser
+from phoenix.server.dml_event import DmlEvent
+from phoenix.server.types import (
+    CanGetLastUpdatedAt,
+    CanPutItem,
+    DbSessionFactory,
+    TokenStore,
+    UserId,
+)
 
 
 @dataclass
-class Context:
-    request: Union[Request, WebSocket]
-    response: Optional[Response]
+class DataLoaders:
+    average_experiment_run_latency: AverageExperimentRunLatencyDataLoader
+    dataset_example_revisions: DatasetExampleRevisionsDataLoader
+    dataset_example_spans: DatasetExampleSpansDataLoader
+    document_evaluation_summaries: DocumentEvaluationSummaryDataLoader
+    document_evaluations: DocumentEvaluationsDataLoader
+    document_retrieval_metrics: DocumentRetrievalMetricsDataLoader
+    annotation_summaries: AnnotationSummaryDataLoader
+    experiment_annotation_summaries: ExperimentAnnotationSummaryDataLoader
+    experiment_error_rates: ExperimentErrorRatesDataLoader
+    experiment_run_annotations: ExperimentRunAnnotations
+    experiment_run_counts: ExperimentRunCountsDataLoader
+    experiment_sequence_number: ExperimentSequenceNumberDataLoader
+    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_by_trace_ids: TraceByTraceIdsDataLoader
+    trace_root_spans: TraceRootSpansDataLoader
+    project_by_name: ProjectByNameDataLoader
+    users: UsersDataLoader
+    user_roles: UserRolesDataLoader
+
+
+class _NoOp:
+    def get(self, *args: Any, **kwargs: Any) -> Any: ...
+    def put(self, *args: Any, **kwargs: Any) -> Any: ...
+
+
+@dataclass
+class Context(BaseContext):
+    db: DbSessionFactory
+    data_loaders: DataLoaders
+    cache_for_dataloaders: Optional[CacheForDataLoaders]
     model: Model
     export_path: Path
+    last_updated_at: CanGetLastUpdatedAt = _NoOp()
+    event_queue: CanPutItem[DmlEvent] = _NoOp()
     corpus: Optional[Model] = None
-    traces: Optional[Traces] = None
+    read_only: bool = False
+    locked: bool = False
+    auth_enabled: bool = False
+    secret: Optional[str] = None
+    token_store: Optional[TokenStore] = None
+
+    def get_secret(self) -> str:
+        """A type-safe way to get the application secret. Throws an error if the secret is not set.
+
+        Returns:
+            str: the phoenix secret
+        """
+        if self.secret is None:
+            raise ValueError(
+                "Application secret not set."
+                " Please set the PHOENIX_SECRET environment variable and re-deploy the application."
+            )
+        return self.secret
+
+    def get_request(self) -> StarletteRequest:
+        """
+        A type-safe way to get the request object. Throws an error if the request is not set.
+        """
+        if not isinstance(request := self.request, StarletteRequest):
+            raise ValueError("no request is set")
+        return request
+
+    def get_response(self) -> StarletteResponse:
+        """
+        A type-safe way to get the response object. Throws an error if the response is not set.
+        """
+        if (response := self.response) is None:
+            raise ValueError("no response is set")
+        return response
+
+    async def is_valid_password(self, password: str, user: models.User) -> bool:
+        return (
+            (hash_ := user.password_hash) is not None
+            and (salt := user.password_salt) is not None
+            and hash_ == await self.hash_password(password, salt)
+        )
+
+    @staticmethod
+    async def hash_password(password: str, salt: bytes) -> bytes:
+        compute = partial(compute_password_hash, password=password, salt=salt)
+        return await get_running_loop().run_in_executor(None, compute)
+
+    async def log_out(self, user_id: int) -> None:
+        assert self.token_store is not None
+        await self.token_store.log_out(UserId(user_id))
+
+    @cached_property
+    def user(self) -> PhoenixUser:
+        return cast(PhoenixUser, self.get_request().user)

phoenix/server/api/dataloaders/__init__.py

@@ -0,0 +1,91 @@
+from dataclasses import dataclass, field
+
+from .annotation_summaries import AnnotationSummaryCache, AnnotationSummaryDataLoader
+from .average_experiment_run_latency import AverageExperimentRunLatencyDataLoader
+from .dataset_example_revisions import DatasetExampleRevisionsDataLoader
+from .dataset_example_spans import DatasetExampleSpansDataLoader
+from .document_evaluation_summaries import (
+    DocumentEvaluationSummaryCache,
+    DocumentEvaluationSummaryDataLoader,
+)
+from .document_evaluations import DocumentEvaluationsDataLoader
+from .document_retrieval_metrics import DocumentRetrievalMetricsDataLoader
+from .experiment_annotation_summaries import ExperimentAnnotationSummaryDataLoader
+from .experiment_error_rates import ExperimentErrorRatesDataLoader
+from .experiment_run_annotations import ExperimentRunAnnotations
+from .experiment_run_counts import ExperimentRunCountsDataLoader
+from .experiment_sequence_number import ExperimentSequenceNumberDataLoader
+from .latency_ms_quantile import LatencyMsQuantileCache, LatencyMsQuantileDataLoader
+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_by_trace_ids import TraceByTraceIdsDataLoader
+from .trace_root_spans import TraceRootSpansDataLoader
+from .user_roles import UserRolesDataLoader
+from .users import UsersDataLoader
+
+__all__ = [
+    "CacheForDataLoaders",
+    "AverageExperimentRunLatencyDataLoader",
+    "DatasetExampleRevisionsDataLoader",
+    "DatasetExampleSpansDataLoader",
+    "DocumentEvaluationSummaryDataLoader",
+    "DocumentEvaluationsDataLoader",
+    "DocumentRetrievalMetricsDataLoader",
+    "AnnotationSummaryDataLoader",
+    "ExperimentAnnotationSummaryDataLoader",
+    "ExperimentErrorRatesDataLoader",
+    "ExperimentRunAnnotations",
+    "ExperimentRunCountsDataLoader",
+    "ExperimentSequenceNumberDataLoader",
+    "LatencyMsQuantileDataLoader",
+    "MinStartOrMaxEndTimeDataLoader",
+    "RecordCountDataLoader",
+    "SessionIODataLoader",
+    "SessionNumTracesDataLoader",
+    "SessionNumTracesWithErrorDataLoader",
+    "SessionTokenUsagesDataLoader",
+    "SessionTraceLatencyMsQuantileDataLoader",
+    "SpanDatasetExamplesDataLoader",
+    "SpanDescendantsDataLoader",
+    "SpanProjectsDataLoader",
+    "TokenCountDataLoader",
+    "TraceByTraceIdsDataLoader",
+    "TraceRootSpansDataLoader",
+    "ProjectByNameDataLoader",
+    "SpanAnnotationsDataLoader",
+    "UsersDataLoader",
+    "UserRolesDataLoader",
+]
+
+
+@dataclass(frozen=True)
+class CacheForDataLoaders:
+    document_evaluation_summary: DocumentEvaluationSummaryCache = field(
+        default_factory=DocumentEvaluationSummaryCache,
+    )
+    annotation_summary: AnnotationSummaryCache = field(
+        default_factory=AnnotationSummaryCache,
+    )
+    latency_ms_quantile: LatencyMsQuantileCache = field(
+        default_factory=LatencyMsQuantileCache,
+    )
+    min_start_or_max_end_time: MinStartOrMaxEndTimeCache = field(
+        default_factory=MinStartOrMaxEndTimeCache,
+    )
+    record_count: RecordCountCache = field(
+        default_factory=RecordCountCache,
+    )
+    token_count: TokenCountCache = field(
+        default_factory=TokenCountCache,
+    )

phoenix/server/api/dataloaders/annotation_summaries.py

@@ -0,0 +1,139 @@
+from collections import defaultdict
+from datetime import datetime
+from typing import Any, Literal, Optional
+
+import pandas as pd
+from aioitertools.itertools import groupby
+from cachetools import LFUCache, TTLCache
+from sqlalchemy import Select, func, or_, select
+from strawberry.dataloader import AbstractCache, DataLoader
+from typing_extensions import TypeAlias, assert_never
+
+from phoenix.db import models
+from phoenix.server.api.dataloaders.cache import TwoTierCache
+from phoenix.server.api.input_types.TimeRange import TimeRange
+from phoenix.server.api.types.AnnotationSummary import AnnotationSummary
+from phoenix.server.types import DbSessionFactory
+from phoenix.trace.dsl import SpanFilter
+
+Kind: TypeAlias = Literal["span", "trace"]
+ProjectRowId: TypeAlias = int
+TimeInterval: TypeAlias = tuple[Optional[datetime], Optional[datetime]]
+FilterCondition: TypeAlias = Optional[str]
+AnnotationName: TypeAlias = str
+
+Segment: TypeAlias = tuple[Kind, ProjectRowId, TimeInterval, FilterCondition]
+Param: TypeAlias = AnnotationName
+
+Key: TypeAlias = tuple[Kind, ProjectRowId, Optional[TimeRange], FilterCondition, AnnotationName]
+Result: TypeAlias = Optional[AnnotationSummary]
+ResultPosition: TypeAlias = int
+DEFAULT_VALUE: Result = None
+
+
+def _cache_key_fn(key: Key) -> tuple[Segment, Param]:
+    kind, project_rowid, time_range, filter_condition, eval_name = key
+    interval = (
+        (time_range.start, time_range.end) if isinstance(time_range, TimeRange) else (None, None)
+    )
+    return (kind, project_rowid, interval, filter_condition), eval_name
+
+
+_Section: TypeAlias = tuple[ProjectRowId, AnnotationName, Kind]
+_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition]
+
+
+class AnnotationSummaryCache(
+    TwoTierCache[Key, Result, _Section, _SubKey],
+):
+    def __init__(self) -> None:
+        super().__init__(
+            # TTL=3600 (1-hour) because time intervals are always moving forward, but
+            # interval endpoints are rounded down to the hour by the UI, so anything
+            # older than an hour most likely won't be a cache-hit anyway.
+            main_cache=TTLCache(maxsize=64 * 32 * 2, ttl=3600),
+            sub_cache_factory=lambda: LFUCache(maxsize=2 * 2),
+        )
+
+    def invalidate_project(self, project_rowid: ProjectRowId) -> None:
+        for section in self._cache.keys():
+            if section[0] == project_rowid:
+                del self._cache[section]
+
+    def _cache_key(self, key: Key) -> tuple[_Section, _SubKey]:
+        (kind, project_rowid, interval, filter_condition), annotation_name = _cache_key_fn(key)
+        return (project_rowid, annotation_name, kind), (interval, filter_condition)
+
+
+class AnnotationSummaryDataLoader(DataLoader[Key, Result]):
+    def __init__(
+        self,
+        db: DbSessionFactory,
+        cache_map: Optional[AbstractCache[Key, Result]] = None,
+    ) -> None:
+        super().__init__(
+            load_fn=self._load_fn,
+            cache_key_fn=_cache_key_fn,
+            cache_map=cache_map,
+        )
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        results: list[Result] = [DEFAULT_VALUE] * len(keys)
+        arguments: defaultdict[
+            Segment,
+            defaultdict[Param, list[ResultPosition]],
+        ] = defaultdict(lambda: defaultdict(list))
+        for position, key in enumerate(keys):
+            segment, param = _cache_key_fn(key)
+            arguments[segment][param].append(position)
+        for segment, params in arguments.items():
+            stmt = _get_stmt(segment, *params.keys())
+            async with self._db() as session:
+                data = await session.stream(stmt)
+                async for annotation_name, group in groupby(data, lambda row: row.name):
+                    summary = AnnotationSummary(pd.DataFrame(group))
+                    for position in params[annotation_name]:
+                        results[position] = summary
+        return results
+
+
+def _get_stmt(
+    segment: Segment,
+    *annotation_names: Param,
+) -> Select[Any]:
+    kind, project_rowid, (start_time, end_time), filter_condition = segment
+    stmt = select()
+    if kind == "span":
+        msa = models.SpanAnnotation
+        name_column, label_column, score_column = msa.name, msa.label, msa.score
+        time_column = models.Span.start_time
+        stmt = stmt.join(models.Span).join_from(models.Span, models.Trace)
+        if filter_condition:
+            sf = SpanFilter(filter_condition)
+            stmt = sf(stmt)
+    elif kind == "trace":
+        mta = models.TraceAnnotation
+        name_column, label_column, score_column = mta.name, mta.label, mta.score
+        time_column = models.Trace.start_time
+        stmt = stmt.join(models.Trace)
+    else:
+        assert_never(kind)
+    stmt = stmt.add_columns(
+        name_column,
+        label_column,
+        func.count().label("record_count"),
+        func.count(label_column).label("label_count"),
+        func.count(score_column).label("score_count"),
+        func.sum(score_column).label("score_sum"),
+    )
+    stmt = stmt.group_by(name_column, label_column)
+    stmt = stmt.order_by(name_column, label_column)
+    stmt = stmt.where(models.Trace.project_rowid == project_rowid)
+    stmt = stmt.where(or_(score_column.is_not(None), label_column.is_not(None)))
+    stmt = stmt.where(name_column.in_(annotation_names))
+    if start_time:
+        stmt = stmt.where(start_time <= time_column)
+    if end_time:
+        stmt = stmt.where(time_column < end_time)
+    return stmt

phoenix/server/api/dataloaders/average_experiment_run_latency.py

@@ -0,0 +1,54 @@
+from typing import Optional
+
+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
+
+ExperimentID: TypeAlias = int
+RunLatency: TypeAlias = Optional[float]
+Key: TypeAlias = ExperimentID
+Result: TypeAlias = RunLatency
+
+
+class AverageExperimentRunLatencyDataLoader(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]:
+        experiment_ids = keys
+        resolved_experiment_ids = (
+            select(models.Experiment.id)
+            .where(models.Experiment.id.in_(set(experiment_ids)))
+            .subquery()
+        )
+        query = (
+            select(
+                resolved_experiment_ids.c.id,
+                func.avg(
+                    func.extract("epoch", models.ExperimentRun.end_time)
+                    - func.extract("epoch", models.ExperimentRun.start_time)
+                ),
+            )
+            .outerjoin_from(
+                from_=resolved_experiment_ids,
+                target=models.ExperimentRun,
+                onclause=resolved_experiment_ids.c.id == models.ExperimentRun.experiment_id,
+            )
+            .group_by(resolved_experiment_ids.c.id)
+        )
+        async with self._db() as session:
+            avg_latencies = {
+                experiment_id: avg_latency
+                async for experiment_id, avg_latency in await session.stream(query)
+            }
+        return [
+            avg_latencies.get(experiment_id, ValueError(f"Unknown experiment: {experiment_id}"))
+            for experiment_id in keys
+        ]

phoenix/server/api/dataloaders/cache/__init__.py

@@ -0,0 +1,3 @@
+from phoenix.server.api.dataloaders.cache.two_tier_cache import TwoTierCache
+
+__all__ = ("TwoTierCache",)