arize-phoenix 4.10.2rc0__py3-none-any.whl → 4.10.2rc2__py3-none-any.whl

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

@@ -1,8 +1,11 @@
+from datetime import datetime
 from random import getrandbits
+from typing import Any, Dict, Optional
 
+from fastapi import APIRouter, HTTPException
+from pydantic import BaseModel, 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
 
@@ -11,6 +14,10 @@ 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 .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())
@@ -24,94 +31,76 @@ def _generate_experiment_name(dataset_name: str) -> str:
     return f"{short_ds_name}-{_short_uuid()}"
 
 
-async def create_experiment(request: Request) -> Response:
+class Experiment(BaseModel):
+    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(BaseModel):
     """
-    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
+    Details of the experiment to be created
     """
-    dataset_globalid = GlobalID.from_id(request.path_params["dataset_id"])
+
+    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)
     try:
-        dataset_id = from_global_id_with_expected_type(dataset_globalid, "Dataset")
+        dataset_rowid = from_global_id_with_expected_type(dataset_globalid, "Dataset")
     except ValueError:
-        return Response(
-            content="Dataset with ID {dataset_globalid} does not exist",
+        raise HTTPException(
+            detail="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")
+    dataset_version_globalid_str = request_body.version_id
     if dataset_version_globalid_str is not None:
         try:
             dataset_version_globalid = GlobalID.from_id(dataset_version_globalid_str)
@@ -119,31 +108,31 @@ async def create_experiment(request: Request) -> Response:
                 dataset_version_globalid, "DatasetVersion"
             )
         except ValueError:
-            return Response(
-                content="DatasetVersion with ID {dataset_version_globalid} does not exist",
+            raise HTTPException(
+                detail="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))
+            await session.execute(select(models.Dataset).where(models.Dataset.id == dataset_rowid))
         ).scalar()
         if result is None:
-            return Response(
-                content=f"Dataset with ID {dataset_globalid} does not exist",
+            raise HTTPException(
+                detail=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)
+                .where(models.DatasetVersion.dataset_id == dataset_rowid)
                 .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",
+                raise HTTPException(
+                    detail=f"Dataset {dataset_globalid} does not have any versions",
                     status_code=HTTP_404_NOT_FOUND,
                 )
             dataset_version_id = dataset_version.id
@@ -154,24 +143,24 @@ async def create_experiment(request: Request) -> Response:
             )
             dataset_version = dataset_version.scalar()
             if not dataset_version:
-                return Response(
-                    content=f"DatasetVersion with ID {dataset_version_globalid} does not exist",
+                raise HTTPException(
+                    detail=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)
+        experiment_name = request_body.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_id=int(dataset_rowid),
             dataset_version_id=int(dataset_version_id),
             name=experiment_name,
-            description=payload.get("description"),
-            repetitions=repetitions,
-            metadata_=metadata,
+            description=request_body.description,
+            repetitions=request_body.repetitions,
+            metadata_=request_body.metadata or {},
             project_name=project_name,
         )
         session.add(experiment)
@@ -198,104 +187,65 @@ async def create_experiment(request: Request) -> Response:
             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})
+    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,
+        )
+    )
 
 
-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"])
+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)
     try:
-        experiment_id = from_global_id_with_expected_type(experiment_globalid, "Experiment")
+        experiment_rowid = from_global_id_with_expected_type(experiment_globalid, "Experiment")
     except ValueError:
-        return Response(
-            content="Experiment with ID {experiment_globalid} does not exist",
+        raise HTTPException(
+            detail="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)
+            select(models.Experiment).where(models.Experiment.id == experiment_rowid)
         )
         experiment = experiment.scalar()
         if not experiment:
-            return Response(
-                content=f"Experiment with ID {experiment_globalid} does not exist",
+            raise HTTPException(
+                detail=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})
+    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,
+        )
+    )