arize-phoenix 4.5.0__py3-none-any.whl → 4.6.1__py3-none-any.whl

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

@@ -0,0 +1,302 @@
+from random import getrandbits
+
+from sqlalchemy import select
+from starlette.requests import Request
+from starlette.responses import JSONResponse, Response
+from starlette.status import HTTP_404_NOT_FOUND
+from strawberry.relay import GlobalID
+
+from phoenix.db import models
+from phoenix.db.helpers import SupportedSQLDialect
+from phoenix.db.insertion.helpers import insert_on_conflict
+from phoenix.server.api.types.node import from_global_id_with_expected_type
+
+
+def _short_uuid() -> str:
+    return str(getrandbits(32).to_bytes(4, "big").hex())
+
+
+def _generate_experiment_name(dataset_name: str) -> str:
+    """
+    Generate a semi-unique name for the experiment.
+    """
+    short_ds_name = dataset_name[:8].replace(" ", "-")
+    return f"{short_ds_name}-{_short_uuid()}"
+
+
+async def create_experiment(request: Request) -> Response:
+    """
+    summary: Create an experiment using a dataset
+    operationId: createExperiment
+    tags:
+      - private
+    parameters:
+      - in: path
+        name: dataset_id
+        required: true
+        description: The ID of the dataset to create an experiment for
+        schema:
+          type: string
+    requestBody:
+      description: Details of the experiment to be created
+      required: true
+      content:
+        application/json:
+          schema:
+            type: object
+            properties:
+              repetitions:
+                type: integer
+                description: Number of times the experiment should be repeated for each example
+                default: 1
+              metadata:
+                type: object
+                description: Metadata for the experiment
+                additionalProperties:
+                  type: string
+              version_id:
+                type: string
+                description: ID of the dataset version to use
+    responses:
+      200:
+        description: Experiment retrieved successfully
+        content:
+          application/json:
+            schema:
+              type: object
+              properties:
+                data:
+                  type: object
+                  properties:
+                    id:
+                      type: string
+                      description: The ID of the experiment
+                    dataset_id:
+                      type: string
+                      description: The ID of the dataset associated with the experiment
+                    dataset_version_id:
+                      type: string
+                      description: The ID of the dataset version associated with the experiment
+                    repetitions:
+                      type: integer
+                      description: Number of times the experiment is repeated
+                    metadata:
+                      type: object
+                      description: Metadata of the experiment
+                      additionalProperties:
+                        type: string
+                    project_name:
+                      type: string
+                      description: The name of the project associated with the experiment
+                    created_at:
+                      type: string
+                      format: date-time
+                      description: The creation timestamp of the experiment
+                    updated_at:
+                      type: string
+                      format: date-time
+                      description: The last update timestamp of the experiment
+      404:
+        description: Dataset or DatasetVersion not found
+    """
+    dataset_globalid = GlobalID.from_id(request.path_params["dataset_id"])
+    try:
+        dataset_id = from_global_id_with_expected_type(dataset_globalid, "Dataset")
+    except ValueError:
+        return Response(
+            content="Dataset with ID {dataset_globalid} does not exist",
+            status_code=HTTP_404_NOT_FOUND,
+        )
+
+    payload = await request.json()
+    repetitions = payload.get("repetitions", 1)
+    metadata = payload.get("metadata") or {}
+    dataset_version_globalid_str = payload.get("version_id")
+    if dataset_version_globalid_str is not None:
+        try:
+            dataset_version_globalid = GlobalID.from_id(dataset_version_globalid_str)
+            dataset_version_id = from_global_id_with_expected_type(
+                dataset_version_globalid, "DatasetVersion"
+            )
+        except ValueError:
+            return Response(
+                content="DatasetVersion with ID {dataset_version_globalid} does not exist",
+                status_code=HTTP_404_NOT_FOUND,
+            )
+
+    async with request.app.state.db() as session:
+        result = (
+            await session.execute(select(models.Dataset).where(models.Dataset.id == dataset_id))
+        ).scalar()
+        if result is None:
+            return Response(
+                content=f"Dataset with ID {dataset_globalid} does not exist",
+                status_code=HTTP_404_NOT_FOUND,
+            )
+        dataset_name = result.name
+        if dataset_version_globalid_str is None:
+            dataset_version_result = await session.execute(
+                select(models.DatasetVersion)
+                .where(models.DatasetVersion.dataset_id == dataset_id)
+                .order_by(models.DatasetVersion.id.desc())
+            )
+            dataset_version = dataset_version_result.scalar()
+            if not dataset_version:
+                return Response(
+                    content=f"Dataset {dataset_globalid} does not have any versions",
+                    status_code=HTTP_404_NOT_FOUND,
+                )
+            dataset_version_id = dataset_version.id
+            dataset_version_globalid = GlobalID("DatasetVersion", str(dataset_version_id))
+        else:
+            dataset_version = await session.execute(
+                select(models.DatasetVersion).where(models.DatasetVersion.id == dataset_version_id)
+            )
+            dataset_version = dataset_version.scalar()
+            if not dataset_version:
+                return Response(
+                    content=f"DatasetVersion with ID {dataset_version_globalid} does not exist",
+                    status_code=HTTP_404_NOT_FOUND,
+                )
+
+        # generate a semi-unique name for the experiment
+        experiment_name = payload.get("name") or _generate_experiment_name(dataset_name)
+        project_name = f"Experiment-{getrandbits(96).to_bytes(12, 'big').hex()}"
+        project_description = (
+            f"dataset_id: {dataset_globalid}\ndataset_version_id: {dataset_version_globalid}"
+        )
+        experiment = models.Experiment(
+            dataset_id=int(dataset_id),
+            dataset_version_id=int(dataset_version_id),
+            name=experiment_name,
+            description=payload.get("description"),
+            repetitions=repetitions,
+            metadata_=metadata,
+            project_name=project_name,
+        )
+        session.add(experiment)
+        await session.flush()
+
+        dialect = SupportedSQLDialect(session.bind.dialect.name)
+        project_rowid = await session.scalar(
+            insert_on_conflict(
+                dialect=dialect,
+                table=models.Project,
+                constraint="uq_projects_name",
+                column_names=("name",),
+                values=dict(
+                    name=project_name,
+                    description=project_description,
+                    created_at=experiment.created_at,
+                    updated_at=experiment.updated_at,
+                ),
+            ).returning(models.Project.id)
+        )
+        assert project_rowid is not None
+
+        experiment_globalid = GlobalID("Experiment", str(experiment.id))
+        if dataset_version_globalid_str is None:
+            dataset_version_globalid = GlobalID(
+                "DatasetVersion", str(experiment.dataset_version_id)
+            )
+        experiment_payload = {
+            "id": str(experiment_globalid),
+            "dataset_id": str(dataset_globalid),
+            "dataset_version_id": str(dataset_version_globalid),
+            "repetitions": experiment.repetitions,
+            "metadata": experiment.metadata_,
+            "project_name": experiment.project_name,
+            "created_at": experiment.created_at.isoformat(),
+            "updated_at": experiment.updated_at.isoformat(),
+        }
+    return JSONResponse(content={"data": experiment_payload})
+
+
+async def read_experiment(request: Request) -> Response:
+    """
+    summary: Get details of a specific experiment
+    operationId: getExperiment
+    tags:
+      - private
+    parameters:
+      - in: path
+        name: experiment_id
+        required: true
+        description: The ID of the experiment to retrieve
+        schema:
+          type: string
+    responses:
+      200:
+        description: Experiment retrieved successfully
+        content:
+          application/json:
+            schema:
+              type: object
+              properties:
+                data:
+                  type: object
+                  properties:
+                    id:
+                      type: string
+                      description: The ID of the experiment
+                    dataset_id:
+                      type: string
+                      description: The ID of the dataset associated with the experiment
+                    dataset_version_id:
+                      type: string
+                      description: The ID of the dataset version associated with the experiment
+                    repetitions:
+                      type: integer
+                      description: Number of times the experiment is repeated
+                    metadata:
+                      type: object
+                      description: Metadata of the experiment
+                      additionalProperties:
+                        type: string
+                    project_name:
+                      type: string
+                      description: The name of the project associated with the experiment
+                    created_at:
+                      type: string
+                      format: date-time
+                      description: The creation timestamp of the experiment
+                    updated_at:
+                      type: string
+                      format: date-time
+                      description: The last update timestamp of the experiment
+      404:
+        description: Experiment not found
+    """
+    experiment_globalid = GlobalID.from_id(request.path_params["experiment_id"])
+    try:
+        experiment_id = from_global_id_with_expected_type(experiment_globalid, "Experiment")
+    except ValueError:
+        return Response(
+            content="Experiment with ID {experiment_globalid} does not exist",
+            status_code=HTTP_404_NOT_FOUND,
+        )
+
+    async with request.app.state.db() as session:
+        experiment = await session.execute(
+            select(models.Experiment).where(models.Experiment.id == experiment_id)
+        )
+        experiment = experiment.scalar()
+        if not experiment:
+            return Response(
+                content=f"Experiment with ID {experiment_globalid} does not exist",
+                status_code=HTTP_404_NOT_FOUND,
+            )
+
+        dataset_globalid = GlobalID("Dataset", str(experiment.dataset_id))
+        dataset_version_globalid = GlobalID("DatasetVersion", str(experiment.dataset_version_id))
+        experiment_payload = {
+            "id": str(experiment_globalid),
+            "dataset_id": str(dataset_globalid),
+            "dataset_version_id": str(dataset_version_globalid),
+            "repetitions": experiment.repetitions,
+            "metadata": experiment.metadata_,
+            "project_name": experiment.project_name,
+            "created_at": experiment.created_at.isoformat(),
+            "updated_at": experiment.updated_at.isoformat(),
+        }
+    return JSONResponse(content={"data": experiment_payload})

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

@@ -19,9 +19,9 @@ async def query_spans_handler(request: Request) -> Response:
     summary: Query spans using query DSL
     operationId: querySpans
     tags:
-      - spans
+      - private
     parameters:
-      - name: project-name
+      - name: project_name
         in: query
         schema:
           type: string
@@ -68,6 +68,8 @@ async def query_spans_handler(request: Request) -> Response:
     responses:
       200:
         description: Success
+      403:
+        description: Forbidden
       404:
         description: Not found
       422:
@@ -76,9 +78,11 @@ async def query_spans_handler(request: Request) -> Response:
     payload = await request.json()
     queries = payload.pop("queries", [])
     project_name = (
-        request.query_params.get("project-name")
-        # read from headers/payload for backward-compatibility
-        or request.headers.get("project-name")
+        request.query_params.get("project_name")
+        or request.query_params.get("project-name")  # for backward compatibility
+        or request.headers.get(
+            "project-name"
+        )  # read from headers/payload for backward-compatibility
         or payload.get("project_name")
         or DEFAULT_PROJECT_NAME
     )

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

@@ -11,7 +11,6 @@ from starlette.datastructures import State
 from starlette.requests import Request
 from starlette.responses import Response
 from starlette.status import (
-    HTTP_403_FORBIDDEN,
     HTTP_415_UNSUPPORTED_MEDIA_TYPE,
     HTTP_422_UNPROCESSABLE_ENTITY,
 )
@@ -25,7 +24,7 @@ async def post_traces(request: Request) -> Response:
     summary: Send traces to Phoenix
     operationId: addTraces
     tags:
-      - traces
+      - private
     requestBody:
       required: true
       content:
@@ -43,8 +42,6 @@ async def post_traces(request: Request) -> Response:
       422:
         description: Request body is invalid
     """
-    if request.app.state.read_only:
-        return Response(status_code=HTTP_403_FORBIDDEN)
     content_type = request.headers.get("content-type")
     if content_type != "application/x-protobuf":
         return Response(

phoenix/server/api/schema.py

@@ -1,308 +1,7 @@
-from collections import defaultdict
-from typing import Dict, List, Optional, Set, Union
-
-import numpy as np
-import numpy.typing as npt
 import strawberry
-from sqlalchemy import delete, select
-from sqlalchemy.orm import contains_eager, load_only
-from strawberry import ID, UNSET
-from strawberry.types import Info
-from typing_extensions import Annotated
-
-from phoenix.config import DEFAULT_PROJECT_NAME
-from phoenix.db import models
-from phoenix.db.insertion.span import ClearProjectSpansEvent
-from phoenix.pointcloud.clustering import Hdbscan
-from phoenix.server.api.context import Context
-from phoenix.server.api.helpers import ensure_list
-from phoenix.server.api.input_types.ClusterInput import ClusterInput
-from phoenix.server.api.input_types.Coordinates import (
-    InputCoordinate2D,
-    InputCoordinate3D,
-)
-from phoenix.server.api.types.Cluster import Cluster, to_gql_clusters
-from phoenix.server.api.types.DatasetRole import AncillaryDatasetRole, DatasetRole
-from phoenix.server.api.types.Dimension import to_gql_dimension
-from phoenix.server.api.types.EmbeddingDimension import (
-    DEFAULT_CLUSTER_SELECTION_EPSILON,
-    DEFAULT_MIN_CLUSTER_SIZE,
-    DEFAULT_MIN_SAMPLES,
-    to_gql_embedding_dimension,
-)
-from phoenix.server.api.types.Event import create_event_id, unpack_event_id
-from phoenix.server.api.types.ExportEventsMutation import ExportEventsMutation
-from phoenix.server.api.types.Functionality import Functionality
-from phoenix.server.api.types.Model import Model
-from phoenix.server.api.types.node import (
-    GlobalID,
-    Node,
-    from_global_id,
-    from_global_id_with_expected_type,
-)
-from phoenix.server.api.types.pagination import (
-    Connection,
-    ConnectionArgs,
-    CursorString,
-    connection_from_list,
-)
-from phoenix.server.api.types.Project import Project
-from phoenix.server.api.types.Span import to_gql_span
-from phoenix.server.api.types.Trace import Trace
-
-
-@strawberry.type
-class Query:
-    @strawberry.field
-    async def projects(
-        self,
-        info: Info[Context, None],
-        first: Optional[int] = 50,
-        last: Optional[int] = UNSET,
-        after: Optional[CursorString] = UNSET,
-        before: Optional[CursorString] = UNSET,
-    ) -> Connection[Project]:
-        args = ConnectionArgs(
-            first=first,
-            after=after if isinstance(after, CursorString) else None,
-            last=last,
-            before=before if isinstance(before, CursorString) else None,
-        )
-        async with info.context.db() as session:
-            projects = await session.scalars(select(models.Project))
-        data = [
-            Project(
-                id_attr=project.id,
-                name=project.name,
-                gradient_start_color=project.gradient_start_color,
-                gradient_end_color=project.gradient_end_color,
-            )
-            for project in projects
-        ]
-        return connection_from_list(data=data, args=args)
-
-    @strawberry.field
-    async def functionality(self, info: Info[Context, None]) -> "Functionality":
-        has_model_inferences = not info.context.model.is_empty
-        async with info.context.db() as session:
-            has_traces = (await session.scalar(select(models.Trace).limit(1))) is not None
-        return Functionality(
-            model_inferences=has_model_inferences,
-            tracing=has_traces,
-        )
-
-    @strawberry.field
-    def model(self) -> Model:
-        return Model()
-
-    @strawberry.field
-    async def node(self, id: GlobalID, info: Info[Context, None]) -> Node:
-        type_name, node_id = from_global_id(str(id))
-        if type_name == "Dimension":
-            dimension = info.context.model.scalar_dimensions[node_id]
-            return to_gql_dimension(node_id, dimension)
-        elif type_name == "EmbeddingDimension":
-            embedding_dimension = info.context.model.embedding_dimensions[node_id]
-            return to_gql_embedding_dimension(node_id, embedding_dimension)
-        elif type_name == "Project":
-            project_stmt = select(
-                models.Project.id,
-                models.Project.name,
-                models.Project.gradient_start_color,
-                models.Project.gradient_end_color,
-            ).where(models.Project.id == node_id)
-            async with info.context.db() as session:
-                project = (await session.execute(project_stmt)).first()
-            if project is None:
-                raise ValueError(f"Unknown project: {id}")
-            return Project(
-                id_attr=project.id,
-                name=project.name,
-                gradient_start_color=project.gradient_start_color,
-                gradient_end_color=project.gradient_end_color,
-            )
-        elif type_name == "Trace":
-            trace_stmt = select(models.Trace.id).where(models.Trace.id == node_id)
-            async with info.context.db() as session:
-                id_attr = await session.scalar(trace_stmt)
-            if id_attr is None:
-                raise ValueError(f"Unknown trace: {id}")
-            return Trace(id_attr=id_attr)
-        elif type_name == "Span":
-            span_stmt = (
-                select(models.Span)
-                .join(models.Trace)
-                .options(contains_eager(models.Span.trace))
-                .where(models.Span.id == node_id)
-            )
-            async with info.context.db() as session:
-                span = await session.scalar(span_stmt)
-            if span is None:
-                raise ValueError(f"Unknown span: {id}")
-            return to_gql_span(span)
-        raise Exception(f"Unknown node type: {type_name}")
-
-    @strawberry.field
-    def clusters(
-        self,
-        clusters: List[ClusterInput],
-    ) -> List[Cluster]:
-        clustered_events: Dict[str, Set[ID]] = defaultdict(set)
-        for i, cluster in enumerate(clusters):
-            clustered_events[cluster.id or str(i)].update(cluster.event_ids)
-        return to_gql_clusters(
-            clustered_events=clustered_events,
-        )
-
-    @strawberry.field
-    def hdbscan_clustering(
-        self,
-        info: Info[Context, None],
-        event_ids: Annotated[
-            List[ID],
-            strawberry.argument(
-                description="Event ID of the coordinates",
-            ),
-        ],
-        coordinates_2d: Annotated[
-            Optional[List[InputCoordinate2D]],
-            strawberry.argument(
-                description="Point coordinates. Must be either 2D or 3D.",
-            ),
-        ] = UNSET,
-        coordinates_3d: Annotated[
-            Optional[List[InputCoordinate3D]],
-            strawberry.argument(
-                description="Point coordinates. Must be either 2D or 3D.",
-            ),
-        ] = UNSET,
-        min_cluster_size: Annotated[
-            int,
-            strawberry.argument(
-                description="HDBSCAN minimum cluster size",
-            ),
-        ] = DEFAULT_MIN_CLUSTER_SIZE,
-        cluster_min_samples: Annotated[
-            int,
-            strawberry.argument(
-                description="HDBSCAN minimum samples",
-            ),
-        ] = DEFAULT_MIN_SAMPLES,
-        cluster_selection_epsilon: Annotated[
-            float,
-            strawberry.argument(
-                description="HDBSCAN cluster selection epsilon",
-            ),
-        ] = DEFAULT_CLUSTER_SELECTION_EPSILON,
-    ) -> List[Cluster]:
-        coordinates_3d = ensure_list(coordinates_3d)
-        coordinates_2d = ensure_list(coordinates_2d)
-
-        if len(coordinates_3d) > 0 and len(coordinates_2d) > 0:
-            raise ValueError("must specify only one of 2D or 3D coordinates")
-
-        if len(coordinates_3d) > 0:
-            coordinates = list(
-                map(
-                    lambda coord: np.array(
-                        [coord.x, coord.y, coord.z],
-                    ),
-                    coordinates_3d,
-                )
-            )
-        else:
-            coordinates = list(
-                map(
-                    lambda coord: np.array(
-                        [coord.x, coord.y],
-                    ),
-                    coordinates_2d,
-                )
-            )
-
-        if len(event_ids) != len(coordinates):
-            raise ValueError(
-                f"length mismatch between "
-                f"event_ids ({len(event_ids)}) "
-                f"and coordinates ({len(coordinates)})"
-            )
-
-        if len(event_ids) == 0:
-            return []
-
-        grouped_event_ids: Dict[
-            Union[DatasetRole, AncillaryDatasetRole],
-            List[ID],
-        ] = defaultdict(list)
-        grouped_coordinates: Dict[
-            Union[DatasetRole, AncillaryDatasetRole],
-            List[npt.NDArray[np.float64]],
-        ] = defaultdict(list)
-
-        for event_id, coordinate in zip(event_ids, coordinates):
-            row_id, dataset_role = unpack_event_id(event_id)
-            grouped_coordinates[dataset_role].append(coordinate)
-            grouped_event_ids[dataset_role].append(create_event_id(row_id, dataset_role))
-
-        stacked_event_ids = (
-            grouped_event_ids[DatasetRole.primary]
-            + grouped_event_ids[DatasetRole.reference]
-            + grouped_event_ids[AncillaryDatasetRole.corpus]
-        )
-        stacked_coordinates = np.stack(
-            grouped_coordinates[DatasetRole.primary]
-            + grouped_coordinates[DatasetRole.reference]
-            + grouped_coordinates[AncillaryDatasetRole.corpus]
-        )
-
-        clusters = Hdbscan(
-            min_cluster_size=min_cluster_size,
-            min_samples=cluster_min_samples,
-            cluster_selection_epsilon=cluster_selection_epsilon,
-        ).find_clusters(stacked_coordinates)
-
-        clustered_events = {
-            str(i): {stacked_event_ids[row_idx] for row_idx in cluster}
-            for i, cluster in enumerate(clusters)
-        }
-
-        return to_gql_clusters(
-            clustered_events=clustered_events,
-        )
-
-
-@strawberry.type
-class Mutation(ExportEventsMutation):
-    @strawberry.mutation
-    async def delete_project(self, info: Info[Context, None], id: GlobalID) -> Query:
-        if info.context.read_only:
-            return Query()
-        node_id = from_global_id_with_expected_type(str(id), "Project")
-        async with info.context.db() as session:
-            project = await session.scalar(
-                select(models.Project)
-                .where(models.Project.id == node_id)
-                .options(load_only(models.Project.name))
-            )
-            if project is None:
-                raise ValueError(f"Unknown project: {id}")
-            if project.name == DEFAULT_PROJECT_NAME:
-                raise ValueError(f"Cannot delete the {DEFAULT_PROJECT_NAME} project")
-            await session.delete(project)
-        return Query()
-
-    @strawberry.mutation
-    async def clear_project(self, info: Info[Context, None], id: GlobalID) -> Query:
-        if info.context.read_only:
-            return Query()
-        project_id = from_global_id_with_expected_type(str(id), "Project")
-        delete_statement = delete(models.Trace).where(models.Trace.project_rowid == project_id)
-        async with info.context.db() as session:
-            await session.execute(delete_statement)
-            if cache := info.context.cache_for_dataloaders:
-                cache.invalidate(ClearProjectSpansEvent(project_rowid=project_id))
-        return Query()
 
+from phoenix.server.api.mutations import Mutation
+from phoenix.server.api.queries import Query
 
 # This is the schema for generating `schema.graphql`.
 # See https://strawberry.rocks/docs/guides/schema-export

phoenix/server/api/types/AnnotatorKind.py

@@ -0,0 +1,10 @@
+from enum import Enum
+
+import strawberry
+
+
+@strawberry.enum
+class AnnotatorKind(Enum):
+    LLM = "LLM"
+    HUMAN = "HUMAN"
+    CODE = "CODE"