arize-phoenix 11.32.1__py3-none-any.whl → 11.34.0__py3-none-any.whl

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

@@ -102,7 +102,7 @@ async def post_evaluations(
             detail="Evaluation name must not be blank/empty",
             status_code=HTTP_422_UNPROCESSABLE_ENTITY,
         )
-    await request.state.queue_evaluation_for_bulk_insert(evaluation)
+    await request.state.enqueue_evaluation(evaluation)
     return Response()
 
 
@@ -221,7 +221,7 @@ async def _add_evaluations(state: State, evaluations: Evaluations) -> None:
                 explanation=explanation,
                 metadata_={},
             )
-            await state.enqueue(document_annotation)
+            await state.enqueue_annotations(document_annotation)
     elif len(names) == 1 and names[0] in ("context.span_id", "span_id"):
         for index, row in dataframe.iterrows():
             score, label, explanation = _get_annotation_result(row)
@@ -235,7 +235,7 @@ async def _add_evaluations(state: State, evaluations: Evaluations) -> None:
                 explanation=explanation,
                 metadata_={},
             )
-            await state.enqueue(span_annotation)
+            await state.enqueue_annotations(span_annotation)
     elif len(names) == 1 and names[0] in ("context.trace_id", "trace_id"):
         for index, row in dataframe.iterrows():
             score, label, explanation = _get_annotation_result(row)
@@ -249,7 +249,7 @@ async def _add_evaluations(state: State, evaluations: Evaluations) -> None:
                 explanation=explanation,
                 metadata_={},
             )
-            await state.enqueue(trace_annotation)
+            await state.enqueue_annotations(trace_annotation)
 
 
 def _get_annotation_result(

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

@@ -27,7 +27,7 @@ class ExperimentRun(V1RoutesBaseModel):
         description="The ID of the dataset example used in the experiment run"
     )
     output: Any = Field(description="The output of the experiment task")
-    repetition_number: int = Field(description="The repetition number of the experiment run")
+    repetition_number: int = Field(description="The repetition number of the experiment run", gt=0)
     start_time: datetime = Field(description="The start time of the experiment run")
     end_time: datetime = Field(description="The end time of the experiment run")
     trace_id: Optional[str] = Field(

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

@@ -46,7 +46,7 @@ class Experiment(V1RoutesBaseModel):
     dataset_version_id: str = Field(
         description="The ID of the dataset version associated with the experiment"
     )
-    repetitions: int = Field(description="Number of times the experiment is repeated")
+    repetitions: int = Field(description="Number of times the experiment is repeated", gt=0)
     metadata: dict[str, Any] = Field(description="Metadata of the experiment")
     project_name: Optional[str] = Field(
         description="The name of the project associated with the experiment"

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

@@ -897,7 +897,7 @@ async def annotate_spans(
         )
     precursors = [d.as_precursor(user_id=user_id) for d in filtered_span_annotations]
     if not sync:
-        await request.state.enqueue(*precursors)
+        await request.state.enqueue_annotations(*precursors)
         return AnnotateSpansResponseBody(data=[])
 
     span_ids = {p.span_id for p in precursors}
@@ -1072,7 +1072,7 @@ async def create_spans(
 
     # All spans are valid, queue them all
     for span_for_insertion, project_name in spans_to_queue:
-        await request.state.queue_span_for_bulk_insert(span_for_insertion, project_name)
+        await request.state.enqueue_span(span_for_insertion, project_name)
 
     return CreateSpansResponseBody(
         total_received=total_received,

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

@@ -18,6 +18,7 @@ from starlette.status import (
     HTTP_404_NOT_FOUND,
     HTTP_415_UNSUPPORTED_MEDIA_TYPE,
     HTTP_422_UNPROCESSABLE_ENTITY,
+    HTTP_503_SERVICE_UNAVAILABLE,
 )
 from strawberry.relay import GlobalID
 
@@ -29,6 +30,7 @@ from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.authorization import is_not_locked
 from phoenix.server.bearer_auth import PhoenixUser
 from phoenix.server.dml_event import SpanDeleteEvent, TraceAnnotationInsertEvent
+from phoenix.server.prometheus import SPAN_QUEUE_REJECTIONS
 from phoenix.trace.otel import decode_otlp_span
 from phoenix.utilities.project import get_project_name
 
@@ -42,9 +44,18 @@ from .utils import (
 router = APIRouter(tags=["traces"])
 
 
+def is_not_at_capacity(request: Request) -> None:
+    if request.app.state.span_queue_is_full():
+        SPAN_QUEUE_REJECTIONS.inc()
+        raise HTTPException(
+            detail="Server is at capacity and cannot process more requests",
+            status_code=HTTP_503_SERVICE_UNAVAILABLE,
+        )
+
+
 @router.post(
     "/traces",
-    dependencies=[Depends(is_not_locked)],
+    dependencies=[Depends(is_not_locked), Depends(is_not_at_capacity)],
     operation_id="addTraces",
     summary="Send traces",
     responses=add_errors_to_responses(
@@ -56,6 +67,10 @@ router = APIRouter(tags=["traces"])
                 ),
             },
             {"status_code": HTTP_422_UNPROCESSABLE_ENTITY, "description": "Invalid request body"},
+            {
+                "status_code": HTTP_503_SERVICE_UNAVAILABLE,
+                "description": "Server is at capacity and cannot process more requests",
+            },
         ]
     ),
     openapi_extra={
@@ -145,7 +160,7 @@ async def annotate_traces(
 
     precursors = [d.as_precursor(user_id=user_id) for d in request_body.data]
     if not sync:
-        await request.state.enqueue(*precursors)
+        await request.state.enqueue_annotations(*precursors)
         return AnnotateTracesResponseBody(data=[])
 
     trace_ids = {p.trace_id for p in precursors}
@@ -193,7 +208,7 @@ async def _add_spans(req: ExportTraceServiceRequest, state: State) -> None:
         for scope_span in resource_spans.scope_spans:
             for otlp_span in scope_span.spans:
                 span = await run_in_threadpool(decode_otlp_span, otlp_span)
-                await state.queue_span_for_bulk_insert(span, project_name)
+                await state.enqueue_span(span, project_name)
 
 
 @router.delete(

phoenix/server/api/types/DatasetExample.py

@@ -10,8 +10,12 @@ from strawberry.types import Info
 
 from phoenix.db import models
 from phoenix.server.api.context import Context
+from phoenix.server.api.exceptions import BadRequest
 from phoenix.server.api.types.DatasetExampleRevision import DatasetExampleRevision
 from phoenix.server.api.types.DatasetVersion import DatasetVersion
+from phoenix.server.api.types.ExperimentRepeatedRunGroup import (
+    ExperimentRepeatedRunGroup,
+)
 from phoenix.server.api.types.ExperimentRun import ExperimentRun, to_gql_experiment_run
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.pagination import (
@@ -65,6 +69,7 @@ class DatasetExample(Node):
         last: Optional[int] = UNSET,
         after: Optional[CursorString] = UNSET,
         before: Optional[CursorString] = UNSET,
+        experiment_ids: Optional[list[GlobalID]] = UNSET,
     ) -> Connection[ExperimentRun]:
         args = ConnectionArgs(
             first=first,
@@ -78,8 +83,51 @@ class DatasetExample(Node):
             .options(joinedload(models.ExperimentRun.trace).load_only(models.Trace.trace_id))
             .join(models.Experiment, models.Experiment.id == models.ExperimentRun.experiment_id)
             .where(models.ExperimentRun.dataset_example_id == example_id)
-            .order_by(models.Experiment.id.desc())
+            .order_by(
+                models.ExperimentRun.experiment_id.asc(),
+                models.ExperimentRun.repetition_number.asc(),
+            )
         )
+        if experiment_ids:
+            experiment_db_ids = [
+                from_global_id_with_expected_type(
+                    global_id=experiment_id,
+                    expected_type_name=models.Experiment.__name__,
+                )
+                for experiment_id in experiment_ids or []
+            ]
+            query = query.where(models.ExperimentRun.experiment_id.in_(experiment_db_ids))
         async with info.context.db() as session:
             runs = (await session.scalars(query)).all()
         return connection_from_list([to_gql_experiment_run(run) for run in runs], args)
+
+    @strawberry.field
+    async def experiment_repeated_run_groups(
+        self,
+        info: Info[Context, None],
+        experiment_ids: list[GlobalID],
+    ) -> list[ExperimentRepeatedRunGroup]:
+        example_rowid = self.id_attr
+        experiment_rowids = []
+        for experiment_id in experiment_ids:
+            try:
+                experiment_rowid = from_global_id_with_expected_type(
+                    global_id=experiment_id,
+                    expected_type_name=models.Experiment.__name__,
+                )
+            except Exception:
+                raise BadRequest(f"Invalid experiment ID: {experiment_id}")
+            experiment_rowids.append(experiment_rowid)
+        repeated_run_groups = (
+            await info.context.data_loaders.experiment_repeated_run_groups.load_many(
+                [(experiment_rowid, example_rowid) for experiment_rowid in experiment_rowids]
+            )
+        )
+        return [
+            ExperimentRepeatedRunGroup(
+                experiment_rowid=group.experiment_rowid,
+                dataset_example_rowid=group.dataset_example_rowid,
+                runs=[to_gql_experiment_run(run) for run in group.runs],
+            )
+            for group in repeated_run_groups
+        ]

phoenix/server/api/types/Experiment.py

@@ -5,13 +5,14 @@ import strawberry
 from sqlalchemy import func, select
 from sqlalchemy.orm import joinedload
 from strawberry import UNSET, Private
-from strawberry.relay import Connection, Node, NodeID
+from strawberry.relay import Connection, GlobalID, Node, NodeID
 from strawberry.scalars import JSON
 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.DatasetVersion import DatasetVersion
 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 (
@@ -32,6 +33,7 @@ class Experiment(Node):
     name: str
     project_name: Optional[str]
     description: Optional[str]
+    dataset_version_id: GlobalID
     metadata: JSON
     created_at: datetime
     updated_at: datetime
@@ -71,7 +73,10 @@ class Experiment(Node):
                 await session.scalars(
                     select(models.ExperimentRun)
                     .where(models.ExperimentRun.experiment_id == experiment_id)
-                    .order_by(models.ExperimentRun.id.desc())
+                    .order_by(
+                        models.ExperimentRun.dataset_example_id.asc(),
+                        models.ExperimentRun.repetition_number.asc(),
+                    )
                     .options(
                         joinedload(models.ExperimentRun.trace).load_only(models.Trace.trace_id)
                     )
@@ -187,6 +192,10 @@ class Experiment(Node):
                 async for token_type, is_prompt, cost, tokens in data
             ]
 
+    @strawberry.field
+    async def repetition_count(self, info: Info[Context, None]) -> int:
+        return await info.context.data_loaders.experiment_repetition_counts.load(self.id_attr)
+
 
 def to_gql_experiment(
     experiment: models.Experiment,
@@ -201,6 +210,7 @@ def to_gql_experiment(
         name=experiment.name,
         project_name=experiment.project_name,
         description=experiment.description,
+        dataset_version_id=GlobalID(DatasetVersion.__name__, str(experiment.dataset_version_id)),
         metadata=experiment.metadata_,
         created_at=experiment.created_at,
         updated_at=experiment.updated_at,

phoenix/server/api/types/ExperimentComparison.py

@@ -1,18 +1,12 @@
 import strawberry
-from strawberry.relay import GlobalID, Node, NodeID
+from strawberry.relay import Node, NodeID
 
 from phoenix.server.api.types.DatasetExample import DatasetExample
-from phoenix.server.api.types.ExperimentRun import ExperimentRun
-
-
-@strawberry.type
-class RunComparisonItem:
-    experiment_id: GlobalID
-    runs: list[ExperimentRun]
+from phoenix.server.api.types.ExperimentRepeatedRunGroup import ExperimentRepeatedRunGroup
 
 
 @strawberry.type
 class ExperimentComparison(Node):
     id_attr: NodeID[int]
     example: DatasetExample
-    run_comparison_items: list[RunComparisonItem]
+    repeated_run_groups: list[ExperimentRepeatedRunGroup]

phoenix/server/api/types/ExperimentRepeatedRunGroup.py

@@ -0,0 +1,146 @@
+import re
+from base64 import b64decode
+from typing import Optional
+
+import strawberry
+from sqlalchemy import func, select
+from strawberry.relay import GlobalID, Node
+from strawberry.types import Info
+from typing_extensions import Self, TypeAlias
+
+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.ExperimentRepeatedRunGroupAnnotationSummary import (
+    ExperimentRepeatedRunGroupAnnotationSummary,
+)
+from phoenix.server.api.types.ExperimentRun import ExperimentRun
+from phoenix.server.api.types.SpanCostDetailSummaryEntry import SpanCostDetailSummaryEntry
+from phoenix.server.api.types.SpanCostSummary import SpanCostSummary
+
+ExperimentRowId: TypeAlias = int
+DatasetExampleRowId: TypeAlias = int
+
+
+@strawberry.type
+class ExperimentRepeatedRunGroup(Node):
+    experiment_rowid: strawberry.Private[ExperimentRowId]
+    dataset_example_rowid: strawberry.Private[DatasetExampleRowId]
+    runs: list[ExperimentRun]
+
+    @classmethod
+    def resolve_id(
+        cls,
+        root: Self,
+        *,
+        info: Info,
+    ) -> str:
+        return (
+            f"experiment_id={root.experiment_rowid}:dataset_example_id={root.dataset_example_rowid}"
+        )
+
+    @strawberry.field
+    def experiment_id(self) -> strawberry.ID:
+        from phoenix.server.api.types.Experiment import Experiment
+
+        return strawberry.ID(str(GlobalID(Experiment.__name__, str(self.experiment_rowid))))
+
+    @strawberry.field
+    async def average_latency_ms(self, info: Info[Context, None]) -> Optional[float]:
+        return await info.context.data_loaders.average_experiment_repeated_run_group_latency.load(
+            (self.experiment_rowid, self.dataset_example_rowid)
+        )
+
+    @strawberry.field
+    async def cost_summary(self, info: Info[Context, None]) -> SpanCostSummary:
+        experiment_id = self.experiment_rowid
+        example_id = self.dataset_example_rowid
+        summary = (
+            await info.context.data_loaders.span_cost_summary_by_experiment_repeated_run_group.load(
+                (experiment_id, example_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.experiment_rowid
+        example_id = self.dataset_example_rowid
+        stmt = (
+            select(
+                models.SpanCostDetail.token_type,
+                models.SpanCostDetail.is_prompt,
+                func.sum(models.SpanCostDetail.cost).label("cost"),
+                func.sum(models.SpanCostDetail.tokens).label("tokens"),
+            )
+            .select_from(models.SpanCostDetail)
+            .join(models.SpanCost, models.SpanCostDetail.span_cost_id == models.SpanCost.id)
+            .join(models.Trace, models.SpanCost.trace_rowid == models.Trace.id)
+            .join(models.ExperimentRun, models.ExperimentRun.trace_id == models.Trace.trace_id)
+            .where(models.ExperimentRun.experiment_id == experiment_id)
+            .where(models.ExperimentRun.dataset_example_id == example_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
+            ]
+
+    @strawberry.field
+    async def annotation_summaries(
+        self,
+        info: Info[Context, None],
+    ) -> list[ExperimentRepeatedRunGroupAnnotationSummary]:
+        loader = info.context.data_loaders.experiment_repeated_run_group_annotation_summaries
+        summaries = await loader.load((self.experiment_rowid, self.dataset_example_rowid))
+        return [
+            ExperimentRepeatedRunGroupAnnotationSummary(
+                annotation_name=summary.annotation_name,
+                mean_score=summary.mean_score,
+            )
+            for summary in summaries
+        ]
+
+
+_EXPERIMENT_REPEATED_RUN_GROUP_NODE_ID_PATTERN = re.compile(
+    r"ExperimentRepeatedRunGroup:experiment_id=(\d+):dataset_example_id=(\d+)"
+)
+
+
+def parse_experiment_repeated_run_group_node_id(
+    node_id: str,
+) -> tuple[ExperimentRowId, DatasetExampleRowId]:
+    decoded_node_id = _base64_decode(node_id)
+    match = re.match(_EXPERIMENT_REPEATED_RUN_GROUP_NODE_ID_PATTERN, decoded_node_id)
+    if not match:
+        raise ValueError(f"Invalid node ID format: {node_id}")
+
+    experiment_id = int(match.group(1))
+    dataset_example_id = int(match.group(2))
+    return experiment_id, dataset_example_id
+
+
+def _base64_decode(string: str) -> str:
+    return b64decode(string.encode()).decode()

phoenix/server/api/types/ExperimentRepeatedRunGroupAnnotationSummary.py

@@ -0,0 +1,9 @@
+from typing import Optional
+
+import strawberry
+
+
+@strawberry.type
+class ExperimentRepeatedRunGroupAnnotationSummary:
+    annotation_name: str
+    mean_score: Optional[float]

phoenix/server/api/types/ExperimentRun.py

@@ -3,7 +3,6 @@ from typing import TYPE_CHECKING, Annotated, Optional
 
 import strawberry
 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
@@ -34,12 +33,17 @@ if TYPE_CHECKING:
 class ExperimentRun(Node):
     id_attr: NodeID[int]
     experiment_id: GlobalID
+    repetition_number: int
     trace_id: Optional[str]
     output: Optional[JSON]
     start_time: datetime
     end_time: datetime
     error: Optional[str]
 
+    @strawberry.field
+    def latency_ms(self) -> float:
+        return (self.end_time - self.start_time).total_seconds() * 1000
+
     @strawberry.field
     async def annotations(
         self,
@@ -78,24 +82,12 @@ class ExperimentRun(Node):
     ]:  # use lazy types to avoid circular import: https://strawberry.rocks/docs/types/lazy
         from phoenix.server.api.types.DatasetExample import DatasetExample
 
-        async with info.context.db() as session:
-            assert (
-                result := await session.execute(
-                    select(models.DatasetExample, models.Experiment.dataset_version_id)
-                    .select_from(models.ExperimentRun)
-                    .join(
-                        models.DatasetExample,
-                        models.DatasetExample.id == models.ExperimentRun.dataset_example_id,
-                    )
-                    .join(
-                        models.Experiment,
-                        models.Experiment.id == models.ExperimentRun.experiment_id,
-                    )
-                    .where(models.ExperimentRun.id == self.id_attr)
-                    .options(load_only(models.DatasetExample.id, models.DatasetExample.created_at))
-                )
-            ) is not None
-            example, version_id = result.first()
+        (
+            example,
+            version_id,
+        ) = await info.context.data_loaders.dataset_examples_and_versions_by_experiment_run.load(
+            self.id_attr
+        )
         return DatasetExample(
             id_attr=example.id,
             created_at=example.created_at,
@@ -165,6 +157,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)),
+        repetition_number=run.repetition_number,
         trace_id=run.trace.trace_id if run.trace else None,
         output=run.output.get("task_output"),
         start_time=run.start_time,

phoenix/server/api/types/Prompt.py

@@ -19,6 +19,7 @@ from phoenix.server.api.types.pagination import (
     connection_from_list,
 )
 
+from .PromptLabel import PromptLabel, to_gql_prompt_label
 from .PromptVersion import (
     PromptVersion,
     to_gql_prompt_version,
@@ -116,6 +117,16 @@ class Prompt(Node):
                 raise NotFound(f"Source prompt not found: {self.source_prompt_id}")
             return to_gql_prompt_from_orm(source_prompt)
 
+    @strawberry.field
+    async def labels(self, info: Info[Context, None]) -> list["PromptLabel"]:
+        async with info.context.db() as session:
+            labels = await session.scalars(
+                select(models.PromptLabel)
+                .join(models.PromptPromptLabel)
+                .where(models.PromptPromptLabel.prompt_id == self.id_attr)
+            )
+            return [to_gql_prompt_label(label) for label in labels]
+
 
 def to_gql_prompt_from_orm(orm_model: "models.Prompt") -> Prompt:
     if not orm_model.source_prompt_id:

phoenix/server/api/types/PromptLabel.py

@@ -1,14 +1,10 @@
 from typing import Optional
 
 import strawberry
-from sqlalchemy import select
 from strawberry.relay import Node, NodeID
-from strawberry.types import Info
 
 from phoenix.db import models
-from phoenix.server.api.context import Context
 from phoenix.server.api.types.Identifier import Identifier
-from phoenix.server.api.types.Prompt import Prompt, to_gql_prompt_from_orm
 
 
 @strawberry.type
@@ -16,21 +12,7 @@ class PromptLabel(Node):
     id_attr: NodeID[int]
     name: Identifier
     description: Optional[str] = None
-
-    @strawberry.field
-    async def prompts(self, info: Info[Context, None]) -> list[Prompt]:
-        async with info.context.db() as session:
-            statement = (
-                select(models.Prompt)
-                .join(
-                    models.PromptPromptLabel, models.Prompt.id == models.PromptPromptLabel.prompt_id
-                )
-                .where(models.PromptPromptLabel.prompt_label_id == self.id_attr)
-            )
-            return [
-                to_gql_prompt_from_orm(prompt_orm)
-                async for prompt_orm in await session.stream_scalars(statement)
-            ]
+    color: str
 
 
 def to_gql_prompt_label(label_orm: models.PromptLabel) -> PromptLabel:
@@ -38,4 +20,5 @@ def to_gql_prompt_label(label_orm: models.PromptLabel) -> PromptLabel:
         id_attr=label_orm.id,
         name=Identifier(label_orm.name),
         description=label_orm.description,
+        color=label_orm.color,
     )

phoenix/server/api/types/node.py

@@ -1,5 +1,15 @@
+import re
+from base64 import b64decode
+
 from strawberry.relay import GlobalID
 
+_GLOBAL_ID_PATTERN = re.compile(r"[a-zA-Z]+:[0-9]+")
+
+
+def is_global_id(node_id: str) -> bool:
+    decoded_node_id = b64decode(node_id).decode()
+    return _GLOBAL_ID_PATTERN.match(decoded_node_id) is not None
+
 
 def from_global_id(global_id: GlobalID) -> tuple[str, int]:
     """