arize-phoenix 10.15.0__py3-none-any.whl → 11.1.0__py3-none-any.whl

phoenix/server/api/dataloaders/span_costs.py

@@ -0,0 +1,35 @@
+from typing import Optional
+
+from sqlalchemy import select
+from sqlalchemy.orm import joinedload, load_only
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+SpanID: TypeAlias = int
+Key: TypeAlias = SpanID
+Result: TypeAlias = Optional[models.SpanCost]
+
+
+class SpanCostsDataLoader(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]:
+        span_ids = list(set(keys))
+        async with self._db() as session:
+            costs = {
+                span.id: span.span_cost
+                async for span in await session.stream_scalars(
+                    select(models.Span)
+                    .where(models.Span.id.in_(span_ids))
+                    .options(
+                        load_only(models.Span.id),
+                        joinedload(models.Span.span_cost),
+                    )
+                )
+            }
+        return [costs.get(span_id) for span_id in keys]

phoenix/server/api/dataloaders/types.py

@@ -0,0 +1,29 @@
+from dataclasses import dataclass, field
+from functools import cached_property
+from typing import Optional
+
+
+@dataclass(frozen=True)
+class CostBreakdown:
+    tokens: Optional[float] = None
+    cost: Optional[float] = None
+
+    @cached_property
+    def cost_per_token(self) -> Optional[float]:
+        if self.tokens and self.cost:
+            return self.cost / self.tokens
+        return None
+
+
+@dataclass(frozen=True)
+class SpanCostSummary:
+    prompt: CostBreakdown = field(default_factory=CostBreakdown)
+    completion: CostBreakdown = field(default_factory=CostBreakdown)
+    total: CostBreakdown = field(default_factory=CostBreakdown)
+
+
+@dataclass(frozen=True)
+class SpanCostDetailSummaryEntry:
+    token_type: str
+    is_prompt: bool
+    value: CostBreakdown = field(default_factory=CostBreakdown)

phoenix/server/api/helpers/playground_clients.py

@@ -463,6 +463,35 @@ class OpenAIBaseStreamingClient(PlaygroundStreamingClient):
         yield LLM_TOKEN_COUNT_COMPLETION, usage.completion_tokens
         yield LLM_TOKEN_COUNT_TOTAL, usage.total_tokens
 
+        if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details is not None:
+            prompt_details = usage.prompt_tokens_details
+            if (
+                hasattr(prompt_details, "cached_tokens")
+                and prompt_details.cached_tokens is not None
+            ):
+                yield LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ, prompt_details.cached_tokens
+            if hasattr(prompt_details, "audio_tokens") and prompt_details.audio_tokens is not None:
+                yield LLM_TOKEN_COUNT_PROMPT_DETAILS_AUDIO, prompt_details.audio_tokens
+
+        if (
+            hasattr(usage, "completion_tokens_details")
+            and usage.completion_tokens_details is not None
+        ):
+            completion_details = usage.completion_tokens_details
+            if (
+                hasattr(completion_details, "reasoning_tokens")
+                and completion_details.reasoning_tokens is not None
+            ):
+                yield (
+                    LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING,
+                    completion_details.reasoning_tokens,
+                )
+            if (
+                hasattr(completion_details, "audio_tokens")
+                and completion_details.audio_tokens is not None
+            ):
+                yield LLM_TOKEN_COUNT_COMPLETION_DETAILS_AUDIO, completion_details.audio_tokens
+
 
 def _get_credential_value(
     credentials: Optional[list[PlaygroundClientCredential]], env_var_name: str
@@ -1115,13 +1144,20 @@ class OpenAIStreamingClient(OpenAIBaseStreamingClient):
     provider_key=GenerativeProviderKey.OPENAI,
     model_names=[
         "o1",
+        "o1-pro",
         "o1-2024-12-17",
+        "o1-pro-2025-03-19",
         "o1-mini",
         "o1-mini-2024-09-12",
         "o1-preview",
         "o1-preview-2024-09-12",
+        "o3",
+        "o3-pro",
+        "o3-2025-04-16",
         "o3-mini",
         "o3-mini-2025-01-31",
+        "o4-mini",
+        "o4-mini-2025-04-16",
     ],
 )
 class OpenAIReasoningStreamingClient(OpenAIStreamingClient):
@@ -1258,6 +1294,35 @@ class OpenAIReasoningStreamingClient(OpenAIStreamingClient):
         yield LLM_TOKEN_COUNT_COMPLETION, usage.completion_tokens
         yield LLM_TOKEN_COUNT_TOTAL, usage.total_tokens
 
+        if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details is not None:
+            prompt_details = usage.prompt_tokens_details
+            if (
+                hasattr(prompt_details, "cached_tokens")
+                and prompt_details.cached_tokens is not None
+            ):
+                yield LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ, prompt_details.cached_tokens
+            if hasattr(prompt_details, "audio_tokens") and prompt_details.audio_tokens is not None:
+                yield LLM_TOKEN_COUNT_PROMPT_DETAILS_AUDIO, prompt_details.audio_tokens
+
+        if (
+            hasattr(usage, "completion_tokens_details")
+            and usage.completion_tokens_details is not None
+        ):
+            completion_details = usage.completion_tokens_details
+            if (
+                hasattr(completion_details, "reasoning_tokens")
+                and completion_details.reasoning_tokens is not None
+            ):
+                yield (
+                    LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING,
+                    completion_details.reasoning_tokens,
+                )
+            if (
+                hasattr(completion_details, "audio_tokens")
+                and completion_details.audio_tokens is not None
+            ):
+                yield LLM_TOKEN_COUNT_COMPLETION_DETAILS_AUDIO, completion_details.audio_tokens
+
 
 @register_llm_client(
     provider_key=GenerativeProviderKey.AZURE_OPENAI,
@@ -1315,12 +1380,6 @@ class AzureOpenAIStreamingClient(OpenAIBaseStreamingClient):
     provider_key=GenerativeProviderKey.ANTHROPIC,
     model_names=[
         PROVIDER_DEFAULT,
-        "claude-sonnet-4-0",
-        "claude-sonnet-4-20250514",
-        "claude-opus-4-0",
-        "claude-opus-4-20250514",
-        "claude-3-7-sonnet-latest",
-        "claude-3-7-sonnet-20250219",
         "claude-3-5-sonnet-latest",
         "claude-3-5-haiku-latest",
         "claude-3-5-sonnet-20241022",
@@ -1421,15 +1480,34 @@ class AnthropicStreamingClient(PlaygroundStreamingClient):
         async with await throttled_stream(**anthropic_params) as stream:
             async for event in stream:
                 if isinstance(event, anthropic_types.RawMessageStartEvent):
-                    self._attributes.update(
-                        {LLM_TOKEN_COUNT_PROMPT: event.message.usage.input_tokens}
-                    )
+                    usage = event.message.usage
+
+                    token_counts: dict[str, Any] = {}
+                    if prompt_tokens := (
+                        (usage.input_tokens or 0)
+                        + (getattr(usage, "cache_creation_input_tokens", 0) or 0)
+                        + (getattr(usage, "cache_read_input_tokens", 0) or 0)
+                    ):
+                        token_counts[LLM_TOKEN_COUNT_PROMPT] = prompt_tokens
+                    if cache_creation_tokens := getattr(usage, "cache_creation_input_tokens", None):
+                        if cache_creation_tokens is not None:
+                            token_counts[LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE] = (
+                                cache_creation_tokens
+                            )
+                    self._attributes.update(token_counts)
                 elif isinstance(event, anthropic_streaming.TextEvent):
                     yield TextChunk(content=event.text)
                 elif isinstance(event, anthropic_streaming.MessageStopEvent):
-                    self._attributes.update(
-                        {LLM_TOKEN_COUNT_COMPLETION: event.message.usage.output_tokens}
-                    )
+                    usage = event.message.usage
+                    output_token_counts: dict[str, Any] = {}
+                    if usage.output_tokens:
+                        output_token_counts[LLM_TOKEN_COUNT_COMPLETION] = usage.output_tokens
+                    if cache_read_tokens := getattr(usage, "cache_read_input_tokens", None):
+                        if cache_read_tokens is not None:
+                            output_token_counts[LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ] = (
+                                cache_read_tokens
+                            )
+                    self._attributes.update(output_token_counts)
                 elif (
                     isinstance(event, anthropic_streaming.ContentBlockStopEvent)
                     and event.content_block.type == "tool_use"
@@ -1514,6 +1592,10 @@ class AnthropicStreamingClient(PlaygroundStreamingClient):
 @register_llm_client(
     provider_key=GenerativeProviderKey.ANTHROPIC,
     model_names=[
+        "claude-sonnet-4-0",
+        "claude-sonnet-4-20250514",
+        "claude-opus-4-0",
+        "claude-opus-4-20250514",
         "claude-3-7-sonnet-latest",
         "claude-3-7-sonnet-20250219",
     ],
@@ -1698,6 +1780,15 @@ LLM_SYSTEM = SpanAttributes.LLM_SYSTEM
 LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT
 LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION
 LLM_TOKEN_COUNT_TOTAL = SpanAttributes.LLM_TOKEN_COUNT_TOTAL
+LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ = SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ
+LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE = (
+    SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE
+)
+LLM_TOKEN_COUNT_PROMPT_DETAILS_AUDIO = SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_AUDIO
+LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING = (
+    SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING
+)
+LLM_TOKEN_COUNT_COMPLETION_DETAILS_AUDIO = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_AUDIO
 
 
 class _HttpxClient(wrapt.ObjectProxy):  # type: ignore

phoenix/server/api/input_types/ProjectSessionSort.py

@@ -13,6 +13,7 @@ class ProjectSessionColumn(Enum):
     endTime = auto()
     tokenCountTotal = auto()
     numTraces = auto()
+    costTotal = auto()
 
     @property
     def data_type(self) -> CursorSortColumnDataType:
@@ -20,6 +21,8 @@ class ProjectSessionColumn(Enum):
             return CursorSortColumnDataType.INT
         if self is ProjectSessionColumn.startTime or self is ProjectSessionColumn.endTime:
             return CursorSortColumnDataType.DATETIME
+        if self is ProjectSessionColumn.costTotal:
+            return CursorSortColumnDataType.FLOAT
         assert_never(self)
 
 

phoenix/server/api/input_types/SpanSort.py

@@ -27,6 +27,7 @@ class SpanColumn(Enum):
     cumulativeTokenCountTotal = auto()
     cumulativeTokenCountPrompt = auto()
     cumulativeTokenCountCompletion = auto()
+    tokenCostTotal = auto()
 
     @property
     def column_name(self) -> str:
@@ -56,6 +57,8 @@ class SpanColumn(Enum):
             expr = models.Span.cumulative_llm_token_count_prompt
         elif self is SpanColumn.cumulativeTokenCountCompletion:
             expr = models.Span.cumulative_llm_token_count_completion
+        elif self is SpanColumn.tokenCostTotal:
+            expr = models.SpanCost.total_cost
         else:
             assert_never(self)
         return expr.label(self.column_name)
@@ -73,12 +76,25 @@ class SpanColumn(Enum):
             or self is SpanColumn.tokenCountTotal
             or self is SpanColumn.tokenCountPrompt
             or self is SpanColumn.tokenCountCompletion
+            or self is SpanColumn.tokenCostTotal
         ):
             return CursorSortColumnDataType.FLOAT
         if self is SpanColumn.startTime or self is SpanColumn.endTime:
             return CursorSortColumnDataType.DATETIME
         assert_never(self)
 
+    def join_tables(self, stmt: Select[Any]) -> Select[Any]:
+        """
+        If needed, joins tables required for the sort column.
+        """
+        if self is SpanColumn.tokenCostTotal:
+            return stmt.join_from(
+                models.Span,
+                models.SpanCost,
+                onclause=models.SpanCost.span_rowid == models.Span.id,
+            )
+        return stmt
+
 
 @strawberry.enum
 class EvalAttr(Enum):
@@ -140,6 +156,7 @@ class SpanSort:
     def update_orm_expr(self, stmt: Select[Any]) -> SpanSortConfig:
         if (col := self.col) and not self.eval_result_key:
             expr = col.orm_expression
+            stmt = col.join_tables(stmt)
             stmt = stmt.add_columns(expr)
             if self.dir == SortDir.desc:
                 expr = desc(expr)

phoenix/server/api/mutations/__init__.py

@@ -8,6 +8,7 @@ from phoenix.server.api.mutations.chat_mutations import (
 from phoenix.server.api.mutations.dataset_mutations import DatasetMutationMixin
 from phoenix.server.api.mutations.experiment_mutations import ExperimentMutationMixin
 from phoenix.server.api.mutations.export_events_mutations import ExportEventsMutationMixin
+from phoenix.server.api.mutations.model_mutations import ModelMutationMixin
 from phoenix.server.api.mutations.project_mutations import ProjectMutationMixin
 from phoenix.server.api.mutations.project_trace_retention_policy_mutations import (
     ProjectTraceRetentionPolicyMutationMixin,
@@ -29,6 +30,7 @@ class Mutation(
     DatasetMutationMixin,
     ExperimentMutationMixin,
     ExportEventsMutationMixin,
+    ModelMutationMixin,
     ProjectMutationMixin,
     ProjectTraceRetentionPolicyMutationMixin,
     PromptMutationMixin,

phoenix/server/api/mutations/chat_mutations.py

@@ -1,4 +1,5 @@
 import asyncio
+import logging
 from dataclasses import asdict, field
 from datetime import datetime, timezone
 from itertools import chain, islice
@@ -73,6 +74,8 @@ from phoenix.utilities.template_formatters import (
     TemplateFormatter,
 )
 
+logger = logging.getLogger(__name__)
+
 initialize_playground_clients()
 
 ChatCompletionMessage = tuple[ChatCompletionMessageRole, str, Optional[str], Optional[List[Any]]]
@@ -450,6 +453,19 @@ class ChatCompletionMutationMixin:
             session.add(trace)
             session.add(span)
             await session.flush()
+            try:
+                span_cost = info.context.span_cost_calculator.calculate_cost(
+                    start_time=span.start_time,
+                    attributes=span.attributes,
+                )
+            except Exception as e:
+                logger.exception(f"Failed to calculate cost for span {span.id}: {e}")
+                span_cost = None
+            if span_cost:
+                span_cost.span_rowid = span.id
+                span_cost.trace_rowid = trace.id
+                session.add(span_cost)
+                await session.flush()
 
         gql_span = Span(span_rowid=span.id, db_span=span)
 
@@ -605,5 +621,6 @@ TOOL_CALL_FUNCTION_ARGUMENTS_JSON = ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUME
 TOOL_JSON_SCHEMA = ToolAttributes.TOOL_JSON_SCHEMA
 PROMPT_TEMPLATE_VARIABLES = SpanAttributes.LLM_PROMPT_TEMPLATE_VARIABLES
 
+LLM_PROVIDER = SpanAttributes.LLM_PROVIDER
 
 PLAYGROUND_PROJECT_NAME = "playground"

phoenix/server/api/mutations/model_mutations.py

@@ -0,0 +1,208 @@
+import re
+from datetime import datetime, timezone
+from typing import Optional
+
+import sqlalchemy as sa
+import strawberry
+from sqlalchemy import delete
+from sqlalchemy.exc import IntegrityError as PostgreSQLIntegrityError
+from sqlalchemy.orm import joinedload
+from sqlean.dbapi2 import IntegrityError as SQLiteIntegrityError  # type: ignore[import-untyped]
+from strawberry.relay import GlobalID
+from strawberry.types import Info
+
+from phoenix.db import models
+from phoenix.server.api.auth import IsNotReadOnly
+from phoenix.server.api.context import Context
+from phoenix.server.api.exceptions import BadRequest, Conflict, NotFound
+from phoenix.server.api.queries import Query
+from phoenix.server.api.types.GenerativeModel import GenerativeModel, to_gql_generative_model
+from phoenix.server.api.types.node import from_global_id_with_expected_type
+from phoenix.server.api.types.TokenPrice import TokenKind
+
+
+@strawberry.input
+class TokenPriceInput:
+    token_type: str
+    cost_per_million_tokens: float
+    kind: TokenKind
+
+    @property
+    def token_prices(self) -> models.TokenPrice:
+        """Generate TokenPrice instances based on the input."""
+        return models.TokenPrice(
+            token_type=self.token_type,
+            is_prompt=self.kind == TokenKind.PROMPT,
+            base_rate=self.cost_per_million_tokens / 1_000_000,
+        )
+
+
+@strawberry.input
+class CreateModelMutationInput:
+    name: str
+    provider: Optional[str] = None
+    name_pattern: str
+    costs: list[TokenPriceInput]
+    start_time: Optional[datetime] = None
+
+
+@strawberry.type
+class CreateModelMutationPayload:
+    model: GenerativeModel
+    query: Query
+
+
+@strawberry.input
+class UpdateModelMutationInput:
+    id: GlobalID
+    name: str
+    provider: Optional[str]
+    name_pattern: str
+    costs: list[TokenPriceInput]
+    start_time: Optional[datetime] = None
+
+
+@strawberry.type
+class UpdateModelMutationPayload:
+    model: GenerativeModel
+    query: Query
+
+
+@strawberry.input
+class DeleteModelMutationInput:
+    id: GlobalID
+
+
+@strawberry.type
+class DeleteModelMutationPayload:
+    model: GenerativeModel
+    query: Query
+
+
+@strawberry.type
+class ModelMutationMixin:
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    async def create_model(
+        self,
+        info: Info[Context, None],
+        input: CreateModelMutationInput,
+    ) -> CreateModelMutationPayload:
+        cost_types = set(cost.token_type for cost in input.costs)
+        if "input" not in cost_types:
+            raise BadRequest("input cost is required")
+        if "output" not in cost_types:
+            raise BadRequest("output cost is required")
+        name_pattern = _compile_regular_expression(input.name_pattern)
+        token_prices = [cost.token_prices for cost in input.costs]
+        model = models.GenerativeModel(
+            name=input.name,
+            provider=input.provider,
+            name_pattern=name_pattern,
+            is_built_in=False,
+            token_prices=token_prices,
+            start_time=input.start_time,
+        )
+        async with info.context.db() as session:
+            session.add(model)
+            try:
+                await session.flush()
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError):
+                raise Conflict(f"Model with name '{input.name}' already exists")
+
+        return CreateModelMutationPayload(
+            model=to_gql_generative_model(model),
+            query=Query(),
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    async def update_model(
+        self,
+        info: Info[Context, None],
+        input: UpdateModelMutationInput,
+    ) -> UpdateModelMutationPayload:
+        try:
+            model_id = from_global_id_with_expected_type(input.id, GenerativeModel.__name__)
+        except ValueError:
+            raise BadRequest(f'Invalid model id: "{input.id}"')
+
+        cost_types = set(cost.token_type for cost in input.costs)
+        if "input" not in cost_types:
+            raise BadRequest("input cost is required")
+        if "output" not in cost_types:
+            raise BadRequest("output cost is required")
+        name_pattern = _compile_regular_expression(input.name_pattern)
+        token_prices = [cost.token_prices for cost in input.costs]
+        async with info.context.db() as session:
+            model = await session.scalar(
+                sa.select(models.GenerativeModel)
+                .where(models.GenerativeModel.deleted_at.is_(None))
+                .where(models.GenerativeModel.id == model_id)
+                .options(joinedload(models.GenerativeModel.token_prices))
+            )
+            if model is None:
+                raise NotFound(f'Model "{input.id}" not found')
+            if model.is_built_in:
+                raise BadRequest("Cannot update built-in model")
+
+            await session.execute(
+                delete(models.TokenPrice).where(models.TokenPrice.model_id == model.id)
+            )
+
+            await session.refresh(model)
+
+            model.name = input.name
+            model.provider = input.provider or ""
+            model.name_pattern = name_pattern
+            model.token_prices = token_prices
+            model.start_time = input.start_time
+            session.add(model)
+            try:
+                await session.flush()
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError):
+                raise Conflict(f"Model with name '{input.name}' already exists")
+            await session.refresh(model)
+
+        return UpdateModelMutationPayload(
+            model=to_gql_generative_model(model),
+            query=Query(),
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    async def delete_model(
+        self,
+        info: Info[Context, None],
+        input: DeleteModelMutationInput,
+    ) -> DeleteModelMutationPayload:
+        try:
+            model_id = from_global_id_with_expected_type(input.id, GenerativeModel.__name__)
+        except ValueError:
+            raise BadRequest(f'Invalid model id: "{input.id}"')
+
+        async with info.context.db() as session:
+            model = await session.scalar(
+                sa.update(models.GenerativeModel)
+                .values(deleted_at=datetime.now(timezone.utc))
+                .where(models.GenerativeModel.deleted_at.is_(None))
+                .where(models.GenerativeModel.id == model_id)
+                .returning(models.GenerativeModel)
+            )
+            if model is None:
+                raise NotFound(f'Model "{input.id}" not found')
+            if model.is_built_in:
+                await session.rollback()
+                raise BadRequest("Cannot delete built-in model")
+        return DeleteModelMutationPayload(
+            model=to_gql_generative_model(model),
+            query=Query(),
+        )
+
+
+def _compile_regular_expression(maybe_regex: str) -> re.Pattern[str]:
+    """
+    Compile the given string as a regular expression.
+    Raises a BadRequest error if the given string is not a valid regex.
+    """
+    try:
+        return re.compile(maybe_regex)
+    except re.error as error:
+        raise BadRequest(f"Invalid regex: {str(error)}")