arize-phoenix 10.14.0__py3-none-any.whl → 11.0.0__py3-none-any.whl

phoenix/db/models.py

@@ -1,12 +1,15 @@
+import re
 from datetime import datetime, timezone
 from typing import Any, Iterable, Literal, Optional, Sequence, TypedDict, cast
 
+import sqlalchemy as sa
 import sqlalchemy.sql as sql
 from openinference.semconv.trace import RerankerAttributes, SpanAttributes
 from sqlalchemy import (
     JSON,
     NUMERIC,
     TIMESTAMP,
+    Boolean,
     CheckConstraint,
     ColumnElement,
     Dialect,
@@ -52,6 +55,10 @@ from phoenix.db.types.annotation_configs import (
 )
 from phoenix.db.types.identifier import Identifier
 from phoenix.db.types.model_provider import ModelProvider
+from phoenix.db.types.token_price_customization import (
+    TokenPriceCustomization,
+    TokenPriceCustomizationParser,
+)
 from phoenix.db.types.trace_retention import TraceRetentionCronExpression, TraceRetentionRule
 from phoenix.server.api.helpers.prompts.models import (
     PromptInvocationParameters,
@@ -391,12 +398,49 @@ class _AnnotationConfig(TypeDecorator[AnnotationConfigType]):
         return AnnotationConfigModel.model_validate(value).root if value is not None else None
 
 
+class _TokenCustomization(TypeDecorator[TokenPriceCustomization]):
+    # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    cache_ok = True
+    impl = JSON
+
+    def process_bind_param(
+        self, value: Optional[TokenPriceCustomization], _: Dialect
+    ) -> Optional[dict[str, Any]]:
+        return value.model_dump() if value is not None else None
+
+    def process_result_value(
+        self, value: Optional[dict[str, Any]], _: Dialect
+    ) -> Optional[TokenPriceCustomization]:
+        return TokenPriceCustomizationParser.parse(value)
+
+
+class _RegexStr(TypeDecorator[re.Pattern[str]]):
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    cache_ok = True
+    impl = String
+
+    def process_bind_param(self, value: Optional[re.Pattern[str]], _: Dialect) -> Optional[str]:
+        if value is None:
+            return None
+        if not isinstance(value, re.Pattern):
+            raise TypeError(f"Expected a regex pattern, got {type(value)}")
+        pattern = value.pattern
+        if not isinstance(pattern, str):
+            raise ValueError(f"Expected a string, got {type(pattern)}")
+        return pattern
+
+    def process_result_value(self, value: Optional[str], _: Dialect) -> Optional[re.Pattern[str]]:
+        if value is None:
+            return None
+        return re.compile(value)
+
+
 class ExperimentRunOutput(TypedDict, total=False):
     task_output: Any
 
 
 class Base(DeclarativeBase):
-    id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    id: Mapped[int] = mapped_column(primary_key=True)
     # Enforce best practices for naming constraints
     # https://alembic.sqlalchemy.org/en/latest/naming.html#integration-of-naming-conventions-into-operations-autogenerate
     metadata = MetaData(
@@ -418,7 +462,6 @@ class Base(DeclarativeBase):
 
 class ProjectTraceRetentionPolicy(Base):
     __tablename__ = "project_trace_retention_policies"
-    id: Mapped[int] = mapped_column(Integer, primary_key=True)
     name: Mapped[str] = mapped_column(String, nullable=False)
     cron_expression: Mapped[TraceRetentionCronExpression] = mapped_column(
         _TraceRetentionCronExpression, nullable=False
@@ -528,6 +571,12 @@ class Trace(Base):
         primaryjoin="foreign(ExperimentRun.trace_id) == Trace.trace_id",
         back_populates="trace",
     )
+    span_costs: Mapped[list["SpanCost"]] = relationship(
+        "SpanCost",
+        back_populates="trace",
+        cascade="all, delete-orphan",
+        uselist=True,
+    )
     __table_args__ = (
         UniqueConstraint(
             "trace_id",
@@ -685,6 +734,7 @@ class Span(Base):
     span_annotations: Mapped[list["SpanAnnotation"]] = relationship(back_populates="span")
     document_annotations: Mapped[list["DocumentAnnotation"]] = relationship(back_populates="span")
     dataset_examples: Mapped[list["DatasetExample"]] = relationship(back_populates="span")
+    span_cost: Mapped[Optional["SpanCost"]] = relationship(back_populates="span")
 
     __table_args__ = (
         UniqueConstraint(
@@ -1303,9 +1353,86 @@ class ApiKey(Base):
     __table_args__ = (dict(sqlite_autoincrement=True),)
 
 
+CostType: TypeAlias = Literal["DEFAULT", "OVERRIDE"]
+
+
+class GenerativeModel(Base):
+    __tablename__ = "generative_models"
+    name: Mapped[str] = mapped_column(String, nullable=False)
+    provider: Mapped[str]
+    start_time: Mapped[Optional[datetime]] = mapped_column(UtcTimeStamp)
+    name_pattern: Mapped[re.Pattern[str]] = mapped_column(_RegexStr, nullable=False)
+    is_built_in: Mapped[bool] = mapped_column(
+        Boolean,
+        nullable=False,
+    )
+    created_at: Mapped[datetime] = mapped_column(
+        UtcTimeStamp,
+        server_default=func.now(),
+    )
+    updated_at: Mapped[datetime] = mapped_column(
+        UtcTimeStamp,
+        server_default=func.now(),
+        onupdate=func.now(),
+    )
+    deleted_at: Mapped[Optional[datetime]] = mapped_column(UtcTimeStamp)
+
+    token_prices: Mapped[list["TokenPrice"]] = relationship(
+        "TokenPrice",
+        back_populates="model",
+        cascade="all, delete-orphan",
+        uselist=True,
+    )
+
+    __table_args__ = (
+        Index(
+            "ix_generative_models_match_criteria",
+            "name_pattern",
+            "provider",
+            "is_built_in",
+            postgresql_where=sa.text("deleted_at IS NULL"),
+            sqlite_where=sa.text("deleted_at IS NULL"),
+            unique=True,
+        ),
+        Index(
+            "ix_generative_models_name_is_built_in",
+            "name",
+            "is_built_in",
+            postgresql_where=sa.text("deleted_at IS NULL"),
+            sqlite_where=sa.text("deleted_at IS NULL"),
+            unique=True,
+        ),
+    )
+
+
+class TokenPrice(Base):
+    __tablename__ = "token_prices"
+    model_id: Mapped[int] = mapped_column(
+        ForeignKey("generative_models.id", ondelete="CASCADE"),
+        nullable=False,
+        index=True,
+    )
+    token_type: Mapped[str]
+    is_prompt: Mapped[bool]
+    base_rate: Mapped[float]
+    customization: Mapped[TokenPriceCustomization] = mapped_column(_TokenCustomization)
+
+    model: Mapped["GenerativeModel"] = relationship(
+        "GenerativeModel",
+        back_populates="token_prices",
+    )
+
+    __table_args__ = (
+        UniqueConstraint(
+            "model_id",
+            "token_type",
+            "is_prompt",
+        ),
+    )
+
+
 class PromptLabel(Base):
     __tablename__ = "prompt_labels"
-
     name: Mapped[str] = mapped_column(String, unique=True, index=True, nullable=False)
     description: Mapped[Optional[str]]
     color: Mapped[str] = mapped_column(String, nullable=True)
@@ -1320,7 +1447,6 @@ class PromptLabel(Base):
 
 class Prompt(Base):
     __tablename__ = "prompts"
-
     source_prompt_id: Mapped[Optional[int]] = mapped_column(
         ForeignKey("prompts.id", ondelete="SET NULL"),
         index=True,
@@ -1358,7 +1484,6 @@ class Prompt(Base):
 
 class PromptPromptLabel(Base):
     __tablename__ = "prompts_prompt_labels"
-
     prompt_label_id: Mapped[int] = mapped_column(
         ForeignKey("prompt_labels.id", ondelete="CASCADE"),
         index=True,
@@ -1458,16 +1583,12 @@ class PromptVersionTag(Base):
 
 class AnnotationConfig(Base):
     __tablename__ = "annotation_configs"
-
-    id: Mapped[int] = mapped_column(primary_key=True)
     name: Mapped[str] = mapped_column(String, nullable=False, unique=True)
     config: Mapped[AnnotationConfigType] = mapped_column(_AnnotationConfig, nullable=False)
 
 
 class ProjectAnnotationConfig(Base):
     __tablename__ = "project_annotation_configs"
-
-    id: Mapped[int] = mapped_column(primary_key=True)
     project_id: Mapped[int] = mapped_column(
         ForeignKey("projects.id", ondelete="CASCADE"), nullable=False, index=True
     )
@@ -1476,3 +1597,140 @@ class ProjectAnnotationConfig(Base):
     )
 
     __table_args__ = (UniqueConstraint("project_id", "annotation_config_id"),)
+
+
+class SpanCost(Base):
+    __tablename__ = "span_costs"
+
+    span_rowid: Mapped[int] = mapped_column(
+        ForeignKey("spans.id", ondelete="CASCADE"),
+        nullable=False,
+    )
+    trace_rowid: Mapped[int] = mapped_column(
+        ForeignKey("traces.id", ondelete="CASCADE"),
+        nullable=False,
+    )
+    span_start_time: Mapped[datetime] = mapped_column(
+        UtcTimeStamp,
+        nullable=False,
+        index=True,
+    )
+    model_id: Mapped[Optional[int]] = mapped_column(
+        sa.Integer,
+        ForeignKey(
+            "generative_models.id",
+            ondelete="RESTRICT",
+        ),
+        nullable=True,
+    )
+    total_cost: Mapped[Optional[float]]
+    total_tokens: Mapped[Optional[float]]
+
+    @hybrid_property
+    def total_cost_per_token(self) -> Optional[float]:
+        return ((self.total_cost or 0) / self.total_tokens) if self.total_tokens else None
+
+    @total_cost_per_token.inplace.expression
+    @classmethod
+    def _total_cost_per_token_expression(cls) -> ColumnElement[Optional[float]]:
+        return sql.case(
+            (
+                sa.and_(cls.total_tokens.isnot(None), cls.total_tokens != 0),
+                cls.total_cost / cls.total_tokens,
+            )
+        )
+
+    prompt_cost: Mapped[Optional[float]]
+    prompt_tokens: Mapped[Optional[float]]
+
+    @hybrid_property
+    def prompt_cost_per_token(self) -> Optional[float]:
+        return ((self.prompt_cost or 0) / self.prompt_tokens) if self.prompt_tokens else None
+
+    @prompt_cost_per_token.inplace.expression
+    @classmethod
+    def _prompt_cost_per_token_expression(cls) -> ColumnElement[Optional[float]]:
+        return sql.case(
+            (
+                sa.and_(cls.prompt_tokens.isnot(None), cls.prompt_tokens != 0),
+                cls.prompt_cost / cls.prompt_tokens,
+            )
+        )
+
+    completion_cost: Mapped[Optional[float]]
+    completion_tokens: Mapped[Optional[float]]
+
+    @hybrid_property
+    def completion_cost_per_token(self) -> Optional[float]:
+        return (
+            ((self.completion_cost or 0) / self.completion_tokens)
+            if self.completion_tokens
+            else None
+        )
+
+    @completion_cost_per_token.inplace.expression
+    @classmethod
+    def _completion_cost_per_token_expression(cls) -> ColumnElement[Optional[float]]:
+        return sql.case(
+            (
+                sa.and_(cls.completion_tokens.isnot(None), cls.completion_tokens != 0),
+                cls.completion_cost / cls.completion_tokens,
+            )
+        )
+
+    span: Mapped["Span"] = relationship("Span", back_populates="span_cost")
+    trace: Mapped["Trace"] = relationship("Trace", back_populates="span_costs")
+    span_cost_details: Mapped[list["SpanCostDetail"]] = relationship(
+        "SpanCostDetail",
+        back_populates="span_cost",
+        cascade="all, delete-orphan",
+        uselist=True,
+    )
+
+    __table_args__ = (
+        Index(
+            "ix_span_costs_model_id_span_start_time",
+            "model_id",
+            "span_start_time",
+        ),
+    )
+
+    def append_detail(self, detail: "SpanCostDetail") -> None:
+        self.span_cost_details.append(detail)
+        if cost := detail.cost:
+            if detail.is_prompt:
+                self.prompt_cost = (self.prompt_cost or 0) + cost
+            else:
+                self.completion_cost = (self.completion_cost or 0) + cost
+            self.total_cost = (self.total_cost or 0) + cost
+        if tokens := detail.tokens:
+            if detail.is_prompt:
+                self.prompt_tokens = (self.prompt_tokens or 0) + tokens
+            else:
+                self.completion_tokens = (self.completion_tokens or 0) + tokens
+            self.total_tokens = (self.total_tokens or 0) + tokens
+
+
+class SpanCostDetail(Base):
+    __tablename__ = "span_cost_details"
+    span_cost_id: Mapped[int] = mapped_column(
+        ForeignKey("span_costs.id", ondelete="CASCADE"),
+        nullable=False,
+        index=True,
+    )
+    token_type: Mapped[str]
+    is_prompt: Mapped[bool]
+
+    cost: Mapped[Optional[float]]
+    tokens: Mapped[Optional[float]]
+    cost_per_token: Mapped[Optional[float]]
+
+    span_cost: Mapped["SpanCost"] = relationship("SpanCost", back_populates="span_cost_details")
+
+    __table_args__ = (
+        UniqueConstraint(
+            "span_cost_id",
+            "token_type",
+            "is_prompt",
+        ),
+    )

phoenix/db/types/model_provider.py

@@ -9,3 +9,4 @@ class ModelProvider(Enum):
     DEEPSEEK = "DEEPSEEK"
     XAI = "XAI"
     OLLAMA = "OLLAMA"
+    AWS = "AWS"

phoenix/db/types/token_price_customization.py

@@ -0,0 +1,29 @@
+from abc import ABC
+from typing import Any, Literal, Optional
+
+from pydantic import BaseModel, ValidationError
+
+
+class TokenPriceCustomization(BaseModel, ABC):
+    model_config = {"extra": "allow"}
+
+
+class ThresholdBasedTokenPriceCustomization(TokenPriceCustomization):
+    type: Literal["threshold_based"] = "threshold_based"
+    key: str
+    threshold: float
+    new_rate: float
+
+
+class TokenPriceCustomizationParser:
+    """Intended to be forward-compatible while maintaining the ability to round-trip."""
+
+    @staticmethod
+    def parse(data: Optional[dict[str, Any]]) -> Optional[TokenPriceCustomization]:
+        if not data:
+            return None
+        try:
+            return ThresholdBasedTokenPriceCustomization.model_validate(data)
+        except ValidationError:
+            pass
+        return TokenPriceCustomization.model_validate(data)

phoenix/server/api/context.py

@@ -28,6 +28,7 @@ from phoenix.server.api.dataloaders import (
     ExperimentRunAnnotations,
     ExperimentRunCountsDataLoader,
     ExperimentSequenceNumberDataLoader,
+    LastUsedTimesByGenerativeModelIdDataLoader,
     LatencyMsQuantileDataLoader,
     MinStartOrMaxEndTimeDataLoader,
     NumChildSpansDataLoader,
@@ -43,6 +44,18 @@ from phoenix.server.api.dataloaders import (
     SessionTraceLatencyMsQuantileDataLoader,
     SpanAnnotationsDataLoader,
     SpanByIdDataLoader,
+    SpanCostBySpanDataLoader,
+    SpanCostDetailsBySpanCostDataLoader,
+    SpanCostDetailSummaryEntriesByGenerativeModelDataLoader,
+    SpanCostDetailSummaryEntriesByProjectSessionDataLoader,
+    SpanCostDetailSummaryEntriesBySpanDataLoader,
+    SpanCostDetailSummaryEntriesByTraceDataLoader,
+    SpanCostSummaryByExperimentDataLoader,
+    SpanCostSummaryByExperimentRunDataLoader,
+    SpanCostSummaryByGenerativeModelDataLoader,
+    SpanCostSummaryByProjectDataLoader,
+    SpanCostSummaryByProjectSessionDataLoader,
+    SpanCostSummaryByTraceDataLoader,
     SpanDatasetExamplesDataLoader,
     SpanDescendantsDataLoader,
     SpanProjectsDataLoader,
@@ -55,6 +68,7 @@ from phoenix.server.api.dataloaders import (
     UsersDataLoader,
 )
 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,23 +82,26 @@ from phoenix.server.types import (
 
 @dataclass
 class DataLoaders:
+    annotation_summaries: AnnotationSummaryDataLoader
     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
+    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_version_sequence_number: PromptVersionSequenceNumberDataLoader
     record_counts: RecordCountDataLoader
@@ -96,6 +113,24 @@ class DataLoaders:
     session_trace_latency_ms_quantile: SessionTraceLatencyMsQuantileDataLoader
     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_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
@@ -104,11 +139,9 @@ class DataLoaders:
     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
+    users: UsersDataLoader
 
 
 class _NoOp:
@@ -123,6 +156,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

phoenix/server/api/dataloaders/__init__.py

@@ -1,5 +1,9 @@
 from dataclasses import dataclass, field
 
+from phoenix.server.api.dataloaders.span_cost_detail_summary_entries_by_project_session import (
+    SpanCostDetailSummaryEntriesByProjectSessionDataLoader,
+)
+
 from .annotation_summaries import AnnotationSummaryCache, AnnotationSummaryDataLoader
 from .average_experiment_run_latency import AverageExperimentRunLatencyDataLoader
 from .dataset_example_revisions import DatasetExampleRevisionsDataLoader
@@ -15,6 +19,7 @@ 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 .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
@@ -30,6 +35,20 @@ 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_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
@@ -42,23 +61,25 @@ from .user_roles import UserRolesDataLoader
 from .users import UsersDataLoader
 
 __all__ = [
-    "CacheForDataLoaders",
+    "AnnotationSummaryDataLoader",
     "AverageExperimentRunLatencyDataLoader",
+    "CacheForDataLoaders",
     "DatasetExampleRevisionsDataLoader",
     "DatasetExampleSpansDataLoader",
     "DocumentEvaluationSummaryDataLoader",
     "DocumentEvaluationsDataLoader",
     "DocumentRetrievalMetricsDataLoader",
-    "AnnotationSummaryDataLoader",
     "ExperimentAnnotationSummaryDataLoader",
     "ExperimentErrorRatesDataLoader",
     "ExperimentRunAnnotations",
     "ExperimentRunCountsDataLoader",
     "ExperimentSequenceNumberDataLoader",
+    "LastUsedTimesByGenerativeModelIdDataLoader",
     "LatencyMsQuantileDataLoader",
     "MinStartOrMaxEndTimeDataLoader",
     "NumChildSpansDataLoader",
     "NumSpansPerTraceDataLoader",
+    "ProjectByNameDataLoader",
     "ProjectIdsByTraceRetentionPolicyIdDataLoader",
     "PromptVersionSequenceNumberDataLoader",
     "RecordCountDataLoader",
@@ -67,7 +88,21 @@ __all__ = [
     "SessionNumTracesWithErrorDataLoader",
     "SessionTokenUsagesDataLoader",
     "SessionTraceLatencyMsQuantileDataLoader",
+    "SpanAnnotationsDataLoader",
     "SpanByIdDataLoader",
+    "SpanCostBySpanDataLoader",
+    "SpanCostDetailSummaryEntriesByGenerativeModelDataLoader",
+    "SpanCostDetailSummaryEntriesByProjectSessionDataLoader",
+    "SpanCostDetailSummaryEntriesBySpanDataLoader",
+    "SpanCostDetailSummaryEntriesByTraceDataLoader",
+    "SpanCostDetailsBySpanCostDataLoader",
+    "SpanCostSummaryByExperimentDataLoader",
+    "SpanCostSummaryByExperimentRunDataLoader",
+    "SpanCostSummaryByGenerativeModelDataLoader",
+    "SpanCostSummaryByProjectDataLoader",
+    "SpanCostSummaryByProjectSessionDataLoader",
+    "SpanCostSummaryByTraceDataLoader",
+    "SpanCostsDataLoader",
     "SpanDatasetExamplesDataLoader",
     "SpanDescendantsDataLoader",
     "SpanProjectsDataLoader",
@@ -76,10 +111,8 @@ __all__ = [
     "TraceByTraceIdsDataLoader",
     "TraceRetentionPolicyIdByProjectIdDataLoader",
     "TraceRootSpansDataLoader",
-    "ProjectByNameDataLoader",
-    "SpanAnnotationsDataLoader",
-    "UsersDataLoader",
     "UserRolesDataLoader",
+    "UsersDataLoader",
 ]
 
 
@@ -103,3 +136,6 @@ class CacheForDataLoaders:
     token_count: TokenCountCache = field(
         default_factory=TokenCountCache,
     )
+    token_cost: SpanCostSummaryCache = field(
+        default_factory=SpanCostSummaryCache,
+    )

phoenix/server/api/dataloaders/last_used_times_by_generative_model_id.py

@@ -0,0 +1,35 @@
+from datetime import datetime
+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
+
+GenerativeModelID: TypeAlias = int
+Key: TypeAlias = GenerativeModelID
+Result: TypeAlias = Optional[datetime]
+
+
+class LastUsedTimesByGenerativeModelIdDataLoader(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]:
+        async with self._db() as session:
+            last_used_times_by_model_id: dict[Key, Result] = {
+                model_id: last_used_time
+                async for model_id, last_used_time in await session.stream(
+                    select(
+                        models.SpanCost.model_id,
+                        func.max(models.SpanCost.span_start_time).label("last_used_time"),
+                    )
+                    .select_from(models.SpanCost)
+                    .where(models.SpanCost.model_id.in_(keys))
+                    .group_by(models.SpanCost.model_id)
+                )
+            }
+        return [last_used_times_by_model_id.get(model_id) for model_id in keys]

phoenix/server/api/dataloaders/span_cost_by_span.py

@@ -0,0 +1,24 @@
+from typing import 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
+
+SpanRowId: TypeAlias = int
+Key: TypeAlias = SpanRowId
+Result: TypeAlias = Optional[models.SpanCost]
+
+
+class SpanCostBySpanDataLoader(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.SpanCost).where(models.SpanCost.span_rowid.in_(keys))
+        async with self._db() as session:
+            result = {sc.span_rowid: sc async for sc in await session.stream_scalars(stmt)}
+        return list(map(result.get, keys))