arize-phoenix 10.0.4__py3-none-any.whl → 12.28.1__py3-none-any.whl

phoenix/server/api/auth.py

@@ -3,8 +3,10 @@ from typing import Any
 
 from strawberry import Info
 from strawberry.permission import BasePermission
+from typing_extensions import override
 
-from phoenix.server.api.exceptions import Unauthorized
+from phoenix.config import get_env_support_email
+from phoenix.server.api.exceptions import InsufficientStorage, Unauthorized
 from phoenix.server.bearer_auth import PhoenixUser
 
 
@@ -20,16 +22,45 @@ class IsNotReadOnly(Authorization):
         return not info.context.read_only
 
 
-class IsLocked(Authorization):
-    """
-    Disables mutations and subscriptions that create or update data but allows
-    queries and delete mutations.
+class IsNotViewer(Authorization):
+    message = "Viewers cannot perform this action"
+
+    def has_permission(self, source: Any, info: Info, **kwargs: Any) -> bool:
+        if not info.context.auth_enabled:
+            return True
+        return isinstance((user := info.context.user), PhoenixUser) and not user.is_viewer
+
+
+class IsLocked(BasePermission):
     """
+    Permission class that restricts data-modifying operations when insufficient storage.
 
-    message = "Operations that write or modify data are locked"
+    When database storage capacity is exceeded, this permission blocks mutations and
+    subscriptions that create or update data, while allowing queries and delete mutations
+    to continue. This prevents database overflow while maintaining read access and the
+    ability to free up space through deletions.
 
+    Raises:
+        InsufficientStorage: When storage capacity is exceeded and data operations
+            are temporarily disabled. The error includes guidance for resolution
+            and support contact information if configured.
+    """
+
+    @override
+    def on_unauthorized(self) -> None:
+        """Create user-friendly error message when storage operations are blocked."""
+        message = (
+            "Database operations are disabled due to insufficient storage. "
+            "Please delete old data or increase storage."
+        )
+        if support_email := get_env_support_email():
+            message += f" Need help? Contact us at {support_email}"
+        raise InsufficientStorage(message)
+
+    @override
     def has_permission(self, source: Any, info: Info, **kwargs: Any) -> bool:
-        return not info.context.locked
+        """Check if database operations are allowed based on storage capacity and lock status."""
+        return not (info.context.db.should_not_insert_or_update or info.context.locked)
 
 
 MSG_ADMIN_ONLY = "Only admin can perform this action"

phoenix/server/api/auth_messages.py

@@ -0,0 +1,46 @@
+# ruff: noqa: E501
+"""
+Authentication error and success message codes.
+
+These codes are used in authentication flows to safely communicate status
+to users via query parameters. Using codes instead of raw messages prevents
+social engineering and phishing attacks.
+
+The messages are passed to the frontend via window.Config to ensure a single
+source of truth between backend and frontend.
+"""
+
+from types import MappingProxyType
+from typing import Literal, Mapping, get_args
+
+# Error code type - used for type hints in redirect functions
+AuthErrorCode = Literal[
+    "unknown_idp",
+    "auth_failed",
+    "invalid_state",
+    "unsafe_return_url",
+    "oauth_error",
+    "no_oidc_support",
+    "missing_email_scope",
+    "email_in_use",
+    "sign_in_not_allowed",
+]
+
+# Error messages - passed to frontend via window.Config.authErrorMessages
+# Backend generates these codes when redirecting users after OAuth errors
+AUTH_ERROR_MESSAGES: Mapping[AuthErrorCode, str] = MappingProxyType(
+    {
+        "unknown_idp": "Unknown identity provider.",
+        "auth_failed": "Authentication failed. Please contact your administrator.",
+        "invalid_state": "Invalid authentication state. Please try again.",
+        "unsafe_return_url": "Invalid return URL. Please try again.",
+        "oauth_error": "Authentication failed. Please try again.",
+        "no_oidc_support": "Your identity provider does not appear to support OpenID Connect. Please contact your administrator.",
+        "missing_email_scope": "Please ensure your identity provider is configured to use the 'email' scope.",
+        "email_in_use": "An account with this email already exists.",
+        "sign_in_not_allowed": "Sign in is not allowed. Please contact your administrator.",
+    }
+)
+
+# Runtime assertion to ensure AUTH_ERROR_MESSAGES keys match AuthErrorCode Literal values
+assert set(AUTH_ERROR_MESSAGES.keys()) == set(get_args(AuthErrorCode))

phoenix/server/api/context.py

@@ -15,19 +15,29 @@ from phoenix.auth import (
 from phoenix.core.model_schema import Model
 from phoenix.db import models
 from phoenix.server.api.dataloaders import (
+    AnnotationConfigsByProjectDataLoader,
     AnnotationSummaryDataLoader,
+    AverageExperimentRepeatedRunGroupLatencyDataLoader,
     AverageExperimentRunLatencyDataLoader,
     CacheForDataLoaders,
+    DatasetDatasetSplitsDataLoader,
     DatasetExampleRevisionsDataLoader,
+    DatasetExamplesAndVersionsByExperimentRunDataLoader,
     DatasetExampleSpansDataLoader,
+    DatasetExampleSplitsDataLoader,
     DocumentEvaluationsDataLoader,
     DocumentEvaluationSummaryDataLoader,
     DocumentRetrievalMetricsDataLoader,
     ExperimentAnnotationSummaryDataLoader,
+    ExperimentDatasetSplitsDataLoader,
     ExperimentErrorRatesDataLoader,
+    ExperimentRepeatedRunGroupAnnotationSummariesDataLoader,
+    ExperimentRepeatedRunGroupsDataLoader,
     ExperimentRunAnnotations,
     ExperimentRunCountsDataLoader,
+    ExperimentRunsByExperimentAndExampleDataLoader,
     ExperimentSequenceNumberDataLoader,
+    LastUsedTimesByGenerativeModelIdDataLoader,
     LatencyMsQuantileDataLoader,
     MinStartOrMaxEndTimeDataLoader,
     NumChildSpansDataLoader,
@@ -36,6 +46,7 @@ from phoenix.server.api.dataloaders import (
     ProjectIdsByTraceRetentionPolicyIdDataLoader,
     PromptVersionSequenceNumberDataLoader,
     RecordCountDataLoader,
+    SessionAnnotationsBySessionDataLoader,
     SessionIODataLoader,
     SessionNumTracesDataLoader,
     SessionNumTracesWithErrorDataLoader,
@@ -43,18 +54,35 @@ from phoenix.server.api.dataloaders import (
     SessionTraceLatencyMsQuantileDataLoader,
     SpanAnnotationsDataLoader,
     SpanByIdDataLoader,
+    SpanCostBySpanDataLoader,
+    SpanCostDetailsBySpanCostDataLoader,
+    SpanCostDetailSummaryEntriesByGenerativeModelDataLoader,
+    SpanCostDetailSummaryEntriesByProjectSessionDataLoader,
+    SpanCostDetailSummaryEntriesBySpanDataLoader,
+    SpanCostDetailSummaryEntriesByTraceDataLoader,
+    SpanCostSummaryByExperimentDataLoader,
+    SpanCostSummaryByExperimentRepeatedRunGroupDataLoader,
+    SpanCostSummaryByExperimentRunDataLoader,
+    SpanCostSummaryByGenerativeModelDataLoader,
+    SpanCostSummaryByProjectDataLoader,
+    SpanCostSummaryByProjectSessionDataLoader,
+    SpanCostSummaryByTraceDataLoader,
     SpanDatasetExamplesDataLoader,
     SpanDescendantsDataLoader,
     SpanProjectsDataLoader,
     TableFieldsDataLoader,
     TokenCountDataLoader,
+    TokenPricesByModelDataLoader,
+    TraceAnnotationsByTraceDataLoader,
     TraceByTraceIdsDataLoader,
     TraceRetentionPolicyIdByProjectIdDataLoader,
     TraceRootSpansDataLoader,
     UserRolesDataLoader,
     UsersDataLoader,
 )
+from phoenix.server.api.dataloaders.dataset_labels import DatasetLabelsDataLoader
 from phoenix.server.bearer_auth import PhoenixUser
+from phoenix.server.daemons.span_cost_calculator import SpanCostCalculator
 from phoenix.server.dml_event import DmlEvent
 from phoenix.server.email.types import EmailSender
 from phoenix.server.types import (
@@ -68,47 +96,107 @@ from phoenix.server.types import (
 
 @dataclass
 class DataLoaders:
+    annotation_configs_by_project: AnnotationConfigsByProjectDataLoader
+    annotation_summaries: AnnotationSummaryDataLoader
+    average_experiment_repeated_run_group_latency: (
+        AverageExperimentRepeatedRunGroupLatencyDataLoader
+    )
     average_experiment_run_latency: AverageExperimentRunLatencyDataLoader
+    dataset_example_fields: TableFieldsDataLoader
     dataset_example_revisions: DatasetExampleRevisionsDataLoader
     dataset_example_spans: DatasetExampleSpansDataLoader
+    dataset_labels: DatasetLabelsDataLoader
+    dataset_label_fields: TableFieldsDataLoader
+    dataset_dataset_splits: DatasetDatasetSplitsDataLoader
+    dataset_examples_and_versions_by_experiment_run: (
+        DatasetExamplesAndVersionsByExperimentRunDataLoader
+    )
+    dataset_example_splits: DatasetExampleSplitsDataLoader
+    dataset_fields: TableFieldsDataLoader
+    dataset_split_fields: TableFieldsDataLoader
+    dataset_version_fields: TableFieldsDataLoader
+    document_annotation_fields: TableFieldsDataLoader
     document_evaluation_summaries: DocumentEvaluationSummaryDataLoader
     document_evaluations: DocumentEvaluationsDataLoader
     document_retrieval_metrics: DocumentRetrievalMetricsDataLoader
-    annotation_summaries: AnnotationSummaryDataLoader
     experiment_annotation_summaries: ExperimentAnnotationSummaryDataLoader
+    experiment_dataset_splits: ExperimentDatasetSplitsDataLoader
     experiment_error_rates: ExperimentErrorRatesDataLoader
+    experiment_fields: TableFieldsDataLoader
+    experiment_repeated_run_group_annotation_summaries: (
+        ExperimentRepeatedRunGroupAnnotationSummariesDataLoader
+    )
+    experiment_repeated_run_groups: ExperimentRepeatedRunGroupsDataLoader
+    experiment_run_annotation_fields: TableFieldsDataLoader
     experiment_run_annotations: ExperimentRunAnnotations
     experiment_run_counts: ExperimentRunCountsDataLoader
+    experiment_run_fields: TableFieldsDataLoader
+    experiment_runs_by_experiment_and_example: ExperimentRunsByExperimentAndExampleDataLoader
     experiment_sequence_number: ExperimentSequenceNumberDataLoader
+    generative_model_fields: TableFieldsDataLoader
+    last_used_times_by_generative_model_id: LastUsedTimesByGenerativeModelIdDataLoader
     latency_ms_quantile: LatencyMsQuantileDataLoader
     min_start_or_max_end_times: MinStartOrMaxEndTimeDataLoader
     num_child_spans: NumChildSpansDataLoader
     num_spans_per_trace: NumSpansPerTraceDataLoader
+    project_by_name: ProjectByNameDataLoader
     project_fields: TableFieldsDataLoader
+    project_trace_retention_policy_fields: TableFieldsDataLoader
     projects_by_trace_retention_policy_id: ProjectIdsByTraceRetentionPolicyIdDataLoader
+    prompt_fields: TableFieldsDataLoader
+    prompt_label_fields: TableFieldsDataLoader
     prompt_version_sequence_number: PromptVersionSequenceNumberDataLoader
+    prompt_version_tag_fields: TableFieldsDataLoader
+    project_session_annotation_fields: TableFieldsDataLoader
+    project_session_fields: TableFieldsDataLoader
     record_counts: RecordCountDataLoader
+    session_annotations_by_session: SessionAnnotationsBySessionDataLoader
     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_annotation_fields: TableFieldsDataLoader
     span_annotations: SpanAnnotationsDataLoader
     span_by_id: SpanByIdDataLoader
+    span_cost_by_span: SpanCostBySpanDataLoader
+    span_cost_detail_fields: TableFieldsDataLoader
+    span_cost_detail_summary_entries_by_generative_model: (
+        SpanCostDetailSummaryEntriesByGenerativeModelDataLoader
+    )
+    span_cost_detail_summary_entries_by_project_session: (
+        SpanCostDetailSummaryEntriesByProjectSessionDataLoader
+    )
+    span_cost_detail_summary_entries_by_span: SpanCostDetailSummaryEntriesBySpanDataLoader
+    span_cost_detail_summary_entries_by_trace: SpanCostDetailSummaryEntriesByTraceDataLoader
+    span_cost_details_by_span_cost: SpanCostDetailsBySpanCostDataLoader
+    span_cost_fields: TableFieldsDataLoader
+    span_cost_summary_by_experiment: SpanCostSummaryByExperimentDataLoader
+    span_cost_summary_by_experiment_repeated_run_group: (
+        SpanCostSummaryByExperimentRepeatedRunGroupDataLoader
+    )
+    span_cost_summary_by_experiment_run: SpanCostSummaryByExperimentRunDataLoader
+    span_cost_summary_by_generative_model: SpanCostSummaryByGenerativeModelDataLoader
+    span_cost_summary_by_project: SpanCostSummaryByProjectDataLoader
+    span_cost_summary_by_project_session: SpanCostSummaryByProjectSessionDataLoader
+    span_cost_summary_by_trace: SpanCostSummaryByTraceDataLoader
     span_dataset_examples: SpanDatasetExamplesDataLoader
     span_descendants: SpanDescendantsDataLoader
     span_fields: TableFieldsDataLoader
     span_projects: SpanProjectsDataLoader
     token_counts: TokenCountDataLoader
+    token_prices_by_model: TokenPricesByModelDataLoader
+    trace_annotation_fields: TableFieldsDataLoader
+    trace_annotations_by_trace: TraceAnnotationsByTraceDataLoader
     trace_by_trace_ids: TraceByTraceIdsDataLoader
     trace_fields: TableFieldsDataLoader
     trace_retention_policy_id_by_project_id: TraceRetentionPolicyIdByProjectIdDataLoader
-    project_trace_retention_policy_fields: TableFieldsDataLoader
     trace_root_spans: TraceRootSpansDataLoader
-    project_by_name: ProjectByNameDataLoader
-    users: UsersDataLoader
     user_roles: UserRolesDataLoader
+    user_api_key_fields: TableFieldsDataLoader
+    user_fields: TableFieldsDataLoader
+    users: UsersDataLoader
 
 
 class _NoOp:
@@ -123,6 +211,7 @@ class Context(BaseContext):
     cache_for_dataloaders: Optional[CacheForDataLoaders]
     model: Model
     export_path: Path
+    span_cost_calculator: SpanCostCalculator
     last_updated_at: CanGetLastUpdatedAt = _NoOp()
     event_queue: CanPutItem[DmlEvent] = _NoOp()
     corpus: Optional[Model] = None
@@ -181,3 +270,10 @@ class Context(BaseContext):
     @cached_property
     def user(self) -> PhoenixUser:
         return cast(PhoenixUser, self.get_request().user)
+
+    @cached_property
+    def user_id(self) -> Optional[int]:
+        try:
+            return int(self.user.identity)
+        except Exception:
+            return None

phoenix/server/api/dataloaders/__init__.py

@@ -1,9 +1,23 @@
 from dataclasses import dataclass, field
 
+from phoenix.server.api.dataloaders.span_cost_detail_summary_entries_by_project_session import (
+    SpanCostDetailSummaryEntriesByProjectSessionDataLoader,
+)
+
+from .annotation_configs_by_project import AnnotationConfigsByProjectDataLoader
 from .annotation_summaries import AnnotationSummaryCache, AnnotationSummaryDataLoader
+from .average_experiment_repeated_run_group_latency import (
+    AverageExperimentRepeatedRunGroupLatencyDataLoader,
+)
 from .average_experiment_run_latency import AverageExperimentRunLatencyDataLoader
+from .dataset_dataset_splits import DatasetDatasetSplitsDataLoader
 from .dataset_example_revisions import DatasetExampleRevisionsDataLoader
 from .dataset_example_spans import DatasetExampleSpansDataLoader
+from .dataset_example_splits import DatasetExampleSplitsDataLoader
+from .dataset_examples_and_versions_by_experiment_run import (
+    DatasetExamplesAndVersionsByExperimentRunDataLoader,
+)
+from .dataset_labels import DatasetLabelsDataLoader
 from .document_evaluation_summaries import (
     DocumentEvaluationSummaryCache,
     DocumentEvaluationSummaryDataLoader,
@@ -11,10 +25,19 @@ from .document_evaluation_summaries import (
 from .document_evaluations import DocumentEvaluationsDataLoader
 from .document_retrieval_metrics import DocumentRetrievalMetricsDataLoader
 from .experiment_annotation_summaries import ExperimentAnnotationSummaryDataLoader
+from .experiment_dataset_splits import ExperimentDatasetSplitsDataLoader
 from .experiment_error_rates import ExperimentErrorRatesDataLoader
+from .experiment_repeated_run_group_annotation_summaries import (
+    ExperimentRepeatedRunGroupAnnotationSummariesDataLoader,
+)
+from .experiment_repeated_run_groups import ExperimentRepeatedRunGroupsDataLoader
 from .experiment_run_annotations import ExperimentRunAnnotations
 from .experiment_run_counts import ExperimentRunCountsDataLoader
+from .experiment_runs_by_experiment_and_example import (
+    ExperimentRunsByExperimentAndExampleDataLoader,
+)
 from .experiment_sequence_number import ExperimentSequenceNumberDataLoader
+from .last_used_times_by_generative_model_id import LastUsedTimesByGenerativeModelIdDataLoader
 from .latency_ms_quantile import LatencyMsQuantileCache, LatencyMsQuantileDataLoader
 from .min_start_or_max_end_times import MinStartOrMaxEndTimeCache, MinStartOrMaxEndTimeDataLoader
 from .num_child_spans import NumChildSpansDataLoader
@@ -23,6 +46,7 @@ from .project_by_name import ProjectByNameDataLoader
 from .project_ids_by_trace_retention_policy_id import ProjectIdsByTraceRetentionPolicyIdDataLoader
 from .prompt_version_sequence_number import PromptVersionSequenceNumberDataLoader
 from .record_counts import RecordCountCache, RecordCountDataLoader
+from .session_annotations_by_session import SessionAnnotationsBySessionDataLoader
 from .session_io import SessionIODataLoader
 from .session_num_traces import SessionNumTracesDataLoader
 from .session_num_traces_with_error import SessionNumTracesWithErrorDataLoader
@@ -30,11 +54,30 @@ from .session_token_usages import SessionTokenUsagesDataLoader
 from .session_trace_latency_ms_quantile import SessionTraceLatencyMsQuantileDataLoader
 from .span_annotations import SpanAnnotationsDataLoader
 from .span_by_id import SpanByIdDataLoader
+from .span_cost_by_span import SpanCostBySpanDataLoader
+from .span_cost_detail_summary_entries_by_generative_model import (
+    SpanCostDetailSummaryEntriesByGenerativeModelDataLoader,
+)
+from .span_cost_detail_summary_entries_by_span import SpanCostDetailSummaryEntriesBySpanDataLoader
+from .span_cost_detail_summary_entries_by_trace import SpanCostDetailSummaryEntriesByTraceDataLoader
+from .span_cost_details_by_span_cost import SpanCostDetailsBySpanCostDataLoader
+from .span_cost_summary_by_experiment import SpanCostSummaryByExperimentDataLoader
+from .span_cost_summary_by_experiment_repeated_run_group import (
+    SpanCostSummaryByExperimentRepeatedRunGroupDataLoader,
+)
+from .span_cost_summary_by_experiment_run import SpanCostSummaryByExperimentRunDataLoader
+from .span_cost_summary_by_generative_model import SpanCostSummaryByGenerativeModelDataLoader
+from .span_cost_summary_by_project import SpanCostSummaryByProjectDataLoader, SpanCostSummaryCache
+from .span_cost_summary_by_project_session import SpanCostSummaryByProjectSessionDataLoader
+from .span_cost_summary_by_trace import SpanCostSummaryByTraceDataLoader
+from .span_costs import SpanCostsDataLoader
 from .span_dataset_examples import SpanDatasetExamplesDataLoader
 from .span_descendants import SpanDescendantsDataLoader
 from .span_projects import SpanProjectsDataLoader
 from .table_fields import TableFieldsDataLoader
 from .token_counts import TokenCountCache, TokenCountDataLoader
+from .token_prices_by_model import TokenPricesByModelDataLoader
+from .trace_annotations_by_trace import TraceAnnotationsByTraceDataLoader
 from .trace_by_trace_ids import TraceByTraceIdsDataLoader
 from .trace_retention_policy_id_by_project_id import TraceRetentionPolicyIdByProjectIdDataLoader
 from .trace_root_spans import TraceRootSpansDataLoader
@@ -42,44 +85,72 @@ from .user_roles import UserRolesDataLoader
 from .users import UsersDataLoader
 
 __all__ = [
-    "CacheForDataLoaders",
+    "AnnotationConfigsByProjectDataLoader",
+    "AnnotationSummaryDataLoader",
+    "AverageExperimentRepeatedRunGroupLatencyDataLoader",
     "AverageExperimentRunLatencyDataLoader",
+    "CacheForDataLoaders",
+    "DatasetDatasetSplitsDataLoader",
     "DatasetExampleRevisionsDataLoader",
     "DatasetExampleSpansDataLoader",
+    "DatasetExamplesAndVersionsByExperimentRunDataLoader",
+    "DatasetExampleSplitsDataLoader",
+    "DatasetLabelsDataLoader",
+    "ExperimentDatasetSplitsDataLoader",
     "DocumentEvaluationSummaryDataLoader",
     "DocumentEvaluationsDataLoader",
     "DocumentRetrievalMetricsDataLoader",
-    "AnnotationSummaryDataLoader",
     "ExperimentAnnotationSummaryDataLoader",
     "ExperimentErrorRatesDataLoader",
+    "ExperimentRepeatedRunGroupsDataLoader",
+    "ExperimentRepeatedRunGroupAnnotationSummariesDataLoader",
     "ExperimentRunAnnotations",
     "ExperimentRunCountsDataLoader",
+    "ExperimentRunsByExperimentAndExampleDataLoader",
     "ExperimentSequenceNumberDataLoader",
+    "LastUsedTimesByGenerativeModelIdDataLoader",
     "LatencyMsQuantileDataLoader",
     "MinStartOrMaxEndTimeDataLoader",
     "NumChildSpansDataLoader",
     "NumSpansPerTraceDataLoader",
+    "ProjectByNameDataLoader",
     "ProjectIdsByTraceRetentionPolicyIdDataLoader",
     "PromptVersionSequenceNumberDataLoader",
     "RecordCountDataLoader",
+    "SessionAnnotationsBySessionDataLoader",
     "SessionIODataLoader",
     "SessionNumTracesDataLoader",
     "SessionNumTracesWithErrorDataLoader",
     "SessionTokenUsagesDataLoader",
     "SessionTraceLatencyMsQuantileDataLoader",
+    "SpanAnnotationsDataLoader",
     "SpanByIdDataLoader",
+    "SpanCostBySpanDataLoader",
+    "SpanCostDetailSummaryEntriesByGenerativeModelDataLoader",
+    "SpanCostDetailSummaryEntriesByProjectSessionDataLoader",
+    "SpanCostDetailSummaryEntriesBySpanDataLoader",
+    "SpanCostDetailSummaryEntriesByTraceDataLoader",
+    "SpanCostDetailsBySpanCostDataLoader",
+    "SpanCostSummaryByExperimentDataLoader",
+    "SpanCostSummaryByExperimentRepeatedRunGroupDataLoader",
+    "SpanCostSummaryByExperimentRunDataLoader",
+    "SpanCostSummaryByGenerativeModelDataLoader",
+    "SpanCostSummaryByProjectDataLoader",
+    "SpanCostSummaryByProjectSessionDataLoader",
+    "SpanCostSummaryByTraceDataLoader",
+    "SpanCostsDataLoader",
     "SpanDatasetExamplesDataLoader",
     "SpanDescendantsDataLoader",
     "SpanProjectsDataLoader",
     "TableFieldsDataLoader",
     "TokenCountDataLoader",
+    "TokenPricesByModelDataLoader",
+    "TraceAnnotationsByTraceDataLoader",
     "TraceByTraceIdsDataLoader",
     "TraceRetentionPolicyIdByProjectIdDataLoader",
     "TraceRootSpansDataLoader",
-    "ProjectByNameDataLoader",
-    "SpanAnnotationsDataLoader",
-    "UsersDataLoader",
     "UserRolesDataLoader",
+    "UsersDataLoader",
 ]
 
 
@@ -103,3 +174,6 @@ class CacheForDataLoaders:
     token_count: TokenCountCache = field(
         default_factory=TokenCountCache,
     )
+    token_cost: SpanCostSummaryCache = field(
+        default_factory=SpanCostSummaryCache,
+    )

phoenix/server/api/dataloaders/annotation_configs_by_project.py

@@ -0,0 +1,31 @@
+from collections import defaultdict
+
+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
+
+ProjectId: TypeAlias = int
+Key: TypeAlias = ProjectId
+Result: TypeAlias = tuple[models.AnnotationConfig, ...]
+
+
+class AnnotationConfigsByProjectDataLoader(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.ProjectAnnotationConfig.project_id, models.AnnotationConfig)
+            .join_from(models.ProjectAnnotationConfig, models.AnnotationConfig)
+            .where(models.ProjectAnnotationConfig.project_id.in_(keys))
+        )
+        results: defaultdict[Key, list[models.AnnotationConfig]] = defaultdict(list)
+        async with self._db() as session:
+            data = await session.stream(stmt)
+            async for id_, config in data:
+                results[id_].append(config)
+        return [tuple(results[k]) for k in keys]

phoenix/server/api/dataloaders/annotation_summaries.py

@@ -13,6 +13,7 @@ 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.session_filters import get_filtered_session_rowids_subquery
 from phoenix.server.types import DbSessionFactory
 from phoenix.trace.dsl import SpanFilter
 
@@ -20,27 +21,41 @@ Kind: TypeAlias = Literal["span", "trace"]
 ProjectRowId: TypeAlias = int
 TimeInterval: TypeAlias = tuple[Optional[datetime], Optional[datetime]]
 FilterCondition: TypeAlias = Optional[str]
+SessionFilterCondition: TypeAlias = Optional[str]
 AnnotationName: TypeAlias = str
 
-Segment: TypeAlias = tuple[Kind, ProjectRowId, TimeInterval, FilterCondition]
+Segment: TypeAlias = tuple[
+    Kind,
+    ProjectRowId,
+    TimeInterval,
+    FilterCondition,
+    SessionFilterCondition,
+]
 Param: TypeAlias = AnnotationName
 
-Key: TypeAlias = tuple[Kind, ProjectRowId, Optional[TimeRange], FilterCondition, AnnotationName]
+Key: TypeAlias = tuple[
+    Kind,
+    ProjectRowId,
+    Optional[TimeRange],
+    FilterCondition,
+    SessionFilterCondition,
+    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
+    kind, project_rowid, time_range, filter_condition, session_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
+    return (kind, project_rowid, interval, filter_condition, session_filter_condition), eval_name
 
 
 _Section: TypeAlias = tuple[ProjectRowId, AnnotationName, Kind]
-_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition]
+_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition, SessionFilterCondition]
 
 
 class AnnotationSummaryCache(
@@ -61,8 +76,21 @@ class AnnotationSummaryCache(
                 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)
+        (
+            (
+                kind,
+                project_rowid,
+                interval,
+                filter_condition,
+                session_filter_condition,
+            ),
+            annotation_name,
+        ) = _cache_key_fn(key)
+        return (project_rowid, annotation_name, kind), (
+            interval,
+            filter_condition,
+            session_filter_condition,
+        )
 
 
 class AnnotationSummaryDataLoader(DataLoader[Key, Result]):
@@ -102,7 +130,9 @@ def _get_stmt(
     segment: Segment,
     *annotation_names: Param,
 ) -> Select[Any]:
-    kind, project_rowid, (start_time, end_time), filter_condition = segment
+    kind, project_rowid, (start_time, end_time), filter_condition, session_filter_condition = (
+        segment
+    )
 
     annotation_model: Union[Type[models.SpanAnnotation], Type[models.TraceAnnotation]]
     entity_model: Union[Type[models.Span], Type[models.Trace]]
@@ -144,6 +174,19 @@ def _get_stmt(
         entity_count_query = entity_count_query.where(
             cast(Type[models.Trace], entity_model).project_rowid == project_rowid
         )
+    else:
+        assert_never(kind)
+
+    if session_filter_condition:
+        filtered_session_rowids = get_filtered_session_rowids_subquery(
+            session_filter_condition=session_filter_condition,
+            project_rowids=[project_rowid],
+            start_time=start_time,
+            end_time=end_time,
+        )
+        entity_count_query = entity_count_query.where(
+            models.Trace.project_session_rowid.in_(filtered_session_rowids)
+        )
 
     entity_count_query = entity_count_query.where(
         or_(score_column.is_not(None), label_column.is_not(None))
@@ -186,6 +229,15 @@ def _get_stmt(
     else:
         assert_never(kind)
 
+    if session_filter_condition:
+        filtered_session_rowids = get_filtered_session_rowids_subquery(
+            session_filter_condition=session_filter_condition,
+            project_rowids=[project_rowid],
+            start_time=start_time,
+            end_time=end_time,
+        )
+        base_stmt = base_stmt.where(models.Trace.project_session_rowid.in_(filtered_session_rowids))
+
     base_stmt = base_stmt.where(or_(score_column.is_not(None), label_column.is_not(None)))
     base_stmt = base_stmt.where(name_column.in_(annotation_names))