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

phoenix/server/api/types/GenerativeModel.py

@@ -1,9 +1,150 @@
+from datetime import datetime
+from enum import Enum
+from typing import Optional
+
 import strawberry
+from openinference.semconv.trace import OpenInferenceLLMProviderValues
+from sqlalchemy import inspect
+from strawberry.relay import Node, NodeID
+from strawberry.types import Info
+from typing_extensions import assert_never
 
+from phoenix.db import models
+from phoenix.server.api.context import Context
+from phoenix.server.api.types.CostBreakdown import CostBreakdown
 from phoenix.server.api.types.GenerativeProvider import GenerativeProviderKey
+from phoenix.server.api.types.ModelInterface import ModelInterface
+from phoenix.server.api.types.SpanCostDetailSummaryEntry import SpanCostDetailSummaryEntry
+from phoenix.server.api.types.SpanCostSummary import SpanCostSummary
+from phoenix.server.api.types.TokenPrice import TokenKind, TokenPrice
+
+
+@strawberry.enum
+class GenerativeModelKind(Enum):
+    CUSTOM = "CUSTOM"
+    BUILT_IN = "BUILT_IN"
 
 
 @strawberry.type
-class GenerativeModel:
+class GenerativeModel(Node, ModelInterface):
+    id_attr: NodeID[int]
     name: str
-    provider_key: GenerativeProviderKey
+    provider: Optional[str]
+    name_pattern: str
+    kind: GenerativeModelKind
+    created_at: datetime
+    updated_at: datetime
+    provider_key: Optional[GenerativeProviderKey]
+    costs: strawberry.Private[Optional[list[models.TokenPrice]]] = None
+    start_time: Optional[datetime] = None
+
+    @strawberry.field
+    async def token_prices(self) -> list[TokenPrice]:
+        if self.costs is None:
+            raise NotImplementedError
+        token_prices: list[TokenPrice] = list()
+        for cost in self.costs:
+            token_prices.append(
+                TokenPrice(
+                    token_type=cost.token_type,
+                    kind=TokenKind.PROMPT if cost.is_prompt else TokenKind.COMPLETION,
+                    cost_per_million_tokens=cost.base_rate * 1_000_000,
+                    cost_per_token=cost.base_rate,
+                )
+            )
+        return token_prices
+
+    @strawberry.field
+    async def cost_summary(self, info: Info[Context, None]) -> SpanCostSummary:
+        loader = info.context.data_loaders.span_cost_summary_by_generative_model
+        summary = await loader.load(self.id_attr)
+        return SpanCostSummary(
+            prompt=CostBreakdown(
+                tokens=summary.prompt.tokens,
+                cost=summary.prompt.cost,
+            ),
+            completion=CostBreakdown(
+                tokens=summary.completion.tokens,
+                cost=summary.completion.cost,
+            ),
+            total=CostBreakdown(
+                tokens=summary.total.tokens,
+                cost=summary.total.cost,
+            ),
+        )
+
+    @strawberry.field
+    async def cost_detail_summary_entries(
+        self,
+        info: Info[Context, None],
+    ) -> list[SpanCostDetailSummaryEntry]:
+        loader = info.context.data_loaders.span_cost_detail_summary_entries_by_generative_model
+        summary = await loader.load(self.id_attr)
+        return [
+            SpanCostDetailSummaryEntry(
+                token_type=entry.token_type,
+                is_prompt=entry.is_prompt,
+                value=CostBreakdown(
+                    tokens=entry.value.tokens,
+                    cost=entry.value.cost,
+                ),
+            )
+            for entry in summary
+        ]
+
+    @strawberry.field
+    async def last_used_at(self, info: Info[Context, None]) -> Optional[datetime]:
+        model_id = self.id_attr
+        return await info.context.data_loaders.last_used_times_by_generative_model_id.load(model_id)
+
+
+def to_gql_generative_model(model: models.GenerativeModel) -> GenerativeModel:
+    costs_are_loaded = isinstance(inspect(model).attrs.token_prices.loaded_value, list)
+    name_pattern = model.name_pattern.pattern
+    assert isinstance(name_pattern, str)
+    return GenerativeModel(
+        id_attr=model.id,
+        name=model.name,
+        provider=model.provider or None,
+        name_pattern=name_pattern,
+        kind=GenerativeModelKind.BUILT_IN if model.is_built_in else GenerativeModelKind.CUSTOM,
+        created_at=model.created_at,
+        updated_at=model.updated_at,
+        start_time=model.start_time,
+        provider_key=_semconv_provider_to_gql_generative_provider_key(model.provider)
+        if model.provider
+        else None,
+        costs=model.token_prices if costs_are_loaded else None,
+    )
+
+
+def _semconv_provider_to_gql_generative_provider_key(
+    semconv_provider_str: str,
+) -> Optional[GenerativeProviderKey]:
+    """
+    Translates a semconv provider string to a GQL GenerativeProviderKey.
+    """
+
+    try:
+        semconv_provider = OpenInferenceLLMProviderValues(semconv_provider_str)
+    except Exception:
+        return None
+    if semconv_provider == OpenInferenceLLMProviderValues.OPENAI:
+        return GenerativeProviderKey.OPENAI
+    if semconv_provider == OpenInferenceLLMProviderValues.ANTHROPIC:
+        return GenerativeProviderKey.ANTHROPIC
+    if semconv_provider == OpenInferenceLLMProviderValues.AZURE:
+        return GenerativeProviderKey.AZURE_OPENAI
+    if semconv_provider == OpenInferenceLLMProviderValues.GOOGLE:
+        return GenerativeProviderKey.GOOGLE
+    if semconv_provider == OpenInferenceLLMProviderValues.DEEPSEEK:
+        return GenerativeProviderKey.DEEPSEEK
+    if semconv_provider == OpenInferenceLLMProviderValues.XAI:
+        return GenerativeProviderKey.XAI
+    if semconv_provider == OpenInferenceLLMProviderValues.AWS:
+        raise NotImplementedError("AWS models are not yet supported")
+    if semconv_provider == OpenInferenceLLMProviderValues.COHERE:
+        raise NotImplementedError("Cohere models are not yet supported")
+    if semconv_provider == OpenInferenceLLMProviderValues.MISTRALAI:
+        raise NotImplementedError("Mistral AI models are not yet supported")
+    assert_never(semconv_provider)

phoenix/server/api/types/{Model.py → InferenceModel.py}

@@ -28,7 +28,7 @@ from .TimeSeries import (
 
 
 @strawberry.type
-class Model:
+class InferenceModel:
     @strawberry.field
     def dimensions(
         self,

phoenix/server/api/types/ModelInterface.py

@@ -0,0 +1,11 @@
+from typing import Optional
+
+import strawberry
+
+from phoenix.server.api.types.GenerativeProvider import GenerativeProviderKey
+
+
+@strawberry.interface
+class ModelInterface:
+    name: str
+    provider_key: Optional[GenerativeProviderKey]

phoenix/server/api/types/PlaygroundModel.py

@@ -0,0 +1,10 @@
+import strawberry
+
+from phoenix.server.api.types.GenerativeProvider import GenerativeProviderKey
+from phoenix.server.api.types.ModelInterface import ModelInterface
+
+
+@strawberry.type
+class PlaygroundModel(ModelInterface):
+    name: str
+    provider_key: GenerativeProviderKey  # PlaygroundModel always has a provider_key

phoenix/server/api/types/Project.py

@@ -28,6 +28,7 @@ from phoenix.server.api.input_types.SpanSort import SpanSort, SpanSortConfig
 from phoenix.server.api.input_types.TimeRange import TimeRange
 from phoenix.server.api.types.AnnotationConfig import AnnotationConfig, to_gql_annotation_config
 from phoenix.server.api.types.AnnotationSummary import AnnotationSummary
+from phoenix.server.api.types.CostBreakdown import CostBreakdown
 from phoenix.server.api.types.DocumentEvaluationSummary import DocumentEvaluationSummary
 from phoenix.server.api.types.pagination import (
     ConnectionArgs,
@@ -40,6 +41,7 @@ from phoenix.server.api.types.pagination import (
 from phoenix.server.api.types.ProjectSession import ProjectSession, to_gql_project_session
 from phoenix.server.api.types.SortDir import SortDir
 from phoenix.server.api.types.Span import Span
+from phoenix.server.api.types.SpanCostSummary import SpanCostSummary
 from phoenix.server.api.types.TimeSeries import TimeSeries, TimeSeriesDataPoint
 from phoenix.server.api.types.Trace import Trace
 from phoenix.server.api.types.ValidationResult import ValidationResult
@@ -175,6 +177,30 @@ class Project(Node):
             ("completion", self.project_rowid, time_range, filter_condition),
         )
 
+    @strawberry.field
+    async def cost_summary(
+        self,
+        info: Info[Context, None],
+        time_range: Optional[TimeRange] = UNSET,
+        filter_condition: Optional[str] = UNSET,
+    ) -> SpanCostSummary:
+        loader = info.context.data_loaders.span_cost_summary_by_project
+        summary = await loader.load((self.project_rowid, time_range, filter_condition))
+        return SpanCostSummary(
+            prompt=CostBreakdown(
+                tokens=summary.prompt.tokens,
+                cost=summary.prompt.cost,
+            ),
+            completion=CostBreakdown(
+                tokens=summary.completion.tokens,
+                cost=summary.completion.cost,
+            ),
+            total=CostBreakdown(
+                tokens=summary.total.tokens,
+                cost=summary.total.cost,
+            ),
+        )
+
     @strawberry.field
     async def latency_ms_quantile(
         self,
@@ -238,6 +264,7 @@ class Project(Node):
     ) -> Connection[Span]:
         stmt = (
             select(models.Span.id)
+            .select_from(models.Span)
             .join(models.Trace)
             .where(models.Trace.project_rowid == self.project_rowid)
         )
@@ -410,6 +437,21 @@ class Project(Node):
                     assert_never(sort.col)
                 key = sort_subq.c.key
                 stmt = stmt.join(sort_subq, table.id == sort_subq.c.id)
+            elif sort.col is ProjectSessionColumn.costTotal:
+                sort_subq = (
+                    select(
+                        models.Trace.project_session_rowid.label("id"),
+                        func.sum(models.SpanCost.total_cost).label("key"),
+                    )
+                    .join_from(
+                        models.Trace,
+                        models.SpanCost,
+                        models.Trace.id == models.SpanCost.trace_rowid,
+                    )
+                    .group_by(models.Trace.project_session_rowid)
+                ).subquery()
+                key = sort_subq.c.key
+                stmt = stmt.join(sort_subq, table.id == sort_subq.c.id)
             else:
                 assert_never(sort.col)
             stmt = stmt.add_columns(key)

phoenix/server/api/types/ProjectSession.py

@@ -9,8 +9,11 @@ from strawberry.relay import Connection, GlobalID, Node, NodeID
 
 from phoenix.db import models
 from phoenix.server.api.context import Context
+from phoenix.server.api.types.CostBreakdown import CostBreakdown
 from phoenix.server.api.types.MimeType import MimeType
 from phoenix.server.api.types.pagination import ConnectionArgs, CursorString, connection_from_list
+from phoenix.server.api.types.SpanCostDetailSummaryEntry import SpanCostDetailSummaryEntry
+from phoenix.server.api.types.SpanCostSummary import SpanCostSummary
 from phoenix.server.api.types.SpanIOValue import SpanIOValue
 from phoenix.server.api.types.TokenUsage import TokenUsage
 
@@ -122,6 +125,47 @@ class ProjectSession(Node):
             (self.id_attr, probability)
         )
 
+    @strawberry.field
+    async def cost_summary(
+        self,
+        info: Info[Context, None],
+    ) -> SpanCostSummary:
+        loader = info.context.data_loaders.span_cost_summary_by_project_session
+        summary = await loader.load(self.id_attr)
+        return SpanCostSummary(
+            prompt=CostBreakdown(
+                tokens=summary.prompt.tokens,
+                cost=summary.prompt.cost,
+            ),
+            completion=CostBreakdown(
+                tokens=summary.completion.tokens,
+                cost=summary.completion.cost,
+            ),
+            total=CostBreakdown(
+                tokens=summary.total.tokens,
+                cost=summary.total.cost,
+            ),
+        )
+
+    @strawberry.field
+    async def cost_detail_summary_entries(
+        self,
+        info: Info[Context, None],
+    ) -> list[SpanCostDetailSummaryEntry]:
+        loader = info.context.data_loaders.span_cost_detail_summary_entries_by_project_session
+        summary = await loader.load(self.id_attr)
+        return [
+            SpanCostDetailSummaryEntry(
+                token_type=entry.token_type,
+                is_prompt=entry.is_prompt,
+                value=CostBreakdown(
+                    tokens=entry.value.tokens,
+                    cost=entry.value.cost,
+                ),
+            )
+            for entry in summary
+        ]
+
 
 def to_gql_project_session(project_session: models.ProjectSession) -> ProjectSession:
     return ProjectSession(

phoenix/server/api/types/Span.py

@@ -19,6 +19,7 @@ from typing_extensions import Annotated, TypeAlias
 import phoenix.trace.schemas as trace_schema
 from phoenix.db import models
 from phoenix.server.api.context import Context
+from phoenix.server.api.dataloaders import types as dataloader_types
 from phoenix.server.api.helpers.dataset_helpers import (
     get_dataset_example_input,
     get_dataset_example_output,
@@ -33,6 +34,7 @@ from phoenix.server.api.input_types.SpanAnnotationSort import (
     SpanAnnotationSort,
 )
 from phoenix.server.api.types.AnnotationSummary import AnnotationSummary
+from phoenix.server.api.types.CostBreakdown import CostBreakdown
 from phoenix.server.api.types.DocumentRetrievalMetrics import DocumentRetrievalMetrics
 from phoenix.server.api.types.Evaluation import DocumentEvaluation
 from phoenix.server.api.types.ExampleRevisionInterface import ExampleRevision
@@ -41,6 +43,8 @@ from phoenix.server.api.types.MimeType import MimeType
 from phoenix.server.api.types.pagination import ConnectionArgs, CursorString, connection_from_list
 from phoenix.server.api.types.SortDir import SortDir
 from phoenix.server.api.types.SpanAnnotation import SpanAnnotation, to_gql_span_annotation
+from phoenix.server.api.types.SpanCostDetailSummaryEntry import SpanCostDetailSummaryEntry
+from phoenix.server.api.types.SpanCostSummary import SpanCostSummary
 from phoenix.server.api.types.SpanIOValue import SpanIOValue, truncate_value
 from phoenix.trace.attributes import get_attribute_value
 
@@ -790,6 +794,139 @@ class Span(Node):
             )
         ]
 
+    @strawberry.field
+    async def cost_summary(self, info: Info[Context, None]) -> Optional[SpanCostSummary]:
+        span_cost = await info.context.data_loaders.span_cost_by_span.load(self.span_rowid)
+        if span_cost is None:
+            return None
+        return SpanCostSummary(
+            prompt=CostBreakdown(
+                tokens=span_cost.prompt_tokens,
+                cost=span_cost.prompt_cost,
+            ),
+            completion=CostBreakdown(
+                tokens=span_cost.completion_tokens,
+                cost=span_cost.completion_cost,
+            ),
+            total=CostBreakdown(
+                tokens=span_cost.total_tokens,
+                cost=span_cost.total_cost,
+            ),
+        )
+
+    @strawberry.field
+    async def cost_detail_summary_entries(
+        self, info: Info[Context, None]
+    ) -> list[SpanCostDetailSummaryEntry]:
+        loader = info.context.data_loaders.span_cost_detail_summary_entries_by_span
+        entries = await loader.load(self.span_rowid)
+        return [
+            SpanCostDetailSummaryEntry(
+                token_type=entry.token_type,
+                is_prompt=entry.is_prompt,
+                value=CostBreakdown(tokens=entry.value.tokens, cost=entry.value.cost),
+            )
+            for entry in entries
+        ]
+
+    @strawberry.field
+    async def cumulative_cost_summary(self, info: Info[Context, None]) -> Optional[SpanCostSummary]:
+        max_depth = 0
+        descendant_rowids = await info.context.data_loaders.span_descendants.load(
+            (self.span_rowid, max_depth)
+        )
+        span_costs = await info.context.data_loaders.span_cost_by_span.load_many(
+            (self.span_rowid, *descendant_rowids)
+        )
+        total_cost: Optional[float] = None
+        total_tokens: Optional[float] = None
+        prompt_cost: Optional[float] = None
+        prompt_tokens: Optional[float] = None
+        completion_cost: Optional[float] = None
+        completion_tokens: Optional[float] = None
+        for span_cost in span_costs:
+            if span_cost is None:
+                continue
+            if span_cost.total_cost is not None:
+                total_cost = (total_cost or 0) + span_cost.total_cost
+            if span_cost.total_tokens is not None:
+                total_tokens = (total_tokens or 0) + span_cost.total_tokens
+            if span_cost.prompt_cost is not None:
+                prompt_cost = (prompt_cost or 0) + span_cost.prompt_cost
+            if span_cost.prompt_tokens is not None:
+                prompt_tokens = (prompt_tokens or 0) + span_cost.prompt_tokens
+            if span_cost.completion_cost is not None:
+                completion_cost = (completion_cost or 0) + span_cost.completion_cost
+            if span_cost.completion_tokens is not None:
+                completion_tokens = (completion_tokens or 0) + span_cost.completion_tokens
+        return SpanCostSummary(
+            prompt=CostBreakdown(
+                tokens=prompt_tokens,
+                cost=prompt_cost,
+            ),
+            completion=CostBreakdown(
+                tokens=completion_tokens,
+                cost=completion_cost,
+            ),
+            total=CostBreakdown(
+                tokens=total_tokens,
+                cost=total_cost,
+            ),
+        )
+
+    @strawberry.field
+    async def cumulative_cost_detail_summary_entries(
+        self, info: Info[Context, None]
+    ) -> list[SpanCostDetailSummaryEntry]:
+        max_depth = 0
+        descendant_rowids = await info.context.data_loaders.span_descendants.load(
+            (self.span_rowid, max_depth)
+        )
+        entry_lists = (
+            await info.context.data_loaders.span_cost_detail_summary_entries_by_span.load_many(
+                (self.span_rowid, *descendant_rowids)
+            )
+        )
+
+        TokenType: TypeAlias = str
+        IsPrompt: TypeAlias = bool
+        grouped_entries: dict[
+            IsPrompt, dict[TokenType, list[dataloader_types.SpanCostDetailSummaryEntry]]
+        ] = {}
+
+        for entries in entry_lists:
+            for entry in entries:
+                is_prompt = entry.is_prompt
+                token_type = entry.token_type
+                if is_prompt not in grouped_entries:
+                    grouped_entries[is_prompt] = {}
+                if token_type not in grouped_entries[is_prompt]:
+                    grouped_entries[is_prompt][token_type] = []
+                grouped_entries[is_prompt][token_type].append(entry)
+
+        result: list[SpanCostDetailSummaryEntry] = []
+        for is_prompt in (True, False):
+            entries_by_token_type = grouped_entries[is_prompt]
+            for token_type, entries in sorted(entries_by_token_type.items()):
+                cost: Optional[float] = None
+                tokens: Optional[float] = None
+                for entry in entries:
+                    if entry.value.cost is not None:
+                        cost = (cost or 0) + entry.value.cost
+                    if entry.value.tokens is not None:
+                        tokens = (tokens or 0) + entry.value.tokens
+                result.append(
+                    SpanCostDetailSummaryEntry(
+                        token_type=token_type,
+                        is_prompt=is_prompt,
+                        value=CostBreakdown(
+                            tokens=tokens,
+                            cost=cost,
+                        ),
+                    )
+                )
+        return result
+
 
 def _hide_embedding_vectors(attributes: Mapping[str, Any]) -> Mapping[str, Any]:
     if not (

phoenix/server/api/types/SpanCostDetailSummaryEntry.py

@@ -0,0 +1,10 @@
+import strawberry
+
+from phoenix.server.api.types.CostBreakdown import CostBreakdown
+
+
+@strawberry.type
+class SpanCostDetailSummaryEntry:
+    token_type: str
+    is_prompt: bool
+    value: CostBreakdown = strawberry.field(default_factory=CostBreakdown)

phoenix/server/api/types/SpanCostSummary.py

@@ -0,0 +1,10 @@
+import strawberry
+
+from phoenix.server.api.types.CostBreakdown import CostBreakdown
+
+
+@strawberry.type
+class SpanCostSummary:
+    prompt: CostBreakdown = strawberry.field(default_factory=CostBreakdown)
+    completion: CostBreakdown = strawberry.field(default_factory=CostBreakdown)
+    total: CostBreakdown = strawberry.field(default_factory=CostBreakdown)

phoenix/server/api/types/TokenPrice.py

@@ -0,0 +1,16 @@
+from enum import Enum
+
+import strawberry
+
+
+class TokenKind(Enum):
+    PROMPT = "prompt"
+    COMPLETION = "completion"
+
+
+@strawberry.type
+class TokenPrice:
+    token_type: str
+    kind: TokenKind
+    cost_per_million_tokens: float
+    cost_per_token: float

phoenix/server/api/types/TokenUsage.py

@@ -3,9 +3,9 @@ import strawberry
 
 @strawberry.type
 class TokenUsage:
-    prompt: int = 0
-    completion: int = 0
+    prompt: float = 0
+    completion: float = 0
 
     @strawberry.field
-    async def total(self) -> int:
+    async def total(self) -> float:
         return self.prompt + self.completion

phoenix/server/api/types/Trace.py

@@ -14,6 +14,7 @@ from typing_extensions import TypeAlias
 from phoenix.db import models
 from phoenix.server.api.context import Context
 from phoenix.server.api.input_types.TraceAnnotationSort import TraceAnnotationSort
+from phoenix.server.api.types.CostBreakdown import CostBreakdown
 from phoenix.server.api.types.pagination import (
     ConnectionArgs,
     CursorString,
@@ -21,6 +22,8 @@ from phoenix.server.api.types.pagination import (
 )
 from phoenix.server.api.types.SortDir import SortDir
 from phoenix.server.api.types.Span import Span
+from phoenix.server.api.types.SpanCostDetailSummaryEntry import SpanCostDetailSummaryEntry
+from phoenix.server.api.types.SpanCostSummary import SpanCostSummary
 from phoenix.server.api.types.TraceAnnotation import TraceAnnotation, to_gql_trace_annotation
 
 if TYPE_CHECKING:
@@ -226,6 +229,44 @@ class Trace(Node):
             annotations = await session.scalars(stmt)
         return [to_gql_trace_annotation(annotation) for annotation in annotations]
 
+    @strawberry.field
+    async def cost_summary(
+        self,
+        info: Info[Context, None],
+    ) -> SpanCostSummary:
+        loader = info.context.data_loaders.span_cost_summary_by_trace
+        summary = await loader.load(self.trace_rowid)
+        return SpanCostSummary(
+            prompt=CostBreakdown(
+                tokens=summary.prompt.tokens,
+                cost=summary.prompt.cost,
+            ),
+            completion=CostBreakdown(
+                tokens=summary.completion.tokens,
+                cost=summary.completion.cost,
+            ),
+            total=CostBreakdown(
+                tokens=summary.total.tokens,
+                cost=summary.total.cost,
+            ),
+        )
+
+    @strawberry.field
+    async def cost_detail_summary_entries(
+        self,
+        info: Info[Context, None],
+    ) -> list[SpanCostDetailSummaryEntry]:
+        loader = info.context.data_loaders.span_cost_detail_summary_entries_by_trace
+        entries = await loader.load(self.trace_rowid)
+        return [
+            SpanCostDetailSummaryEntry(
+                token_type=entry.token_type,
+                is_prompt=entry.is_prompt,
+                value=CostBreakdown(tokens=entry.value.tokens, cost=entry.value.cost),
+            )
+            for entry in entries
+        ]
+
 
 INPUT_VALUE = SpanAttributes.INPUT_VALUE.split(".")
 OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE.split(".")