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

phoenix/server/api/queries.py

@@ -1,3 +1,4 @@
+import re
 from collections import defaultdict
 from datetime import datetime
 from typing import Iterable, Iterator, Optional, Union, cast
@@ -21,12 +22,6 @@ from phoenix.config import (
 from phoenix.db import models
 from phoenix.db.constants import DEFAULT_PROJECT_TRACE_RETENTION_POLICY_ID
 from phoenix.db.helpers import SupportedSQLDialect, exclude_experiment_projects
-from phoenix.db.models import DatasetExample as OrmExample
-from phoenix.db.models import DatasetExampleRevision as OrmRevision
-from phoenix.db.models import DatasetVersion as OrmVersion
-from phoenix.db.models import Experiment as OrmExperiment
-from phoenix.db.models import ExperimentRun as OrmExperimentRun
-from phoenix.db.models import Trace as OrmTrace
 from phoenix.pointcloud.clustering import Hdbscan
 from phoenix.server.api.auth import MSG_ADMIN_ONLY, IsAdmin
 from phoenix.server.api.context import Context
@@ -62,12 +57,13 @@ from phoenix.server.api.types.Experiment import Experiment
 from phoenix.server.api.types.ExperimentComparison import ExperimentComparison, RunComparisonItem
 from phoenix.server.api.types.ExperimentRun import ExperimentRun, to_gql_experiment_run
 from phoenix.server.api.types.Functionality import Functionality
-from phoenix.server.api.types.GenerativeModel import GenerativeModel
+from phoenix.server.api.types.GenerativeModel import GenerativeModel, to_gql_generative_model
 from phoenix.server.api.types.GenerativeProvider import GenerativeProvider, GenerativeProviderKey
+from phoenix.server.api.types.InferenceModel import InferenceModel
 from phoenix.server.api.types.InferencesRole import AncillaryInferencesRole, InferencesRole
-from phoenix.server.api.types.Model import Model
 from phoenix.server.api.types.node import from_global_id, from_global_id_with_expected_type
 from phoenix.server.api.types.pagination import ConnectionArgs, CursorString, connection_from_list
+from phoenix.server.api.types.PlaygroundModel import PlaygroundModel
 from phoenix.server.api.types.Project import Project
 from phoenix.server.api.types.ProjectSession import ProjectSession, to_gql_project_session
 from phoenix.server.api.types.ProjectTraceRetentionPolicy import ProjectTraceRetentionPolicy
@@ -115,20 +111,39 @@ class Query:
         ]
 
     @strawberry.field
-    async def models(self, input: Optional[ModelsInput] = None) -> list[GenerativeModel]:
+    async def generative_models(
+        self,
+        info: Info[Context, None],
+    ) -> list[GenerativeModel]:
+        async with info.context.db() as session:
+            result = await session.scalars(
+                select(models.GenerativeModel)
+                .where(models.GenerativeModel.deleted_at.is_(None))
+                .order_by(
+                    models.GenerativeModel.is_built_in.asc(),  # display custom models first
+                    models.GenerativeModel.provider.nullslast(),
+                    models.GenerativeModel.name,
+                )
+                .options(joinedload(models.GenerativeModel.token_prices))
+            )
+
+        return [to_gql_generative_model(model) for model in result.unique()]
+
+    @strawberry.field
+    async def playground_models(self, input: Optional[ModelsInput] = None) -> list[PlaygroundModel]:
         if input is not None and input.provider_key is not None:
             supported_model_names = PLAYGROUND_CLIENT_REGISTRY.list_models(input.provider_key)
             supported_models = [
-                GenerativeModel(name=model_name, provider_key=input.provider_key)
+                PlaygroundModel(name=model_name, provider_key=input.provider_key)
                 for model_name in supported_model_names
             ]
             return supported_models
 
         registered_models = PLAYGROUND_CLIENT_REGISTRY.list_all_models()
-        all_models: list[GenerativeModel] = []
+        all_models: list[PlaygroundModel] = []
         for provider_key, model_name in registered_models:
             if model_name is not None and provider_key is not None:
-                all_models.append(GenerativeModel(name=model_name, provider_key=provider_key))
+                all_models.append(PlaygroundModel(name=model_name, provider_key=provider_key))
         return all_models
 
     @strawberry.field
@@ -330,7 +345,7 @@ class Query:
             )
 
         experiment_ids_ = [
-            from_global_id_with_expected_type(experiment_id, OrmExperiment.__name__)
+            from_global_id_with_expected_type(experiment_id, models.Experiment.__name__)
             for experiment_id in experiment_ids
         ]
         if len(set(experiment_ids_)) != len(experiment_ids_):
@@ -340,18 +355,18 @@ class Query:
             validation_result = (
                 await session.execute(
                     select(
-                        func.count(distinct(OrmVersion.dataset_id)),
-                        func.max(OrmVersion.dataset_id),
-                        func.max(OrmVersion.id),
-                        func.count(OrmExperiment.id),
+                        func.count(distinct(models.DatasetVersion.dataset_id)),
+                        func.max(models.DatasetVersion.dataset_id),
+                        func.max(models.DatasetVersion.id),
+                        func.count(models.Experiment.id),
                     )
-                    .select_from(OrmVersion)
+                    .select_from(models.DatasetVersion)
                     .join(
-                        OrmExperiment,
-                        OrmExperiment.dataset_version_id == OrmVersion.id,
+                        models.Experiment,
+                        models.Experiment.dataset_version_id == models.DatasetVersion.id,
                     )
                     .where(
-                        OrmExperiment.id.in_(experiment_ids_),
+                        models.Experiment.id.in_(experiment_ids_),
                     )
                 )
             ).first()
@@ -365,29 +380,33 @@ class Query:
                 raise ValueError("Unable to resolve one or more experiment IDs.")
 
             revision_ids = (
-                select(func.max(OrmRevision.id))
-                .join(OrmExample, OrmExample.id == OrmRevision.dataset_example_id)
+                select(func.max(models.DatasetExampleRevision.id))
+                .join(
+                    models.DatasetExample,
+                    models.DatasetExample.id == models.DatasetExampleRevision.dataset_example_id,
+                )
                 .where(
                     and_(
-                        OrmRevision.dataset_version_id <= version_id,
-                        OrmExample.dataset_id == dataset_id,
+                        models.DatasetExampleRevision.dataset_version_id <= version_id,
+                        models.DatasetExample.dataset_id == dataset_id,
                     )
                 )
-                .group_by(OrmRevision.dataset_example_id)
+                .group_by(models.DatasetExampleRevision.dataset_example_id)
                 .scalar_subquery()
             )
             examples_query = (
-                select(OrmExample)
-                .distinct(OrmExample.id)
+                select(models.DatasetExample)
+                .distinct(models.DatasetExample.id)
                 .join(
-                    OrmRevision,
+                    models.DatasetExampleRevision,
                     onclause=and_(
-                        OrmExample.id == OrmRevision.dataset_example_id,
-                        OrmRevision.id.in_(revision_ids),
-                        OrmRevision.revision_kind != "DELETE",
+                        models.DatasetExample.id
+                        == models.DatasetExampleRevision.dataset_example_id,
+                        models.DatasetExampleRevision.id.in_(revision_ids),
+                        models.DatasetExampleRevision.revision_kind != "DELETE",
                     ),
                 )
-                .order_by(OrmExample.id.desc())
+                .order_by(models.DatasetExample.id.desc())
             )
 
             if filter_condition:
@@ -401,18 +420,20 @@ class Query:
 
             ExampleID: TypeAlias = int
             ExperimentID: TypeAlias = int
-            runs: defaultdict[ExampleID, defaultdict[ExperimentID, list[OrmExperimentRun]]] = (
+            runs: defaultdict[ExampleID, defaultdict[ExperimentID, list[models.ExperimentRun]]] = (
                 defaultdict(lambda: defaultdict(list))
             )
             async for run in await session.stream_scalars(
-                select(OrmExperimentRun)
+                select(models.ExperimentRun)
                 .where(
                     and_(
-                        OrmExperimentRun.dataset_example_id.in_(example.id for example in examples),
-                        OrmExperimentRun.experiment_id.in_(experiment_ids_),
+                        models.ExperimentRun.dataset_example_id.in_(
+                            example.id for example in examples
+                        ),
+                        models.ExperimentRun.experiment_id.in_(experiment_ids_),
                     )
                 )
-                .options(joinedload(OrmExperimentRun.trace).load_only(OrmTrace.trace_id))
+                .options(joinedload(models.ExperimentRun.trace).load_only(models.Trace.trace_id))
             ):
                 runs[run.dataset_example_id][run.experiment_id].append(run)
 
@@ -460,7 +481,7 @@ class Query:
             compile_sqlalchemy_filter_condition(
                 filter_condition=condition,
                 experiment_ids=[
-                    from_global_id_with_expected_type(experiment_id, OrmExperiment.__name__)
+                    from_global_id_with_expected_type(experiment_id, models.Experiment.__name__)
                     for experiment_id in experiment_ids
                 ],
             )
@@ -482,8 +503,8 @@ class Query:
         )
 
     @strawberry.field
-    def model(self) -> Model:
-        return Model()
+    def model(self) -> InferenceModel:
+        return InferenceModel()
 
     @strawberry.field
     async def node(self, id: GlobalID, info: Info[Context, None]) -> Node:
@@ -658,6 +679,18 @@ class Query:
                 if not trace_annotation:
                     raise NotFound(f"Unknown trace annotation: {id}")
             return to_gql_trace_annotation(trace_annotation)
+        elif type_name == GenerativeModel.__name__:
+            async with info.context.db() as session:
+                stmt = (
+                    select(models.GenerativeModel)
+                    .where(models.GenerativeModel.deleted_at.is_(None))
+                    .where(models.GenerativeModel.id == node_id)
+                    .options(joinedload(models.GenerativeModel.token_prices))
+                )
+                model = await session.scalar(stmt)
+                if not model:
+                    raise NotFound(f"Unknown model: {id}")
+            return to_gql_generative_model(model)
         raise NotFound(f"Unknown node type: {type_name}")
 
     @strawberry.field
@@ -964,6 +997,14 @@ class Query:
             for table_name, num_bytes in stats
         ]
 
+    @strawberry.field
+    def validate_regular_expression(self, regex: str) -> ValidationResult:
+        try:
+            re.compile(regex)
+            return ValidationResult(is_valid=True, error_message=None)
+        except re.error as error:
+            return ValidationResult(is_valid=False, error_message=str(error))
+
 
 def _consolidate_sqlite_db_table_stats(
     stats: Iterable[tuple[str, int]],

phoenix/server/api/routers/v1/traces.py

@@ -3,7 +3,6 @@ import zlib
 from typing import Any, Literal, Optional
 
 from fastapi import APIRouter, BackgroundTasks, Header, HTTPException, Query
-from google.protobuf.json_format import MessageToJson
 from google.protobuf.message import DecodeError
 from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
     ExportTraceServiceRequest,
@@ -14,7 +13,7 @@ from sqlalchemy import insert, select
 from starlette.concurrency import run_in_threadpool
 from starlette.datastructures import State
 from starlette.requests import Request
-from starlette.responses import JSONResponse
+from starlette.responses import Response
 from starlette.status import (
     HTTP_404_NOT_FOUND,
     HTTP_415_UNSUPPORTED_MEDIA_TYPE,
@@ -66,7 +65,7 @@ async def post_traces(
     background_tasks: BackgroundTasks,
     content_type: Optional[str] = Header(default=None),
     content_encoding: Optional[str] = Header(default=None),
-) -> JSONResponse:
+) -> Response:
     if content_type != "application/x-protobuf":
         raise HTTPException(
             detail=f"Unsupported content type: {content_type}",
@@ -91,7 +90,15 @@ async def post_traces(
             status_code=HTTP_422_UNPROCESSABLE_ENTITY,
         )
     background_tasks.add_task(_add_spans, req, request.state)
-    return JSONResponse(MessageToJson(ExportTraceServiceResponse()))
+
+    # "The server MUST use the same Content-Type in the response as it received in the request"
+    response_message = ExportTraceServiceResponse()
+    response_bytes = response_message.SerializeToString()
+    return Response(
+        content=response_bytes,
+        media_type="application/x-protobuf",
+        status_code=200,
+    )
 
 
 class TraceAnnotationResult(V1RoutesBaseModel):

phoenix/server/api/subscriptions.py

@@ -62,6 +62,7 @@ from phoenix.server.api.types.Experiment import to_gql_experiment
 from phoenix.server.api.types.ExperimentRun import to_gql_experiment_run
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.Span import Span
+from phoenix.server.daemons.span_cost_calculator import SpanCostCalculator
 from phoenix.server.dml_event import SpanInsertEvent
 from phoenix.server.types import DbSessionFactory
 from phoenix.utilities.template_formatters import (
@@ -173,6 +174,19 @@ class Subscription:
             db_span = get_db_span(span, db_trace)
             session.add(db_span)
             await session.flush()
+            try:
+                span_cost = info.context.span_cost_calculator.calculate_cost(
+                    start_time=db_span.start_time,
+                    attributes=span.attributes,
+                )
+            except Exception as e:
+                logger.exception(f"Failed to calculate cost for span {db_span.id}: {e}")
+                span_cost = None
+            if span_cost:
+                span_cost.span_rowid = db_span.id
+                span_cost.trace_rowid = db_span.trace_rowid
+                session.add(span_cost)
+
         info.context.event_queue.put(SpanInsertEvent(ids=(playground_project_id,)))
         yield ChatCompletionSubscriptionResult(span=Span(span_rowid=db_span.id, db_span=db_span))
 
@@ -372,14 +386,18 @@ class Subscription:
                     and not write_already_in_progress
                 ):
                     result_payloads_stream = _chat_completion_result_payloads(
-                        db=info.context.db, results=_drain_no_wait(results)
+                        db=info.context.db,
+                        results=_drain_no_wait(results),
+                        span_cost_calculator=info.context.span_cost_calculator,
                     )
                     task = _create_task_with_timeout(result_payloads_stream)
                     in_progress.append((None, result_payloads_stream, task))
                     last_write_time = datetime.now()
         if remaining_results := await _drain(results):
             async for result_payload in _chat_completion_result_payloads(
-                db=info.context.db, results=remaining_results
+                db=info.context.db,
+                results=remaining_results,
+                span_cost_calculator=info.context.span_cost_calculator,
             ):
                 yield result_payload
 
@@ -463,6 +481,7 @@ async def _chat_completion_result_payloads(
     *,
     db: DbSessionFactory,
     results: Sequence[ChatCompletionResult],
+    span_cost_calculator: SpanCostCalculator,
 ) -> ChatStream:
     if not results:
         return
@@ -470,6 +489,19 @@ async def _chat_completion_result_payloads(
         for _, span, run in results:
             if span:
                 session.add(span)
+                await session.flush()
+                try:
+                    span_cost = 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 = span.trace_rowid
+                    session.add(span_cost)
             session.add(run)
         await session.flush()
     for example_id, span, run in results:
@@ -594,3 +626,5 @@ LLM_OUTPUT_MESSAGES = SpanAttributes.LLM_OUTPUT_MESSAGES
 LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION
 LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT
 PROMPT_TEMPLATE_VARIABLES = SpanAttributes.LLM_PROMPT_TEMPLATE_VARIABLES
+LLM_MODEL_NAME = SpanAttributes.LLM_MODEL_NAME
+LLM_PROVIDER = SpanAttributes.LLM_PROVIDER

phoenix/server/api/types/CostBreakdown.py

@@ -0,0 +1,15 @@
+from typing import Optional
+
+import strawberry
+
+
+@strawberry.type
+class CostBreakdown:
+    tokens: Optional[float] = None
+    cost: Optional[float] = None
+
+    @strawberry.field
+    def cost_per_token(self) -> Optional[float]:
+        if self.tokens and self.cost:
+            return self.cost / self.tokens
+        return None

phoenix/server/api/types/Experiment.py

@@ -2,8 +2,9 @@ from datetime import datetime
 from typing import ClassVar, Optional
 
 import strawberry
-from sqlalchemy import select
+from sqlalchemy import func, select
 from sqlalchemy.orm import joinedload
+from sqlalchemy.sql.functions import coalesce
 from strawberry import UNSET, Private
 from strawberry.relay import Connection, Node, NodeID
 from strawberry.scalars import JSON
@@ -11,6 +12,7 @@ from strawberry.types import Info
 
 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.ExperimentAnnotationSummary import ExperimentAnnotationSummary
 from phoenix.server.api.types.ExperimentRun import ExperimentRun, to_gql_experiment_run
 from phoenix.server.api.types.pagination import (
@@ -19,6 +21,8 @@ from phoenix.server.api.types.pagination import (
     connection_from_list,
 )
 from phoenix.server.api.types.Project import Project
+from phoenix.server.api.types.SpanCostDetailSummaryEntry import SpanCostDetailSummaryEntry
+from phoenix.server.api.types.SpanCostSummary import SpanCostSummary
 
 
 @strawberry.type
@@ -130,6 +134,60 @@ class Experiment(Node):
     def last_updated_at(self, info: Info[Context, None]) -> Optional[datetime]:
         return info.context.last_updated_at.get(self._table, self.id_attr)
 
+    @strawberry.field
+    async def cost_summary(self, info: Info[Context, None]) -> SpanCostSummary:
+        experiment_id = self.id_attr
+        summary = await info.context.data_loaders.span_cost_summary_by_experiment.load(
+            experiment_id
+        )
+        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]:
+        experiment_id = self.id_attr
+
+        stmt = (
+            select(
+                models.SpanCostDetail.token_type,
+                models.SpanCostDetail.is_prompt,
+                coalesce(func.sum(models.SpanCostDetail.cost), 0).label("cost"),
+                coalesce(func.sum(models.SpanCostDetail.tokens), 0).label("tokens"),
+            )
+            .select_from(models.SpanCostDetail)
+            .join(models.SpanCost, models.SpanCostDetail.span_cost_id == models.SpanCost.id)
+            .join(models.Span, models.SpanCost.span_rowid == models.Span.id)
+            .join(models.Trace, models.Span.trace_rowid == models.Trace.id)
+            .join(models.ExperimentRun, models.ExperimentRun.trace_id == models.Trace.trace_id)
+            .where(models.ExperimentRun.experiment_id == experiment_id)
+            .group_by(models.SpanCostDetail.token_type, models.SpanCostDetail.is_prompt)
+        )
+
+        async with info.context.db() as session:
+            data = await session.stream(stmt)
+            return [
+                SpanCostDetailSummaryEntry(
+                    token_type=token_type,
+                    is_prompt=is_prompt,
+                    value=CostBreakdown(tokens=tokens, cost=cost),
+                )
+                async for token_type, is_prompt, cost, tokens in data
+            ]
+
 
 def to_gql_experiment(
     experiment: models.Experiment,

phoenix/server/api/types/ExperimentRun.py

@@ -2,8 +2,9 @@ from datetime import datetime
 from typing import TYPE_CHECKING, Annotated, Optional
 
 import strawberry
-from sqlalchemy import select
+from sqlalchemy import func, select
 from sqlalchemy.orm import load_only
+from sqlalchemy.sql.functions import coalesce
 from strawberry import UNSET
 from strawberry.relay import Connection, GlobalID, Node, NodeID
 from strawberry.scalars import JSON
@@ -11,6 +12,7 @@ from strawberry.types import Info
 
 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.ExperimentRunAnnotation import (
     ExperimentRunAnnotation,
     to_gql_experiment_run_annotation,
@@ -20,6 +22,8 @@ from phoenix.server.api.types.pagination import (
     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.Trace import Trace
 
 if TYPE_CHECKING:
@@ -98,6 +102,58 @@ class ExperimentRun(Node):
             version_id=version_id,
         )
 
+    @strawberry.field
+    async def cost_summary(self, info: Info[Context, None]) -> SpanCostSummary:
+        run_id = self.id_attr
+        summary = await info.context.data_loaders.span_cost_summary_by_experiment_run.load(run_id)
+        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]:
+        run_id = self.id_attr
+
+        stmt = (
+            select(
+                models.SpanCostDetail.token_type,
+                models.SpanCostDetail.is_prompt,
+                coalesce(func.sum(models.SpanCostDetail.cost), 0).label("cost"),
+                coalesce(func.sum(models.SpanCostDetail.tokens), 0).label("tokens"),
+            )
+            .select_from(models.SpanCostDetail)
+            .join(models.SpanCost, models.SpanCostDetail.span_cost_id == models.SpanCost.id)
+            .join(models.Span, models.SpanCost.span_rowid == models.Span.id)
+            .join(models.Trace, models.Span.trace_rowid == models.Trace.id)
+            .join(models.ExperimentRun, models.ExperimentRun.trace_id == models.Trace.trace_id)
+            .where(models.ExperimentRun.id == run_id)
+            .group_by(models.SpanCostDetail.token_type, models.SpanCostDetail.is_prompt)
+        )
+
+        async with info.context.db() as session:
+            data = await session.stream(stmt)
+            return [
+                SpanCostDetailSummaryEntry(
+                    token_type=token_type,
+                    is_prompt=is_prompt,
+                    value=CostBreakdown(tokens=tokens, cost=cost),
+                )
+                async for token_type, is_prompt, cost, tokens in data
+            ]
+
 
 def to_gql_experiment_run(run: models.ExperimentRun) -> ExperimentRun:
     """
@@ -109,9 +165,7 @@ def to_gql_experiment_run(run: models.ExperimentRun) -> ExperimentRun:
     return ExperimentRun(
         id_attr=run.id,
         experiment_id=GlobalID(Experiment.__name__, str(run.experiment_id)),
-        trace_id=trace_id
-        if (trace := run.trace) and (trace_id := trace.trace_id) is not None
-        else None,
+        trace_id=run.trace.trace_id if run.trace else None,
         output=run.output.get("task_output"),
         start_time=run.start_time,
         end_time=run.end_time,