arize-phoenix 4.12.1rc1__py3-none-any.whl → 4.14.1__py3-none-any.whl

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

@@ -1,11 +1,8 @@
-from datetime import datetime
 from random import getrandbits
-from typing import Any, Dict, Optional
 
-from fastapi import APIRouter, HTTPException
-from pydantic import Field
 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
 
@@ -14,11 +11,6 @@ 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
 
-from .pydantic_compat import V1RoutesBaseModel
-from .utils import ResponseBody, add_errors_to_responses
-
-router = APIRouter(tags=["experiments"], include_in_schema=False)
-
 
 def _short_uuid() -> str:
     return str(getrandbits(32).to_bytes(4, "big").hex())
@@ -32,76 +24,94 @@ def _generate_experiment_name(dataset_name: str) -> str:
     return f"{short_ds_name}-{_short_uuid()}"
 
 
-class Experiment(V1RoutesBaseModel):
-    id: str = Field(description="The ID of the experiment")
-    dataset_id: str = Field(description="The ID of the dataset associated with the experiment")
-    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")
-    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"
-    )
-    created_at: datetime = Field(description="The creation timestamp of the experiment")
-    updated_at: datetime = Field(description="The last update timestamp of the experiment")
-
-
-class CreateExperimentRequestBody(V1RoutesBaseModel):
+async def create_experiment(request: Request) -> Response:
     """
-    Details of the experiment to be created
+    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
     """
-
-    name: Optional[str] = Field(
-        default=None,
-        description=("Name of the experiment (if omitted, a random name will be generated)"),
-    )
-    description: Optional[str] = Field(
-        default=None, description="An optional description of the experiment"
-    )
-    metadata: Optional[Dict[str, Any]] = Field(
-        default=None, description="Metadata for the experiment"
-    )
-    version_id: Optional[str] = Field(
-        default=None,
-        description=(
-            "ID of the dataset version over which the experiment will be run "
-            "(if omitted, the latest version will be used)"
-        ),
-    )
-    repetitions: int = Field(
-        default=1, description="Number of times the experiment should be repeated for each example"
-    )
-
-
-class CreateExperimentResponseBody(ResponseBody[Experiment]):
-    pass
-
-
-@router.post(
-    "/datasets/{dataset_id}/experiments",
-    operation_id="createExperiment",
-    summary="Create experiment on a dataset",
-    responses=add_errors_to_responses(
-        [{"status_code": HTTP_404_NOT_FOUND, "description": "Dataset or DatasetVersion not found"}]
-    ),
-    response_description="Experiment retrieved successfully",
-)
-async def create_experiment(
-    request: Request,
-    dataset_id: str,
-    request_body: CreateExperimentRequestBody,
-) -> CreateExperimentResponseBody:
-    dataset_globalid = GlobalID.from_id(dataset_id)
+    dataset_globalid = GlobalID.from_id(request.path_params["dataset_id"])
     try:
-        dataset_rowid = from_global_id_with_expected_type(dataset_globalid, "Dataset")
+        dataset_id = from_global_id_with_expected_type(dataset_globalid, "Dataset")
     except ValueError:
-        raise HTTPException(
-            detail="Dataset with ID {dataset_globalid} does not exist",
+        return Response(
+            content="Dataset with ID {dataset_globalid} does not exist",
             status_code=HTTP_404_NOT_FOUND,
         )
 
-    dataset_version_globalid_str = request_body.version_id
+    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)
@@ -109,31 +119,31 @@ async def create_experiment(
                 dataset_version_globalid, "DatasetVersion"
             )
         except ValueError:
-            raise HTTPException(
-                detail="DatasetVersion with ID {dataset_version_globalid} does not exist",
+            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_rowid))
+            await session.execute(select(models.Dataset).where(models.Dataset.id == dataset_id))
         ).scalar()
         if result is None:
-            raise HTTPException(
-                detail=f"Dataset with ID {dataset_globalid} does not exist",
+            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_rowid)
+                .where(models.DatasetVersion.dataset_id == dataset_id)
                 .order_by(models.DatasetVersion.id.desc())
             )
             dataset_version = dataset_version_result.scalar()
             if not dataset_version:
-                raise HTTPException(
-                    detail=f"Dataset {dataset_globalid} does not have any versions",
+                return Response(
+                    content=f"Dataset {dataset_globalid} does not have any versions",
                     status_code=HTTP_404_NOT_FOUND,
                 )
             dataset_version_id = dataset_version.id
@@ -144,24 +154,24 @@ async def create_experiment(
             )
             dataset_version = dataset_version.scalar()
             if not dataset_version:
-                raise HTTPException(
-                    detail=f"DatasetVersion with ID {dataset_version_globalid} does not exist",
+                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 = request_body.name or _generate_experiment_name(dataset_name)
+        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_rowid),
+            dataset_id=int(dataset_id),
             dataset_version_id=int(dataset_version_id),
             name=experiment_name,
-            description=request_body.description,
-            repetitions=request_body.repetitions,
-            metadata_=request_body.metadata or {},
+            description=payload.get("description"),
+            repetitions=repetitions,
+            metadata_=metadata,
             project_name=project_name,
         )
         session.add(experiment)
@@ -188,65 +198,104 @@ async def create_experiment(
             dataset_version_globalid = GlobalID(
                 "DatasetVersion", str(experiment.dataset_version_id)
             )
-    return CreateExperimentResponseBody(
-        data=Experiment(
-            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,
-            updated_at=experiment.updated_at,
-        )
-    )
+        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})
 
 
-class GetExperimentResponseBody(ResponseBody[Experiment]):
-    pass
-
-
-@router.get(
-    "/experiments/{experiment_id}",
-    operation_id="getExperiment",
-    summary="Get experiment by ID",
-    responses=add_errors_to_responses(
-        [{"status_code": HTTP_404_NOT_FOUND, "description": "Experiment not found"}]
-    ),
-    response_description="Experiment retrieved successfully",
-)
-async def get_experiment(request: Request, experiment_id: str) -> GetExperimentResponseBody:
-    experiment_globalid = GlobalID.from_id(experiment_id)
+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_rowid = from_global_id_with_expected_type(experiment_globalid, "Experiment")
+        experiment_id = from_global_id_with_expected_type(experiment_globalid, "Experiment")
     except ValueError:
-        raise HTTPException(
-            detail="Experiment with ID {experiment_globalid} does not exist",
+        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_rowid)
+            select(models.Experiment).where(models.Experiment.id == experiment_id)
         )
         experiment = experiment.scalar()
         if not experiment:
-            raise HTTPException(
-                detail=f"Experiment with ID {experiment_globalid} does not exist",
+            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))
-    return GetExperimentResponseBody(
-        data=Experiment(
-            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,
-            updated_at=experiment.updated_at,
-        )
-    )
+        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})